Re: [PATCH v3 5/5] rust: rbtree: add `RBTree::entry`

[Benno Lossin]

On 18.04.24 16:15, Matt Gilbride wrote:
> @@ -332,63 +338,54 @@ pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTree
>          // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere
>          // in the subtree of `parent` that `child_field_of_parent` points at. Once
>          // we find an empty subtree, we can insert the new node using `rb_link_node`.
> -        let mut parent = core::ptr::null_mut();
>          let mut child_field_of_parent: &mut *mut bindings::rb_node = &mut self.root.rb_node;
> -        while !child_field_of_parent.is_null() {
> -            parent = *child_field_of_parent;
> +        let mut parent = core::ptr::null_mut();

Nit: why are you moving this line below `child_field_of_parent`? Just an
artifact of rebasing?

> +        while !(*child_field_of_parent).is_null() {
> +            let curr = *child_field_of_parent;
> +            // SAFETY: All links fields we create are in a `Node<K, V>`.
> +            let node = unsafe { container_of!(curr, Node<K, V>, links) };

[...]

> @@ -1119,3 +1099,177 @@ 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> {}
> +
> +impl<K, V> RBTreeNode<K, V> {
> +    /// Drop the key and value, but keep the allocation.
> +    ///
> +    /// It then becomes a reservation that can be re-initialised into a different node (i.e., with
> +    /// a different key and/or value).
> +    ///
> +    /// The existing key and value are dropped in-place as part of this operation, that is, memory
> +    /// may be freed (but only for the key/value; memory for the node itself is kept for reuse).
> +    pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> {
> +        let raw = Box::into_raw(self.node);
> +        let mut ret = RBTreeNodeReservation {
> +            // SAFETY: The pointer came from a valid `Node`, which has the same layout as
> +            // `MaybeUninit<Node>`.
> +            node: unsafe { Box::from_raw(raw as _) },
> +        };
> +        // SAFETY: Although the type is `MaybeUninit<Node>`, we know it has been initialised
> +        // because it came from a `Node`. So it is safe to drop it.
> +        unsafe { core::ptr::drop_in_place::<Node<K, V>>(ret.node.as_mut_ptr()) };
> +        ret
> +    }

With my patch [1] this can be simplified.

[1]: https://lore.kernel.org/rust-for-linux/20240425213419.3904105-1-benno.lossin@xxxxxxxxx/

> +}
> +
> +/// A view into a single entry in a map, which may either be vacant or occupied.
> +///
> +/// This enum is constructed from the [`entry`] method on [`RBTree`].

You could just write [`RBTree::entry`].

> +///
> +/// [`entry`]: fn@RBTree::entry
> +pub enum Entry<'a, K, V> {
> +    /// This [`RBTree`] does not have a node with this key.
> +    Vacant(VacantEntry<'a, K, V>),
> +    /// This [`RBTree`] already has a node with this key.
> +    Occupied(OccupiedEntry<'a, K, V>),
> +}

[...]

> +impl<'a, K, V> RawVacantEntry<'a, K, V> {
> +    /// Inserts the given node into the [`RBTree`] at this entry.
> +    ///
> +    /// The `node` must have a key such that inserting it here does not break the ordering of this
> +    /// [`RBTree`].
> +    fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V {
> +        let node = Box::into_raw(node.node);
> +
> +        // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when
> +        // the node is removed or replaced.
> +        let node_links = unsafe { addr_of_mut!((*node).links) };
> +
> +        // INVARIANT: We are linking in a new node, which is valid. It remains valid because we
> +        // "forgot" it with `Box::into_raw`.
> +        // SAFETY: All pointers are null or valid in an appropriate way.

I don't like the formulation "valid in an appropriate way", since if you
don't know what the appropriate way is, this doesn't help you.

> +        unsafe { bindings::rb_link_node(node_links, self.parent, self.child_field_of_parent) };
> +
> +        // SAFETY: All pointers are valid. `node` has just been inserted into the tree.
> +        unsafe { bindings::rb_insert_color(node_links, &mut self.rbtree.root) };
> +
> +        // SAFETY: The node is valid until we remove it from the tree.
> +        unsafe { &mut (*node).value }
> +    }
> +}
> +
> +impl<'a, K, V> VacantEntry<'a, K, V> {
> +    /// Inserts the given node into the [`RBTree`] at this entry.
> +    pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V {
> +        self.raw.insert(reservation.into_node(self.key, value))
> +    }
> +}
> +
> +/// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum.
> +///
> +/// # Invariants
> +/// - `node_links` is a valid, non-null pointer to a tree node.

It should be the same tree as `self.rbtree`, right? (I see you calling
`rb_replace_node` below with the rbtree root used)

-- 
Cheers,
Benno

> +pub struct OccupiedEntry<'a, K, V> {
> +    rbtree: &'a mut RBTree<K, V>,
> +    /// The node that this entry corresponds to.
> +    node_links: *mut bindings::rb_node,
> +}