Re: [PATCH v11 4/5] rust: rbtree: add cursor

[Benno Lossin <benno.lossin@proton.me>]

On 21.08.24 19:42, Matt Gilbride wrote:
> +    /// Remove the current node from the tree.
> +    ///
> +    /// Returns a tuple where the first element is a cursor to the next node, if it exists,
> +    /// else the previous node, else [`None`] (if the tree becomes empty). The second element
> +    /// is the removed node.
> +    pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) {
> +        let prev = self.get_neighbor_raw(Direction::Prev);
> +        let next = self.get_neighbor_raw(Direction::Next);
> +        // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
> +        // point to the links field of `Node<K, V>` objects.
> +        let this = unsafe { container_of!(self.current.as_ptr(), Node<K, V>, links) }.cast_mut();
> +        // SAFETY: `this` is valid by the type invariants as described above.
> +        let node = unsafe { Box::from_raw(this) };
> +        let node = RBTreeNode { node };
> +        // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
> +        // the tree cannot change. By the tree invariant, all nodes are valid.
> +        unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.tree.root)) };
> +
> +        let current = match (prev, next) {
> +            (_, Some(next)) => next,
> +            (Some(prev), None) => prev,
> +            (None, None) => {
> +                return (None, node);
> +            }
> +        };
> +
> +        (
> +            // INVARIANT:
> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`.
> +            // - Due to the function signature, `self` is an owned [`Cursor`],
> +            //   and [`Cursor`]s are only created via functions with a mutable reference
> +            //   to an [`RBTree`].

You missed this stale invariant.

> +            Some(Self {
> +                current,
> +                tree: self.tree,
> +            }),
> +            node,
> +        )
> +    }
> +
> +    /// Remove the previous node, returning it if it exists.
> +    pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> {
> +        self.remove_neighbor(Direction::Prev)
> +    }
> +
> +    /// Remove the next node, returning it if it exists.
> +    pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> {
> +        self.remove_neighbor(Direction::Next)
> +    }
> +
> +    fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> {
> +        if let Some(neighbor) = self.get_neighbor_raw(direction) {
> +            let neighbor = neighbor.as_ptr();
> +            // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
> +            // the tree cannot change. By the tree invariant, all nodes are valid.
> +            unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.root)) };
> +            // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
> +            // point to the links field of `Node<K, V>` objects.
> +            let this = unsafe { container_of!(neighbor, Node<K, V>, links) }.cast_mut();
> +            // SAFETY: `this` is valid by the type invariants as described above.
> +            let node = unsafe { Box::from_raw(this) };
> +            return Some(RBTreeNode { node });
> +        }
> +        None
> +    }
> +
> +    /// Move the cursor to the previous node, returning [`None`] if it doesn't exist.
> +    pub fn move_prev(self) -> Option<Self> {
> +        self.mv(Direction::Prev)
> +    }
> +
> +    /// Move the cursor to the next node, returning [`None`] if it doesn't exist.
> +    pub fn move_next(self) -> Option<Self> {
> +        self.mv(Direction::Next)
> +    }
> +
> +    fn mv(self, direction: Direction) -> Option<Self> {
> +        // INVARIANT:
> +        // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`.
> +        // - Due to the function signature, `self` is an owned [`Cursor`],
> +        //   and [`Cursor`]s are only created via functions with a mutable reference
> +        //   to an [`RBTree`].

And this one too.

With those fixed:

Reviewed-by: Benno Lossin <benno.lossin@proton.me>

---
Cheers,
Benno

> +        self.get_neighbor_raw(direction).map(|neighbor| Self {
> +            tree: self.tree,
> +            current: neighbor,
> +        })
> +    }