I was trying to implement Monte Carlo Tree Search, with some methods on the nodes in my tree data structure accessing a child node, continuing to select child nodes to arbitrary depth, and eventually one child node backpropagating some data up through all of its parent nodes, which are all mutated in this process. The tree structure also gets pruned in a way that some parent nodes will be deleted and all of that parent's children will also need to be deleted, which I thought made it a poor fit for some arena memory approach.
I ended up implementing this with this smart pointer, but it allows for aliasable mutability. Is this type of setup bad practice?
use std::ptr::NonNull;
/// Smart pointer to use in tree data structures
/// The `TreePtr` type is a strong pointer which can be used safely
/// This type owns its contents, and can be used for parent nodes to
/// own their children
///
/// Will drop/cleanup its contents when out of scope.
pub struct TreePtr<T> {
target: Option<NonNull<T>>,
}
/// Weak version of `TreePtr`, does not own its contents
/// Can be used by nodes to refer to themselves or their parent
///
/// For safe usage, the weak pointer should be used for self-referential pointing
/// or for pointing to a structure which owns the strong version of the pointer.
/// This ensures the required behavior of the parent pointer being dropped before the weak version.
///
/// Allows `Send` but is not thread safe. Use wisely.
pub struct TreePtrWeak<T> {
target: Option<NonNull<T>>,
}
impl<T> Drop for TreePtr<T> {
fn drop(&mut self) {
if let Some(target) = self.target {
unsafe {
Box::from_raw(NonNull::as_ptr(target));
}
}
}
}
impl<T> TreePtr<T> {
/// Get a strong `TreePtr` with no contents
/// A null `TreePtr` cannot be dereferenced, and will panic
/// Therefore, you must ensure to immediately change this pointer to refer to something
/// before attempting to dereference it
pub unsafe fn null() -> Self {
TreePtr { target: None }
}
/// Allocate `val` on the stack and return a strong `TreePtr` to it
pub fn new(val: T) -> Self {
// Will always be some, pointer isn't null
let target = NonNull::new(Box::leak(Box::new(val)));
TreePtr { target }
}
/// Get weak `TreePtrWeak` corresponding to strong `TreePtr`
/// Cannot be safely dereferenced after `ptr` is dropped
/// To use safely, `ptr` should be owned by the same structure as the returned `TreePtrWeak`,
/// or the returned `TreePtrWeak` should be otherwise indirectly owned by the owner of `ptr`
pub fn downgrade(ptr: &Self) -> TreePtrWeak<T> {
TreePtrWeak { target: ptr.target }
}
}
impl<T> Clone for TreePtrWeak<T> {
/// Create a new `TreePtrWeak` to the same underlying data without copying that data
/// Ensure the cloned pointer meets the same safety requirements with respect to the parent
/// `TreePtr` owning the underlying data.
fn clone(&self) -> Self {
unsafe {
TreePtrWeak {
target: self
.target
.and_then(|ptr| Some(NonNull::new_unchecked(ptr.as_ptr()))),
}
}
}
}
impl<T> std::ops::Deref for TreePtr<T> {
type Target = T;
/// Dereference the pointer
///
/// ### Panics
/// This operation will panic if used on the output of `TreePtr::null()`
/// Instead, replace that pointer with something else before dereferencing
fn deref(&self) -> &Self::Target {
unsafe {
&*self
.target
.expect("'Null' CleanupPtr dereferenced")
.as_ptr()
}
}
}
impl<T> std::ops::DerefMut for TreePtr<T> {
/// Mutably dereference the pointer
///
/// ### Panics
/// This operation will panic if used on the output of `TreePtr::null()`
/// Instead, replace that pointer with something else before dereferencing
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe {
&mut *self
.target
.expect("'Null' CleanupPtr dereferenced")
.as_ptr()
}
}
}
impl<T> TreePtrWeak<T> {
/// Unsafely dereference a `TreePtrWeak`
/// This is **undefined behavior** if the parent `TreePtr`
/// the `TreePtrWeak` was cloned from has been dropped.
/// You must ensure the `TreePtrWeak` is dropped with its parent,
/// and therefore cannot be dereferenced after its parent is dropped.
///
/// ### Panics
/// This operation will panic if the parent `TreePtr` was the output of `TreePtr::null()`
pub unsafe fn deref(&self) -> &T {
&*self
.target
.expect("'Null' CleanupPtrWeak dereferenced")
.as_ptr()
}
/// Unsafely mutably dereference a `TreePtrWeak`
/// This is **undefined behavior** if the parent `TreePtr`
/// the `TreePtrWeak` was cloned from has been dropped.
/// You must ensure the `TreePtrWeak` is dropped with its parent,
/// and therefore cannot be dereferenced after its parent is dropped.
///
/// ### Panics
/// This operation will panic if the parent `TreePtr` was the output of `TreePtr::null()`
pub unsafe fn deref_mut(&mut self) -> &mut T {
&mut *self
.target
.expect("'Null' CleanupPtrWeak dereferenced")
.as_ptr()
}
}
unsafe impl<T> Send for TreePtrWeak<T> {}
For additional clarification, the pointer was used only as an element of this struct, where children are always dropped when their parents are dropped because their parents own the strong pointers:
pub struct MCTSNode {
pub parent: Option<TreePtrWeak<MCTSNode>>,
pub this: TreePtrWeak<MCTSNode>,
pub children: Vec<TreePtr<MCTSNode>>,
...
}