Re: [PATCH v3 4/5] rust: rbtree: add `RBTreeCursor`

[Benno Lossin]

On 18.04.24 16:15, Matt Gilbride wrote:
> @@ -400,6 +442,72 @@ fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
>      pub fn remove(&mut self, key: &K) -> Option<V> {
>          self.remove_node(key).map(|node| node.node.value)
>      }
> +
> +    /// Returns a cursor over the tree nodes based on the given key.
> +    ///
> +    /// If the given key exists, the cursor starts there.
> +    /// Otherwise it starts with the first larger key in sort order.
> +    /// If there is no larger key, it returns [`None`].
> +    pub fn cursor_lower_bound(&mut self, key: &K) -> Option<RBTreeCursor<'_, K, V>>
> +    where
> +        K: Ord,
> +    {
> +        let mut node = self.root.rb_node;
> +        let mut best_match: Option<NonNull<Node<K, V>>> = None;
> +        while !node.is_null() {
> +            // 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!(node, Node<K, V>, links) }.cast_mut();
> +            // SAFETY: `this` is a non-null node so it is valid by the type invariants.
> +            let this_key = unsafe { &(*this).key };
> +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> +            let left_child = unsafe { (*node).rb_left };
> +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> +            let right_child = unsafe { (*node).rb_right };

Since you have this pattern multiple times, I think you could have a
single function that walks the tree and takes care of most of the
`unsafe` stuff. A good starting point might be this:

     unsafe fn walk<F, R>(node: *mut bindings::rb_node, dir: F) -> R
     where
                /*     this,             key */
         F: FnMut(*mut bindings::rb_node, &K) -> Either<Direction, R>;

> +            if key == this_key {
> +                // INVARIANT:
> +                // - `self.root` and `node` are valid pointers.
> +                // - `self.root` is the root of an [`RBTree`].
> +                // - `node` is a valid node in an [`RBTree`].
> +                // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> +                //   borrows from `self`.
> +                return Some(RBTreeCursor {
> +                    _tree: PhantomData,
> +                    root: addr_of_mut!(self.root),
> +                    current: node,
> +                });
> +            } else {
> +                node = if key > this_key {
> +                    right_child
> +                } else {
> +                    let is_better_match = match best_match {
> +                        None => true,
> +                        Some(best) => {
> +                            // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> +                            let best_key = unsafe { &(*best.as_ptr()).key };
> +                            best_key > this_key
> +                        }
> +                    };
> +                    if is_better_match {
> +                        best_match = NonNull::new(this);
> +                    }
> +                    left_child
> +                }
> +            };
> +        }
> +        // INVARIANT:
> +        // - `self.root` and `best` are valid pointers.
> +        // - `self.root` is the root of an [`RBTree`].
> +        // - `best` is a valid node in an [`RBTree`].
> +        // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> +        //   borrows from `self`.
> +        best_match.map(|best| RBTreeCursor {
> +            _tree: PhantomData,
> +            root: addr_of_mut!(self.root),
> +            // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> +            current: unsafe { addr_of_mut!((*best.as_ptr()).links) },
> +        })
> +    }
>  }
> 
>  impl<K, V> Default for RBTree<K, V> {

[...]

> +/// # Invariants
> +/// - `root` and `current` are valid pointers.
> +/// - `root` points to the `root` node of an [`RBTree`].
> +/// - `current` points to a node that is in the same [`RBTree`] that `root` is pointing to.
> +/// - A cursor must borrow the [`RBTree`] containing `root` and `current` mutably.
> +pub struct RBTreeCursor<'a, K, V> {
> +    _tree: PhantomData<&'a RBTree<K, V>>,

Why is this not `&'a mut RBTree<K, V>`?

> +    root: *mut bindings::rb_root,
> +    current: *mut bindings::rb_node,
> +}
> +
> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
> +// so it has the same thread safety requirements as mutable references.
> +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}
> +
> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
> +// so it has the same thread safety requirements as mutable references.
> +unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeCursor<'a, K, V> {}
> +
> +impl<'a, K, V> RBTreeCursor<'a, K, V> {
> +    /// The current node
> +    pub fn current(&self) -> (&K, &V) {
> +        // SAFETY:
> +        // - `self.current` is a valid node by the type invariants.
> +        // - We have an immutable reference by the function signature.
> +        unsafe { Self::to_key_value(self.current) }
> +    }
> +
> +    /// The current node, with a mutable value
> +    pub fn current_mut(&mut self) -> (&K, &mut V) {
> +        // SAFETY:
> +        // - `self.current` is a valid node by the type invariants.
> +        // - We have an mutable reference by the function signature.
> +        unsafe { Self::to_key_value_mut(self.current) }
> +    }
> +
> +    /// Remove the current node from the tree.
> +    ///
> +    /// Returns a cursor to the next node, if it exists,
> +    /// else the previous node. Returns [`None`] if the tree
> +    /// becomes empty.
> +    pub fn remove_current(self) -> Option<Self> {
> +        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, Node<K, V>, links) }.cast_mut();
> +        // 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, self.root) };
> +
> +        let current = match (prev, next) {
> +            (_, Some(next)) => next,
> +            (Some(prev), None) => prev,
> +            (None, None) => {
> +                return None;
> +            }
> +        };
> +
> +        // INVARIANT:
> +        // - `self.root` and `current` are valid pointers.
> +        // - `self.root` is the root of an [`RBTree`].
> +        // - `current` is a valid node in an [`RBTree`].
> +        // - Due to the function signature, `self` is an owned [`RBTreeCursor`],
> +        //   and [`RBTreeCursor`]s are only created via functions with a mutable reference
> +        //   to an [`RBTree`].
> +        Some(Self {
> +            current,
> +            _tree: self._tree,
> +            root: self.root,
> +        })
> +    }
> +
> +    /// Remove the previous node, returning it if it exists.
> +    pub fn remove_prev(&mut self) -> Option<(K, V)> {

Why do these functions not return `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<(K, V)> {
> +        self.remove_neighbor(Direction::Next)
> +    }
> +
> +    fn remove_neighbor(&mut self, direction: Direction) -> Option<(K, V)> {
> +        if let Some(neighbor) = self.get_neighbor_raw(direction) {
> +            // 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, self.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 n = unsafe { Box::from_raw(this) };
> +            return Some((n.key, n.value));
> +        }
> +        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> {

Does it hurt to name this `move`?

> +        // INVARIANT:
> +        // - `self.root` and `neighbor` are valid pointers.
> +        // - `self.root` is the root of an [`RBTree`].
> +        // - `neighbor` is a valid node in an [`RBTree`].
> +        // - Due to the function signature, `self` is an owned [`RBTreeCursor`],
> +        //   and [`RBTreeCursor`]s are only created via functions with a mutable reference
> +        //   to an [`RBTree`].
> +        self.get_neighbor_raw(direction).map(|neighbor| Self {
> +            _tree: self._tree,
> +            root: self.root,
> +            current: neighbor,
> +        })
> +    }

[...]

> +    /// SAFETY:
> +    /// - `node` must be a valid pointer to a node in an [`RBTree`].
> +    /// - The caller has immutable access to `node` for the duration of 'a.
> +    unsafe fn to_key_value(node: *mut bindings::rb_node) -> (&'a K, &'a V) {
> +        // 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!(node, Node<K, V>, links) };
> +        // SAFETY: The passed `node` is the current node or a non-null neighbor,
> +        // thus `this` is valid by the type invariants.
> +        let k = unsafe { &(*this).key };
> +        // SAFETY: The passed `node` is the current node or a non-null neighbor,
> +        // thus `this` is valid by the type invariants.
> +        let v = unsafe { &(*this).value };
> +        (k, v)
> +    }
> +
> +    /// SAFETY:
> +    /// - `node` must be a valid pointer to a node in an [`RBTree`].
> +    /// - The caller has mutable access to `node` for the duration of 'a.
> +    unsafe fn to_key_value_mut(node: *mut bindings::rb_node) -> (&'a K, &'a mut V) {
> +        // 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!(node, Node<K, V>, links) }.cast_mut();
> +        // SAFETY: The passed `node` is the current node or a non-null neighbor,
> +        // thus `this` is valid by the type invariants.
> +        let k = unsafe { &(*this).key };
> +        // SAFETY: The passed `node` is the current node or a non-null neighbor,
> +        // thus `this` is valid by the type invariants.
> +        let v = unsafe { &mut (*this).value };
> +        (k, v)
> +    }

You can create a single function that does the `container_of!` stuff and
that returns `(*mut K, *mut V)` and implement these two in terms of that
one.

-- 
Cheers,
Benno

> +}
> +
> +/// Direction for [`RBTreeCursor`] operations.
> +enum Direction {
> +    /// the node immediately before, in sort order
> +    Prev,
> +    /// the node immediately after, in sort order
> +    Next,
> +}
> +
>  impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
>      type Item = (&'a K, &'a V);
>      type IntoIter = RBTreeIterator<'a, K, V>;
> 
> --
> 2.44.0.769.g3c40516874-goog
>