Re: [PATCH v4 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Boqun Feng <boqun.feng@gmail.com>]

On Mon, Jun 03, 2024 at 04:05:17PM +0000, Matt Gilbride wrote:
[...]
> +/// A memory reservation for a red-black tree node.
> +///
> +///
> +/// It contains the memory needed to hold a node that can be inserted into a red-black tree. One
> +/// can be obtained by directly allocating it ([`RBTreeNodeReservation::new`]).
> +pub struct RBTreeNodeReservation<K, V> {
> +    node: Box<MaybeUninit<Node<K, V>>>,
> +}
> +
> +impl<K, V> RBTreeNodeReservation<K, V> {
> +    /// Allocates memory for a node to be eventually initialised and inserted into the tree via a
> +    /// call to [`RBTree::insert`].
> +    pub fn new(flags: Flags) -> Result<RBTreeNodeReservation<K, V>> {
> +        Ok(RBTreeNodeReservation {
> +            node: Box::new_uninit(flags)?,
> +        })
> +    }
> +}
> +
> +// SAFETY: This doesn't actually contain K or V, and is just a memory allocation. Those can always
> +// be moved across threads.
> +unsafe impl<K, V> Send for RBTreeNodeReservation<K, V> {}
> +
> +// SAFETY: This doesn't actually contain K or V, and is just a memory allocation.
> +unsafe impl<K, V> Sync for RBTreeNodeReservation<K, V> {}
> +
> +impl<K, V> RBTreeNodeReservation<K, V> {
> +    /// Initialises a node reservation.
> +    ///
> +    /// It then becomes an [`RBTreeNode`] that can be inserted into a tree.
> +    pub fn into_node(mut self, key: K, value: V) -> RBTreeNode<K, V> {
> +        let node_ptr = self.node.as_mut_ptr();
> +        // SAFETY: `node_ptr` is a valid pointer to a tree node.
> +        unsafe {
> +            node_ptr.write(Node {
> +                key,
> +                value,
> +                links: bindings::rb_node::default(),
> +            })
> +        }
> +        RBTreeNode {
> +            // SAFETY: The pointer came from a `MaybeUninit<Node>` whose fields have all been
> +            // initialised. Additionally, it has the same layout as `Node`.
> +            node: unsafe { Box::<MaybeUninit<_>>::assume_init(self.node) },
> +        }

nit: could you use Box::write()[1] here? It saves two `unsafe` blocks.

[1]: https://doc.rust-lang.org/std/boxed/struct.Box.html#method.write 

Regards,
Boqun

> +    }
> +}
> +
> +/// A red-black tree node.
> +///
> +/// The node is fully initialised (with key and value) and can be inserted into a tree without any
> +/// extra allocations or failure paths.
> +pub struct RBTreeNode<K, V> {
> +    node: Box<Node<K, V>>,
> +}
> +
> +impl<K, V> RBTreeNode<K, V> {
> +    /// Allocates and initialises a node that can be inserted into the tree via
> +    /// [`RBTree::insert`].
> +    pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> {
> +        Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value))
> +    }
> +}
> +
> +// SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across
> +// threads.
> +unsafe impl<K: Send, V: Send> Send for RBTreeNode<K, V> {}
> +
> +// SAFETY: If K and V can be accessed without synchronization, then it's also okay to access
> +// [`RBTreeNode`] without synchronization.
> +unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {}
> +
> +struct Node<K, V> {
> +    links: bindings::rb_node,
> +    key: K,
> +    value: V,
> +}
> 
> -- 
> 2.45.1.288.g0e0cd299f1-goog
>