rust: rbtree: reduce unsafe blocks on pointer derefs

Refactors parts of the get() and find_best_match()
traversal logic to minimize the scope of unsafe blocks
and avoid duplicating same safety comments.

One of the removed comments was also misleading:

    // SAFETY: `node` is a non-null node...
    Ordering::Equal => return Some(unsafe { &(*this).value }),

as `node` should have been `this`.

No functional changes intended; this is purely a safety
improvement that reduces the amount of unsafe blocks
while keeping all invariants intact.

[ Alice writes:

    "One consequence of creating a &_ to the bindings::rb_node struct means
    that we assert immutability for the entire struct and not just the
    rb_left/rb_right fields, but I have verified that this is ok."

      - Miguel ]

Signed-off-by: Onur Özkan <work@onurozkan.dev>
Reviewed-by: Charalampos Mitrodimas <charmitro@posteo.net>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Link: https://patch.msgid.link/20251113144547.502-1-work@onurozkan.dev
[ Reworded title and replaced `cursor_lower_bound()` with
  `find_best_match()` in message. - Miguel ]
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

1 files changed, 15 insertions, 12 deletions

diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
index 4729eb56827a..ed3582d88e4e 100644
--- a/rust/kernel/rbtree.rs
+++ b/rust/kernel/rbtree.rs
@@ -414,14 +414,17 @@ where
             // 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: `this` is a non-null node so it is valid by the type invariants.
-            node = match key.cmp(unsafe { &(*this).key }) {
-                // SAFETY: `node` is a non-null node so it is valid by the type invariants.
-                Ordering::Less => unsafe { (*node).rb_left },
-                // SAFETY: `node` is a non-null node so it is valid by the type invariants.
-                Ordering::Greater => unsafe { (*node).rb_right },
-                // SAFETY: `node` is a non-null node so it is valid by the type invariants.
-                Ordering::Equal => return Some(unsafe { &(*this).value }),
+            let this_ref = unsafe { &*this };
+
+            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
+            let node_ref = unsafe { &*node };
+
+            node = match key.cmp(&this_ref.key) {
+                Ordering::Less => node_ref.rb_left,
+                Ordering::Greater => node_ref.rb_right,
+                Ordering::Equal => return Some(&this_ref.value),
             }
         }
         None
@@ -498,10 +501,10 @@ where
             let this = unsafe { container_of!(node, Node<K, V>, links) };
             // 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 };
+            let node_ref = unsafe { &*node };
+
             match key.cmp(this_key) {
                 Ordering::Equal => {
                     // SAFETY: `this` is a non-null node so it is valid by the type invariants.
@@ -509,7 +512,7 @@ where
                     break;
                 }
                 Ordering::Greater => {
-                    node = right_child;
+                    node = node_ref.rb_right;
                 }
                 Ordering::Less => {
                     let is_better_match = match best_key {
@@ -521,7 +524,7 @@ where
                         // SAFETY: `this` is a non-null node so it is valid by the type invariants.
                         best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) });
                     }
-                    node = left_child;
+                    node = node_ref.rb_left;
                 }
             };
         }