Re: [PATCH v8 6/6] rust: rbtree: add `RBTree::entry`

[Alice Ryhl <aliceryhl@google.com>]

On Mon, Aug 5, 2024 at 10:03 PM Benno Lossin <benno.lossin@proton.me> wrote:
>
> On 27.07.24 22:30, Matt Gilbride wrote:
> > From: Alice Ryhl <aliceryhl@google.com>
> >
> > This mirrors the entry API [1] from the Rust standard library on
> > `RBTree`. This API can be used to access the entry at a specific key and
> > make modifications depending on whether the key is vacant or occupied.
> > This API is useful because it can often be used to avoid traversing the
> > tree multiple times.
> >
> > This is used by binder to look up and conditionally access or insert a
> > value, depending on whether it is there or not [2].
> >
> > Link: https://doc.rust-lang.org/stable/std/collections/btree_map/enum.Entry.html [1]
> > Link: https://android-review.googlesource.com/c/kernel/common/+/2849906 [2]
> > Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> > Tested-by: Alice Ryhl <aliceryhl@google.com>
> > Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> > ---
> >  rust/kernel/rbtree.rs | 302 +++++++++++++++++++++++++++++++++++++-------------
> >  1 file changed, 227 insertions(+), 75 deletions(-)
> >
> > diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> > index 5d37aa373685..428f8be8f3a2 100644
> > --- a/rust/kernel/rbtree.rs
> > +++ b/rust/kernel/rbtree.rs
> > @@ -295,12 +295,19 @@ pub fn try_create_and_insert(
> >      /// key/value pair). Returns [`None`] if a node with the same key didn't already exist.
> >      ///
> >      /// This function always succeeds.
> > -    pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
> > -        let node = Box::into_raw(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) };
> > +    pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
> > +        match self.raw_entry(&node.node.key) {
> > +            RawEntry::Occupied(entry) => Some(entry.replace(node)),
> > +            RawEntry::Vacant(entry) => {
> > +                entry.insert(node);
> > +                None
> > +            }
> > +        }
> > +    }
> >
> > +    fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
> > +        let raw_self: *mut RBTree<K, V> = self;
> > +        // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`.
> >          // The parameters of `rb_link_node` are as follows:
> >          // - `node`: A pointer to an uninitialized node being inserted.
> >          // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be
> > @@ -319,62 +326,56 @@ pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTree
> >          // 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 child_field_of_parent: &mut *mut bindings::rb_node =
> > +            // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above).
> > +            unsafe { &mut (*raw_self).root.rb_node };
> > +        while !(*child_field_of_parent).is_null() {
>
> Why do you manually dereference `child_field_of_parent` here?

I think it helps for clarity. It makes it clear to the reader that
`child_field_of_parent` isn't null, but that it's pointing at a null
pointer.

> > +impl<'a, K, V> OccupiedEntry<'a, K, V> {
> > +    fn node_ptr(&self) -> *mut Node<K, V> {
> > +        // SAFETY: By the type invariant of `Self`, all `node_links` pointers stored in `self`
> > +        // point to the links field of `Node<K, V>` objects.
> > +        unsafe { container_of!(self.node_links, Node<K, V>, links) }.cast_mut()
>
> You should not call `cast_mut` here, see below
>
> > +    }
> > +
> > +    /// Gets a reference to the value in the entry.
> > +    pub fn get(&self) -> &V {
> > +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
>
> Can you add a `# Guarantees` section to `node_ptr` that states exactly
> this?
>
> > +        unsafe { &(*self.node_ptr()).value }
> > +    }
> > +
> > +    /// Gets a mutable reference to the value in the entry.
> > +    pub fn get_mut(&mut self) -> &mut V {
> > +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
> > +        unsafe { &mut (*self.node_ptr()).value }
>
> This is sadly UB, you are creating a `&mut` reference from a pointer
> that was derived from a `&` reference:
> - `node_ptr` takes `&self`, thus it converts the `&mut self` to that.
> - `container_of!` inside of `node_ptr` is used to create a read-only
>   pointer to the `links` field (it is casted to `*mut`, but that doesn't
>   change the fact that you are only allowed to use it for reads)
> - `get_mut` turns it again into a `&mut` reference.
>
> One solution is to make `note_ptr` take a `*mut Self`/`*const Self`.

We will get rid of node_ptr and duplicate its contents into get and get_mut.

> > +    }
> > +
> > +    /// Converts the entry into a mutable reference to its value.
> > +    ///
> > +    /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`].
> > +    pub fn into_mut(self) -> &'a mut V {
> > +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
> > +        unsafe { &mut (*self.node_ptr()).value }
> > +    }
> > +
> > +    /// Remove this entry from the [`RBTree`].
> > +    pub fn remove_node(self) -> RBTreeNode<K, V> {
> > +        // SAFETY: The node is a node in the tree, so it is valid.
> > +        unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.root) };
> > +
> > +        // INVARIANT: The node is being returned and the caller may free it, however, it was
> > +        // removed from the tree. So the invariants still hold.
> > +        RBTreeNode {
> > +            // SAFETY: The node was a node in the tree, but we removed it, so we can convert it
> > +            // back into a box.
> > +            node: unsafe { Box::from_raw(self.node_ptr()) },
> > +        }
> > +    }
> > +
> > +    /// Takes the value of the entry out of the map, and returns it.
> > +    pub fn remove(self) -> V {
> > +        self.remove_node().node.value
> > +    }
> > +
> > +    /// Swap the current node for the provided node.
> > +    ///
> > +    /// The key of both nodes must be equal.
>
> Is this a safety requirement? Ie if it is violated, can memory bugs
> occur, or is it only going to lead to logic bugs?

No, it's not a safety requirement.

Alice