rust: rbtree: add red-black tree implementation backed by the C version

+1

rust/helpers/helpers.c

··· 15 15 #include "kunit.c" 16 16 #include "mutex.c" 17 17 #include "page.c" 18 + #include "rbtree.c" 18 19 #include "refcount.c" 19 20 #include "signal.c" 20 21 #include "slab.c"

+9

rust/helpers/rbtree.c

··· 1 + // SPDX-License-Identifier: GPL-2.0 2 + 3 + #include <linux/rbtree.h> 4 + 5 + void rust_helper_rb_link_node(struct rb_node *node, struct rb_node *parent, 6 + struct rb_node **rb_link) 7 + { 8 + rb_link_node(node, parent, rb_link); 9 + }

+1

rust/kernel/lib.rs

··· 44 44 pub mod page; 45 45 pub mod prelude; 46 46 pub mod print; 47 + pub mod rbtree; 47 48 mod static_assert; 48 49 #[doc(hidden)] 49 50 pub mod std_vendor;

+432

rust/kernel/rbtree.rs

··· 1 + // SPDX-License-Identifier: GPL-2.0 2 + 3 + //! Red-black trees. 4 + //! 5 + //! C header: [`include/linux/rbtree.h`](srctree/include/linux/rbtree.h) 6 + //! 7 + //! Reference: <https://docs.kernel.org/core-api/rbtree.html> 8 + 9 + use crate::{alloc::Flags, bindings, container_of, error::Result, prelude::*}; 10 + use alloc::boxed::Box; 11 + use core::{ 12 + cmp::{Ord, Ordering}, 13 + marker::PhantomData, 14 + mem::MaybeUninit, 15 + ptr::{addr_of_mut, NonNull}, 16 + }; 17 + 18 + /// A red-black tree with owned nodes. 19 + /// 20 + /// It is backed by the kernel C red-black trees. 21 + /// 22 + /// # Examples 23 + /// 24 + /// In the example below we do several operations on a tree. We note that insertions may fail if 25 + /// the system is out of memory. 26 + /// 27 + /// ``` 28 + /// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation}}; 29 + /// 30 + /// // Create a new tree. 31 + /// let mut tree = RBTree::new(); 32 + /// 33 + /// // Insert three elements. 34 + /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 35 + /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 36 + /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 37 + /// 38 + /// // Check the nodes we just inserted. 39 + /// { 40 + /// assert_eq!(tree.get(&10).unwrap(), &100); 41 + /// assert_eq!(tree.get(&20).unwrap(), &200); 42 + /// assert_eq!(tree.get(&30).unwrap(), &300); 43 + /// } 44 + /// 45 + /// // Replace one of the elements. 46 + /// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?; 47 + /// 48 + /// // Check that the tree reflects the replacement. 49 + /// { 50 + /// assert_eq!(tree.get(&10).unwrap(), &1000); 51 + /// assert_eq!(tree.get(&20).unwrap(), &200); 52 + /// assert_eq!(tree.get(&30).unwrap(), &300); 53 + /// } 54 + /// 55 + /// // Change the value of one of the elements. 56 + /// *tree.get_mut(&30).unwrap() = 3000; 57 + /// 58 + /// // Check that the tree reflects the update. 59 + /// { 60 + /// assert_eq!(tree.get(&10).unwrap(), &1000); 61 + /// assert_eq!(tree.get(&20).unwrap(), &200); 62 + /// assert_eq!(tree.get(&30).unwrap(), &3000); 63 + /// } 64 + /// 65 + /// // Remove an element. 66 + /// tree.remove(&10); 67 + /// 68 + /// // Check that the tree reflects the removal. 69 + /// { 70 + /// assert_eq!(tree.get(&10), None); 71 + /// assert_eq!(tree.get(&20).unwrap(), &200); 72 + /// assert_eq!(tree.get(&30).unwrap(), &3000); 73 + /// } 74 + /// 75 + /// # Ok::<(), Error>(()) 76 + /// ``` 77 + /// 78 + /// In the example below, we first allocate a node, acquire a spinlock, then insert the node into 79 + /// the tree. This is useful when the insertion context does not allow sleeping, for example, when 80 + /// holding a spinlock. 81 + /// 82 + /// ``` 83 + /// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode}, sync::SpinLock}; 84 + /// 85 + /// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result { 86 + /// // Pre-allocate node. This may fail (as it allocates memory). 87 + /// let node = RBTreeNode::new(10, 100, flags::GFP_KERNEL)?; 88 + /// 89 + /// // Insert node while holding the lock. It is guaranteed to succeed with no allocation 90 + /// // attempts. 91 + /// let mut guard = tree.lock(); 92 + /// guard.insert(node); 93 + /// Ok(()) 94 + /// } 95 + /// ``` 96 + /// 97 + /// In the example below, we reuse an existing node allocation from an element we removed. 98 + /// 99 + /// ``` 100 + /// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNodeReservation}}; 101 + /// 102 + /// // Create a new tree. 103 + /// let mut tree = RBTree::new(); 104 + /// 105 + /// // Insert three elements. 106 + /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 107 + /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 108 + /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 109 + /// 110 + /// // Check the nodes we just inserted. 111 + /// { 112 + /// assert_eq!(tree.get(&10).unwrap(), &100); 113 + /// assert_eq!(tree.get(&20).unwrap(), &200); 114 + /// assert_eq!(tree.get(&30).unwrap(), &300); 115 + /// } 116 + /// 117 + /// // Remove a node, getting back ownership of it. 118 + /// let existing = tree.remove(&30).unwrap(); 119 + /// 120 + /// // Check that the tree reflects the removal. 121 + /// { 122 + /// assert_eq!(tree.get(&10).unwrap(), &100); 123 + /// assert_eq!(tree.get(&20).unwrap(), &200); 124 + /// assert_eq!(tree.get(&30), None); 125 + /// } 126 + /// 127 + /// // Create a preallocated reservation that we can re-use later. 128 + /// let reservation = RBTreeNodeReservation::new(flags::GFP_KERNEL)?; 129 + /// 130 + /// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to 131 + /// // succeed (no memory allocations). 132 + /// tree.insert(reservation.into_node(15, 150)); 133 + /// 134 + /// // Check that the tree reflect the new insertion. 135 + /// { 136 + /// assert_eq!(tree.get(&10).unwrap(), &100); 137 + /// assert_eq!(tree.get(&15).unwrap(), &150); 138 + /// assert_eq!(tree.get(&20).unwrap(), &200); 139 + /// } 140 + /// 141 + /// # Ok::<(), Error>(()) 142 + /// ``` 143 + /// 144 + /// # Invariants 145 + /// 146 + /// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always 147 + /// valid, and pointing to a field of our internal representation of a node. 148 + pub struct RBTree<K, V> { 149 + root: bindings::rb_root, 150 + _p: PhantomData<Node<K, V>>, 151 + } 152 + 153 + // SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its 154 + // fields, so we use the same Send condition as would be used for a struct with K and V fields. 155 + unsafe impl<K: Send, V: Send> Send for RBTree<K, V> {} 156 + 157 + // SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its 158 + // fields, so we use the same Sync condition as would be used for a struct with K and V fields. 159 + unsafe impl<K: Sync, V: Sync> Sync for RBTree<K, V> {} 160 + 161 + impl<K, V> RBTree<K, V> { 162 + /// Creates a new and empty tree. 163 + pub fn new() -> Self { 164 + Self { 165 + // INVARIANT: There are no nodes in the tree, so the invariant holds vacuously. 166 + root: bindings::rb_root::default(), 167 + _p: PhantomData, 168 + } 169 + } 170 + } 171 + 172 + impl<K, V> RBTree<K, V> 173 + where 174 + K: Ord, 175 + { 176 + /// Tries to insert a new value into the tree. 177 + /// 178 + /// It overwrites a node if one already exists with the same key and returns it (containing the 179 + /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. 180 + /// 181 + /// Returns an error if it cannot allocate memory for the new node. 182 + pub fn try_create_and_insert( 183 + &mut self, 184 + key: K, 185 + value: V, 186 + flags: Flags, 187 + ) -> Result<Option<RBTreeNode<K, V>>> { 188 + Ok(self.insert(RBTreeNode::new(key, value, flags)?)) 189 + } 190 + 191 + /// Inserts a new node into the tree. 192 + /// 193 + /// It overwrites a node if one already exists with the same key and returns it (containing the 194 + /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. 195 + /// 196 + /// This function always succeeds. 197 + pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> { 198 + let node = Box::into_raw(node); 199 + // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when 200 + // the node is removed or replaced. 201 + let node_links = unsafe { addr_of_mut!((*node).links) }; 202 + 203 + // The parameters of `bindings::rb_link_node` are as follows: 204 + // - `node`: A pointer to an uninitialized node being inserted. 205 + // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be 206 + // null, and `node` will become a child of `parent` by replacing that child pointer 207 + // with a pointer to `node`. 208 + // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This 209 + // specifies which child of `parent` should hold `node` after this call. The 210 + // value of `*rb_link` must be null before the call to `rb_link_node`. If the 211 + // red/black tree is empty, then it’s also possible for `parent` to be null. In 212 + // this case, `rb_link` is a pointer to the `root` field of the red/black tree. 213 + // 214 + // We will traverse the tree looking for a node that has a null pointer as its child, 215 + // representing an empty subtree where we can insert our new node. We need to make sure 216 + // that we preserve the ordering of the nodes in the tree. In each iteration of the loop 217 + // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere 218 + // in the subtree of `parent` that `child_field_of_parent` points at. Once 219 + // we find an empty subtree, we can insert the new node using `rb_link_node`. 220 + let mut parent = core::ptr::null_mut(); 221 + let mut child_field_of_parent: &mut *mut bindings::rb_node = &mut self.root.rb_node; 222 + while !child_field_of_parent.is_null() { 223 + parent = *child_field_of_parent; 224 + 225 + // We need to determine whether `node` should be the left or right child of `parent`, 226 + // so we will compare with the `key` field of `parent` a.k.a. `this` below. 227 + // 228 + // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 229 + // point to the links field of `Node<K, V>` objects. 230 + let this = unsafe { container_of!(parent, Node<K, V>, links) }; 231 + 232 + // SAFETY: `this` is a non-null node so it is valid by the type invariants. `node` is 233 + // valid until the node is removed. 234 + match unsafe { (*node).key.cmp(&(*this).key) } { 235 + // We would like `node` to be the left child of `parent`. Move to this child to check 236 + // whether we can use it, or continue searching, at the next iteration. 237 + // 238 + // SAFETY: `parent` is a non-null node so it is valid by the type invariants. 239 + Ordering::Less => child_field_of_parent = unsafe { &mut (*parent).rb_left }, 240 + // We would like `node` to be the right child of `parent`. Move to this child to check 241 + // whether we can use it, or continue searching, at the next iteration. 242 + // 243 + // SAFETY: `parent` is a non-null node so it is valid by the type invariants. 244 + Ordering::Greater => child_field_of_parent = unsafe { &mut (*parent).rb_right }, 245 + Ordering::Equal => { 246 + // There is an existing node in the tree with this key, and that node is 247 + // `parent`. Thus, we are replacing parent with a new node. 248 + // 249 + // INVARIANT: We are replacing an existing node with a new one, which is valid. 250 + // It remains valid because we "forgot" it with `Box::into_raw`. 251 + // SAFETY: All pointers are non-null and valid. 252 + unsafe { bindings::rb_replace_node(parent, node_links, &mut self.root) }; 253 + 254 + // INVARIANT: The node is being returned and the caller may free it, however, 255 + // it was removed from the tree. So the invariants still hold. 256 + return Some(RBTreeNode { 257 + // SAFETY: `this` was a node in the tree, so it is valid. 258 + node: unsafe { Box::from_raw(this.cast_mut()) }, 259 + }); 260 + } 261 + } 262 + } 263 + 264 + // INVARIANT: We are linking in a new node, which is valid. It remains valid because we 265 + // "forgot" it with `Box::into_raw`. 266 + // SAFETY: All pointers are non-null and valid (`*child_field_of_parent` is null, but `child_field_of_parent` is a 267 + // mutable reference). 268 + unsafe { bindings::rb_link_node(node_links, parent, child_field_of_parent) }; 269 + 270 + // SAFETY: All pointers are valid. `node` has just been inserted into the tree. 271 + unsafe { bindings::rb_insert_color(node_links, &mut self.root) }; 272 + None 273 + } 274 + 275 + /// Returns a node with the given key, if one exists. 276 + fn find(&self, key: &K) -> Option<NonNull<Node<K, V>>> { 277 + let mut node = self.root.rb_node; 278 + while !node.is_null() { 279 + // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 280 + // point to the links field of `Node<K, V>` objects. 281 + let this = unsafe { container_of!(node, Node<K, V>, links) }; 282 + // SAFETY: `this` is a non-null node so it is valid by the type invariants. 283 + node = match key.cmp(unsafe { &(*this).key }) { 284 + // SAFETY: `node` is a non-null node so it is valid by the type invariants. 285 + Ordering::Less => unsafe { (*node).rb_left }, 286 + // SAFETY: `node` is a non-null node so it is valid by the type invariants. 287 + Ordering::Greater => unsafe { (*node).rb_right }, 288 + Ordering::Equal => return NonNull::new(this.cast_mut()), 289 + } 290 + } 291 + None 292 + } 293 + 294 + /// Returns a reference to the value corresponding to the key. 295 + pub fn get(&self, key: &K) -> Option<&V> { 296 + // SAFETY: The `find` return value is a node in the tree, so it is valid. 297 + self.find(key).map(|node| unsafe { &node.as_ref().value }) 298 + } 299 + 300 + /// Returns a mutable reference to the value corresponding to the key. 301 + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { 302 + // SAFETY: The `find` return value is a node in the tree, so it is valid. 303 + self.find(key) 304 + .map(|mut node| unsafe { &mut node.as_mut().value }) 305 + } 306 + 307 + /// Removes the node with the given key from the tree. 308 + /// 309 + /// It returns the node that was removed if one exists, or [`None`] otherwise. 310 + fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> { 311 + let mut node = self.find(key)?; 312 + 313 + // SAFETY: The `find` return value is a node in the tree, so it is valid. 314 + unsafe { bindings::rb_erase(&mut node.as_mut().links, &mut self.root) }; 315 + 316 + // INVARIANT: The node is being returned and the caller may free it, however, it was 317 + // removed from the tree. So the invariants still hold. 318 + Some(RBTreeNode { 319 + // SAFETY: The `find` return value was a node in the tree, so it is valid. 320 + node: unsafe { Box::from_raw(node.as_ptr()) }, 321 + }) 322 + } 323 + 324 + /// Removes the node with the given key from the tree. 325 + /// 326 + /// It returns the value that was removed if one exists, or [`None`] otherwise. 327 + pub fn remove(&mut self, key: &K) -> Option<V> { 328 + self.remove_node(key).map(|node| node.node.value) 329 + } 330 + } 331 + 332 + impl<K, V> Default for RBTree<K, V> { 333 + fn default() -> Self { 334 + Self::new() 335 + } 336 + } 337 + 338 + impl<K, V> Drop for RBTree<K, V> { 339 + fn drop(&mut self) { 340 + // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`. 341 + let mut next = unsafe { bindings::rb_first_postorder(&self.root) }; 342 + 343 + // INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid. 344 + while !next.is_null() { 345 + // SAFETY: All links fields we create are in a `Node<K, V>`. 346 + let this = unsafe { container_of!(next, Node<K, V>, links) }; 347 + 348 + // Find out what the next node is before disposing of the current one. 349 + // SAFETY: `next` and all nodes in postorder are still valid. 350 + next = unsafe { bindings::rb_next_postorder(next) }; 351 + 352 + // INVARIANT: This is the destructor, so we break the type invariant during clean-up, 353 + // but it is not observable. The loop invariant is still maintained. 354 + 355 + // SAFETY: `this` is valid per the loop invariant. 356 + unsafe { drop(Box::from_raw(this.cast_mut())) }; 357 + } 358 + } 359 + } 360 + 361 + /// A memory reservation for a red-black tree node. 362 + /// 363 + /// 364 + /// It contains the memory needed to hold a node that can be inserted into a red-black tree. One 365 + /// can be obtained by directly allocating it ([`RBTreeNodeReservation::new`]). 366 + pub struct RBTreeNodeReservation<K, V> { 367 + node: Box<MaybeUninit<Node<K, V>>>, 368 + } 369 + 370 + impl<K, V> RBTreeNodeReservation<K, V> { 371 + /// Allocates memory for a node to be eventually initialised and inserted into the tree via a 372 + /// call to [`RBTree::insert`]. 373 + pub fn new(flags: Flags) -> Result<RBTreeNodeReservation<K, V>> { 374 + Ok(RBTreeNodeReservation { 375 + node: <Box<_> as BoxExt<_>>::new_uninit(flags)?, 376 + }) 377 + } 378 + } 379 + 380 + // SAFETY: This doesn't actually contain K or V, and is just a memory allocation. Those can always 381 + // be moved across threads. 382 + unsafe impl<K, V> Send for RBTreeNodeReservation<K, V> {} 383 + 384 + // SAFETY: This doesn't actually contain K or V, and is just a memory allocation. 385 + unsafe impl<K, V> Sync for RBTreeNodeReservation<K, V> {} 386 + 387 + impl<K, V> RBTreeNodeReservation<K, V> { 388 + /// Initialises a node reservation. 389 + /// 390 + /// It then becomes an [`RBTreeNode`] that can be inserted into a tree. 391 + pub fn into_node(self, key: K, value: V) -> RBTreeNode<K, V> { 392 + let node = Box::write( 393 + self.node, 394 + Node { 395 + key, 396 + value, 397 + links: bindings::rb_node::default(), 398 + }, 399 + ); 400 + RBTreeNode { node } 401 + } 402 + } 403 + 404 + /// A red-black tree node. 405 + /// 406 + /// The node is fully initialised (with key and value) and can be inserted into a tree without any 407 + /// extra allocations or failure paths. 408 + pub struct RBTreeNode<K, V> { 409 + node: Box<Node<K, V>>, 410 + } 411 + 412 + impl<K, V> RBTreeNode<K, V> { 413 + /// Allocates and initialises a node that can be inserted into the tree via 414 + /// [`RBTree::insert`]. 415 + pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> { 416 + Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value)) 417 + } 418 + } 419 + 420 + // SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across 421 + // threads. 422 + unsafe impl<K: Send, V: Send> Send for RBTreeNode<K, V> {} 423 + 424 + // SAFETY: If K and V can be accessed without synchronization, then it's also okay to access 425 + // [`RBTreeNode`] without synchronization. 426 + unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {} 427 + 428 + struct Node<K, V> { 429 + links: bindings::rb_node, 430 + key: K, 431 + value: V, 432 + }

Configure Feed

Configure Feed