FILE
rbtree.rs
File preview with syntax highlighting.
linux-6.19.14/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 9use crate::{alloc::Flags, bindings, container_of, error::Result, prelude::*}; 10use core::{ 11 cmp::{Ord, Ordering}, 12 marker::PhantomData, 13 mem::MaybeUninit, 14 ptr::{addr_of_mut, from_mut, NonNull}, 15}; 16 17/// A red-black tree with owned nodes. 18/// 19/// It is backed by the kernel C red-black trees. 20/// 21/// # Examples 22/// 23/// In the example below we do several operations on a tree. We note that insertions may fail if 24/// the system is out of memory. 25/// 26/// ``` 27/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation}}; 28/// 29/// // Create a new tree. 30/// let mut tree = RBTree::new(); 31/// 32/// // Insert three elements. 33/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 34/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 35/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 36/// 37/// // Check the nodes we just inserted. 38/// { 39/// assert_eq!(tree.get(&10), Some(&100)); 40/// assert_eq!(tree.get(&20), Some(&200)); 41/// assert_eq!(tree.get(&30), Some(&300)); 42/// } 43/// 44/// // Iterate over the nodes we just inserted. 45/// { 46/// let mut iter = tree.iter(); 47/// assert_eq!(iter.next(), Some((&10, &100))); 48/// assert_eq!(iter.next(), Some((&20, &200))); 49/// assert_eq!(iter.next(), Some((&30, &300))); 50/// assert!(iter.next().is_none()); 51/// } 52/// 53/// // Print all elements. 54/// for (key, value) in &tree { 55/// pr_info!("{} = {}\n", key, value); 56/// } 57/// 58/// // Replace one of the elements. 59/// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?; 60/// 61/// // Check that the tree reflects the replacement. 62/// { 63/// let mut iter = tree.iter(); 64/// assert_eq!(iter.next(), Some((&10, &1000))); 65/// assert_eq!(iter.next(), Some((&20, &200))); 66/// assert_eq!(iter.next(), Some((&30, &300))); 67/// assert!(iter.next().is_none()); 68/// } 69/// 70/// // Change the value of one of the elements. 71/// *tree.get_mut(&30).unwrap() = 3000; 72/// 73/// // Check that the tree reflects the update. 74/// { 75/// let mut iter = tree.iter(); 76/// assert_eq!(iter.next(), Some((&10, &1000))); 77/// assert_eq!(iter.next(), Some((&20, &200))); 78/// assert_eq!(iter.next(), Some((&30, &3000))); 79/// assert!(iter.next().is_none()); 80/// } 81/// 82/// // Remove an element. 83/// tree.remove(&10); 84/// 85/// // Check that the tree reflects the removal. 86/// { 87/// let mut iter = tree.iter(); 88/// assert_eq!(iter.next(), Some((&20, &200))); 89/// assert_eq!(iter.next(), Some((&30, &3000))); 90/// assert!(iter.next().is_none()); 91/// } 92/// 93/// # Ok::<(), Error>(()) 94/// ``` 95/// 96/// In the example below, we first allocate a node, acquire a spinlock, then insert the node into 97/// the tree. This is useful when the insertion context does not allow sleeping, for example, when 98/// holding a spinlock. 99/// 100/// ``` 101/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode}, sync::SpinLock}; 102/// 103/// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result { 104/// // Pre-allocate node. This may fail (as it allocates memory). 105/// let node = RBTreeNode::new(10, 100, flags::GFP_KERNEL)?; 106/// 107/// // Insert node while holding the lock. It is guaranteed to succeed with no allocation 108/// // attempts. 109/// let mut guard = tree.lock(); 110/// guard.insert(node); 111/// Ok(()) 112/// } 113/// ``` 114/// 115/// In the example below, we reuse an existing node allocation from an element we removed. 116/// 117/// ``` 118/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNodeReservation}}; 119/// 120/// // Create a new tree. 121/// let mut tree = RBTree::new(); 122/// 123/// // Insert three elements. 124/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 125/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 126/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 127/// 128/// // Check the nodes we just inserted. 129/// { 130/// let mut iter = tree.iter(); 131/// assert_eq!(iter.next(), Some((&10, &100))); 132/// assert_eq!(iter.next(), Some((&20, &200))); 133/// assert_eq!(iter.next(), Some((&30, &300))); 134/// assert!(iter.next().is_none()); 135/// } 136/// 137/// // Remove a node, getting back ownership of it. 138/// let existing = tree.remove(&30); 139/// 140/// // Check that the tree reflects the removal. 141/// { 142/// let mut iter = tree.iter(); 143/// assert_eq!(iter.next(), Some((&10, &100))); 144/// assert_eq!(iter.next(), Some((&20, &200))); 145/// assert!(iter.next().is_none()); 146/// } 147/// 148/// // Create a preallocated reservation that we can re-use later. 149/// let reservation = RBTreeNodeReservation::new(flags::GFP_KERNEL)?; 150/// 151/// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to 152/// // succeed (no memory allocations). 153/// tree.insert(reservation.into_node(15, 150)); 154/// 155/// // Check that the tree reflect the new insertion. 156/// { 157/// let mut iter = tree.iter(); 158/// assert_eq!(iter.next(), Some((&10, &100))); 159/// assert_eq!(iter.next(), Some((&15, &150))); 160/// assert_eq!(iter.next(), Some((&20, &200))); 161/// assert!(iter.next().is_none()); 162/// } 163/// 164/// # Ok::<(), Error>(()) 165/// ``` 166/// 167/// # Invariants 168/// 169/// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always 170/// valid, and pointing to a field of our internal representation of a node. 171pub struct RBTree<K, V> { 172 root: bindings::rb_root, 173 _p: PhantomData<Node<K, V>>, 174} 175 176// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its 177// fields, so we use the same Send condition as would be used for a struct with K and V fields. 178unsafe impl<K: Send, V: Send> Send for RBTree<K, V> {} 179 180// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its 181// fields, so we use the same Sync condition as would be used for a struct with K and V fields. 182unsafe impl<K: Sync, V: Sync> Sync for RBTree<K, V> {} 183 184impl<K, V> RBTree<K, V> { 185 /// Creates a new and empty tree. 186 pub fn new() -> Self { 187 Self { 188 // INVARIANT: There are no nodes in the tree, so the invariant holds vacuously. 189 root: bindings::rb_root::default(), 190 _p: PhantomData, 191 } 192 } 193 194 /// Returns true if this tree is empty. 195 #[inline] 196 pub fn is_empty(&self) -> bool { 197 self.root.rb_node.is_null() 198 } 199 200 /// Returns an iterator over the tree nodes, sorted by key. 201 pub fn iter(&self) -> Iter<'_, K, V> { 202 Iter { 203 _tree: PhantomData, 204 // INVARIANT: 205 // - `self.root` is a valid pointer to a tree root. 206 // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid. 207 iter_raw: IterRaw { 208 // SAFETY: by the invariants, all pointers are valid. 209 next: unsafe { bindings::rb_first(&self.root) }, 210 _phantom: PhantomData, 211 }, 212 } 213 } 214 215 /// Returns a mutable iterator over the tree nodes, sorted by key. 216 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { 217 IterMut { 218 _tree: PhantomData, 219 // INVARIANT: 220 // - `self.root` is a valid pointer to a tree root. 221 // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid. 222 iter_raw: IterRaw { 223 // SAFETY: by the invariants, all pointers are valid. 224 next: unsafe { bindings::rb_first(from_mut(&mut self.root)) }, 225 _phantom: PhantomData, 226 }, 227 } 228 } 229 230 /// Returns an iterator over the keys of the nodes in the tree, in sorted order. 231 pub fn keys(&self) -> impl Iterator<Item = &'_ K> { 232 self.iter().map(|(k, _)| k) 233 } 234 235 /// Returns an iterator over the values of the nodes in the tree, sorted by key. 236 pub fn values(&self) -> impl Iterator<Item = &'_ V> { 237 self.iter().map(|(_, v)| v) 238 } 239 240 /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key. 241 pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> { 242 self.iter_mut().map(|(_, v)| v) 243 } 244 245 /// Returns a cursor over the tree nodes, starting with the smallest key. 246 pub fn cursor_front_mut(&mut self) -> Option<CursorMut<'_, K, V>> { 247 let root = addr_of_mut!(self.root); 248 // SAFETY: `self.root` is always a valid root node. 249 let current = unsafe { bindings::rb_first(root) }; 250 NonNull::new(current).map(|current| { 251 // INVARIANT: 252 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 253 CursorMut { 254 current, 255 tree: self, 256 } 257 }) 258 } 259 260 /// Returns an immutable cursor over the tree nodes, starting with the smallest key. 261 pub fn cursor_front(&self) -> Option<Cursor<'_, K, V>> { 262 let root = &raw const self.root; 263 // SAFETY: `self.root` is always a valid root node. 264 let current = unsafe { bindings::rb_first(root) }; 265 NonNull::new(current).map(|current| { 266 // INVARIANT: 267 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 268 Cursor { 269 current, 270 _tree: PhantomData, 271 } 272 }) 273 } 274 275 /// Returns a cursor over the tree nodes, starting with the largest key. 276 pub fn cursor_back_mut(&mut self) -> Option<CursorMut<'_, K, V>> { 277 let root = addr_of_mut!(self.root); 278 // SAFETY: `self.root` is always a valid root node. 279 let current = unsafe { bindings::rb_last(root) }; 280 NonNull::new(current).map(|current| { 281 // INVARIANT: 282 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 283 CursorMut { 284 current, 285 tree: self, 286 } 287 }) 288 } 289 290 /// Returns a cursor over the tree nodes, starting with the largest key. 291 pub fn cursor_back(&self) -> Option<Cursor<'_, K, V>> { 292 let root = &raw const self.root; 293 // SAFETY: `self.root` is always a valid root node. 294 let current = unsafe { bindings::rb_last(root) }; 295 NonNull::new(current).map(|current| { 296 // INVARIANT: 297 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 298 Cursor { 299 current, 300 _tree: PhantomData, 301 } 302 }) 303 } 304} 305 306impl<K, V> RBTree<K, V> 307where 308 K: Ord, 309{ 310 /// Tries to insert a new value into the tree. 311 /// 312 /// It overwrites a node if one already exists with the same key and returns it (containing the 313 /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. 314 /// 315 /// Returns an error if it cannot allocate memory for the new node. 316 pub fn try_create_and_insert( 317 &mut self, 318 key: K, 319 value: V, 320 flags: Flags, 321 ) -> Result<Option<RBTreeNode<K, V>>> { 322 Ok(self.insert(RBTreeNode::new(key, value, flags)?)) 323 } 324 325 /// Inserts a new node into the tree. 326 /// 327 /// It overwrites a node if one already exists with the same key and returns it (containing the 328 /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. 329 /// 330 /// This function always succeeds. 331 pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> { 332 match self.raw_entry(&node.node.key) { 333 RawEntry::Occupied(entry) => Some(entry.replace(node)), 334 RawEntry::Vacant(entry) => { 335 entry.insert(node); 336 None 337 } 338 } 339 } 340 341 fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> { 342 let raw_self: *mut RBTree<K, V> = self; 343 // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`. 344 // The parameters of `bindings::rb_link_node` are as follows: 345 // - `node`: A pointer to an uninitialized node being inserted. 346 // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be 347 // null, and `node` will become a child of `parent` by replacing that child pointer 348 // with a pointer to `node`. 349 // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This 350 // specifies which child of `parent` should hold `node` after this call. The 351 // value of `*rb_link` must be null before the call to `rb_link_node`. If the 352 // red/black tree is empty, then it’s also possible for `parent` to be null. In 353 // this case, `rb_link` is a pointer to the `root` field of the red/black tree. 354 // 355 // We will traverse the tree looking for a node that has a null pointer as its child, 356 // representing an empty subtree where we can insert our new node. We need to make sure 357 // that we preserve the ordering of the nodes in the tree. In each iteration of the loop 358 // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere 359 // in the subtree of `parent` that `child_field_of_parent` points at. Once 360 // we find an empty subtree, we can insert the new node using `rb_link_node`. 361 let mut parent = core::ptr::null_mut(); 362 let mut child_field_of_parent: &mut *mut bindings::rb_node = 363 // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above). 364 unsafe { &mut (*raw_self).root.rb_node }; 365 while !(*child_field_of_parent).is_null() { 366 let curr = *child_field_of_parent; 367 // SAFETY: All links fields we create are in a `Node<K, V>`. 368 let node = unsafe { container_of!(curr, Node<K, V>, links) }; 369 370 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 371 match key.cmp(unsafe { &(*node).key }) { 372 // SAFETY: `curr` is a non-null node so it is valid by the type invariants. 373 Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left }, 374 // SAFETY: `curr` is a non-null node so it is valid by the type invariants. 375 Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).rb_right }, 376 Ordering::Equal => { 377 return RawEntry::Occupied(OccupiedEntry { 378 rbtree: self, 379 node_links: curr, 380 }) 381 } 382 } 383 parent = curr; 384 } 385 386 RawEntry::Vacant(RawVacantEntry { 387 rbtree: raw_self, 388 parent, 389 child_field_of_parent, 390 _phantom: PhantomData, 391 }) 392 } 393 394 /// Gets the given key's corresponding entry in the map for in-place manipulation. 395 pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { 396 match self.raw_entry(&key) { 397 RawEntry::Occupied(entry) => Entry::Occupied(entry), 398 RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }), 399 } 400 } 401 402 /// Used for accessing the given node, if it exists. 403 pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> { 404 match self.raw_entry(key) { 405 RawEntry::Occupied(entry) => Some(entry), 406 RawEntry::Vacant(_entry) => None, 407 } 408 } 409 410 /// Returns a reference to the value corresponding to the key. 411 pub fn get(&self, key: &K) -> Option<&V> { 412 let mut node = self.root.rb_node; 413 while !node.is_null() { 414 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 415 // point to the links field of `Node<K, V>` objects. 416 let this = unsafe { container_of!(node, Node<K, V>, links) }; 417 // SAFETY: `this` is a non-null node so it is valid by the type invariants. 418 node = match key.cmp(unsafe { &(*this).key }) { 419 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 420 Ordering::Less => unsafe { (*node).rb_left }, 421 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 422 Ordering::Greater => unsafe { (*node).rb_right }, 423 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 424 Ordering::Equal => return Some(unsafe { &(*this).value }), 425 } 426 } 427 None 428 } 429 430 /// Returns a mutable reference to the value corresponding to the key. 431 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { 432 self.find_mut(key).map(|node| node.into_mut()) 433 } 434 435 /// Removes the node with the given key from the tree. 436 /// 437 /// It returns the node that was removed if one exists, or [`None`] otherwise. 438 pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> { 439 self.find_mut(key).map(OccupiedEntry::remove_node) 440 } 441 442 /// Removes the node with the given key from the tree. 443 /// 444 /// It returns the value that was removed if one exists, or [`None`] otherwise. 445 pub fn remove(&mut self, key: &K) -> Option<V> { 446 self.find_mut(key).map(OccupiedEntry::remove) 447 } 448 449 /// Returns a cursor over the tree nodes based on the given key. 450 /// 451 /// If the given key exists, the cursor starts there. 452 /// Otherwise it starts with the first larger key in sort order. 453 /// If there is no larger key, it returns [`None`]. 454 pub fn cursor_lower_bound_mut(&mut self, key: &K) -> Option<CursorMut<'_, K, V>> 455 where 456 K: Ord, 457 { 458 let best = self.find_best_match(key)?; 459 460 NonNull::new(best.as_ptr()).map(|current| { 461 // INVARIANT: 462 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 463 CursorMut { 464 current, 465 tree: self, 466 } 467 }) 468 } 469 470 /// Returns a cursor over the tree nodes based on the given key. 471 /// 472 /// If the given key exists, the cursor starts there. 473 /// Otherwise it starts with the first larger key in sort order. 474 /// If there is no larger key, it returns [`None`]. 475 pub fn cursor_lower_bound(&self, key: &K) -> Option<Cursor<'_, K, V>> 476 where 477 K: Ord, 478 { 479 let best = self.find_best_match(key)?; 480 481 NonNull::new(best.as_ptr()).map(|current| { 482 // INVARIANT: 483 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 484 Cursor { 485 current, 486 _tree: PhantomData, 487 } 488 }) 489 } 490 491 fn find_best_match(&self, key: &K) -> Option<NonNull<bindings::rb_node>> { 492 let mut node = self.root.rb_node; 493 let mut best_key: Option<&K> = None; 494 let mut best_links: Option<NonNull<bindings::rb_node>> = None; 495 while !node.is_null() { 496 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 497 // point to the links field of `Node<K, V>` objects. 498 let this = unsafe { container_of!(node, Node<K, V>, links) }; 499 // SAFETY: `this` is a non-null node so it is valid by the type invariants. 500 let this_key = unsafe { &(*this).key }; 501 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 502 let left_child = unsafe { (*node).rb_left }; 503 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 504 let right_child = unsafe { (*node).rb_right }; 505 match key.cmp(this_key) { 506 Ordering::Equal => { 507 // SAFETY: `this` is a non-null node so it is valid by the type invariants. 508 best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) }); 509 break; 510 } 511 Ordering::Greater => { 512 node = right_child; 513 } 514 Ordering::Less => { 515 let is_better_match = match best_key { 516 None => true, 517 Some(best) => best > this_key, 518 }; 519 if is_better_match { 520 best_key = Some(this_key); 521 // SAFETY: `this` is a non-null node so it is valid by the type invariants. 522 best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) }); 523 } 524 node = left_child; 525 } 526 }; 527 } 528 best_links 529 } 530} 531 532impl<K, V> Default for RBTree<K, V> { 533 fn default() -> Self { 534 Self::new() 535 } 536} 537 538impl<K, V> Drop for RBTree<K, V> { 539 fn drop(&mut self) { 540 // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`. 541 let mut next = unsafe { bindings::rb_first_postorder(&self.root) }; 542 543 // INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid. 544 while !next.is_null() { 545 // SAFETY: All links fields we create are in a `Node<K, V>`. 546 let this = unsafe { container_of!(next, Node<K, V>, links) }; 547 548 // Find out what the next node is before disposing of the current one. 549 // SAFETY: `next` and all nodes in postorder are still valid. 550 next = unsafe { bindings::rb_next_postorder(next) }; 551 552 // INVARIANT: This is the destructor, so we break the type invariant during clean-up, 553 // but it is not observable. The loop invariant is still maintained. 554 555 // SAFETY: `this` is valid per the loop invariant. 556 unsafe { drop(KBox::from_raw(this)) }; 557 } 558 } 559} 560 561/// A bidirectional mutable cursor over the tree nodes, sorted by key. 562/// 563/// # Examples 564/// 565/// In the following example, we obtain a cursor to the first element in the tree. 566/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree. 567/// 568/// ``` 569/// use kernel::{alloc::flags, rbtree::RBTree}; 570/// 571/// // Create a new tree. 572/// let mut tree = RBTree::new(); 573/// 574/// // Insert three elements. 575/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 576/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 577/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 578/// 579/// // Get a cursor to the first element. 580/// let mut cursor = tree.cursor_front_mut().unwrap(); 581/// let mut current = cursor.current(); 582/// assert_eq!(current, (&10, &100)); 583/// 584/// // Move the cursor, updating it to the 2nd element. 585/// cursor = cursor.move_next().unwrap(); 586/// current = cursor.current(); 587/// assert_eq!(current, (&20, &200)); 588/// 589/// // Peek at the next element without impacting the cursor. 590/// let next = cursor.peek_next().unwrap(); 591/// assert_eq!(next, (&30, &300)); 592/// current = cursor.current(); 593/// assert_eq!(current, (&20, &200)); 594/// 595/// // Moving past the last element causes the cursor to return [`None`]. 596/// cursor = cursor.move_next().unwrap(); 597/// current = cursor.current(); 598/// assert_eq!(current, (&30, &300)); 599/// let cursor = cursor.move_next(); 600/// assert!(cursor.is_none()); 601/// 602/// # Ok::<(), Error>(()) 603/// ``` 604/// 605/// A cursor can also be obtained at the last element in the tree. 606/// 607/// ``` 608/// use kernel::{alloc::flags, rbtree::RBTree}; 609/// 610/// // Create a new tree. 611/// let mut tree = RBTree::new(); 612/// 613/// // Insert three elements. 614/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 615/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 616/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 617/// 618/// let mut cursor = tree.cursor_back_mut().unwrap(); 619/// let current = cursor.current(); 620/// assert_eq!(current, (&30, &300)); 621/// 622/// # Ok::<(), Error>(()) 623/// ``` 624/// 625/// Obtaining a cursor returns [`None`] if the tree is empty. 626/// 627/// ``` 628/// use kernel::rbtree::RBTree; 629/// 630/// let mut tree: RBTree<u16, u16> = RBTree::new(); 631/// assert!(tree.cursor_front_mut().is_none()); 632/// 633/// # Ok::<(), Error>(()) 634/// ``` 635/// 636/// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary node in the tree. 637/// 638/// ``` 639/// use kernel::{alloc::flags, rbtree::RBTree}; 640/// 641/// // Create a new tree. 642/// let mut tree = RBTree::new(); 643/// 644/// // Insert five elements. 645/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 646/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 647/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 648/// tree.try_create_and_insert(40, 400, flags::GFP_KERNEL)?; 649/// tree.try_create_and_insert(50, 500, flags::GFP_KERNEL)?; 650/// 651/// // If the provided key exists, a cursor to that key is returned. 652/// let cursor = tree.cursor_lower_bound(&20).unwrap(); 653/// let current = cursor.current(); 654/// assert_eq!(current, (&20, &200)); 655/// 656/// // If the provided key doesn't exist, a cursor to the first larger element in sort order is returned. 657/// let cursor = tree.cursor_lower_bound(&25).unwrap(); 658/// let current = cursor.current(); 659/// assert_eq!(current, (&30, &300)); 660/// 661/// // If there is no larger key, [`None`] is returned. 662/// let cursor = tree.cursor_lower_bound(&55); 663/// assert!(cursor.is_none()); 664/// 665/// # Ok::<(), Error>(()) 666/// ``` 667/// 668/// The cursor allows mutation of values in the tree. 669/// 670/// ``` 671/// use kernel::{alloc::flags, rbtree::RBTree}; 672/// 673/// // Create a new tree. 674/// let mut tree = RBTree::new(); 675/// 676/// // Insert three elements. 677/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 678/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 679/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 680/// 681/// // Retrieve a cursor. 682/// let mut cursor = tree.cursor_front_mut().unwrap(); 683/// 684/// // Get a mutable reference to the current value. 685/// let (k, v) = cursor.current_mut(); 686/// *v = 1000; 687/// 688/// // The updated value is reflected in the tree. 689/// let updated = tree.get(&10).unwrap(); 690/// assert_eq!(updated, &1000); 691/// 692/// # Ok::<(), Error>(()) 693/// ``` 694/// 695/// It also allows node removal. The following examples demonstrate the behavior of removing the current node. 696/// 697/// ``` 698/// use kernel::{alloc::flags, rbtree::RBTree}; 699/// 700/// // Create a new tree. 701/// let mut tree = RBTree::new(); 702/// 703/// // Insert three elements. 704/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 705/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 706/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 707/// 708/// // Remove the first element. 709/// let mut cursor = tree.cursor_front_mut().unwrap(); 710/// let mut current = cursor.current(); 711/// assert_eq!(current, (&10, &100)); 712/// cursor = cursor.remove_current().0.unwrap(); 713/// 714/// // If a node exists after the current element, it is returned. 715/// current = cursor.current(); 716/// assert_eq!(current, (&20, &200)); 717/// 718/// // Get a cursor to the last element, and remove it. 719/// cursor = tree.cursor_back_mut().unwrap(); 720/// current = cursor.current(); 721/// assert_eq!(current, (&30, &300)); 722/// 723/// // Since there is no next node, the previous node is returned. 724/// cursor = cursor.remove_current().0.unwrap(); 725/// current = cursor.current(); 726/// assert_eq!(current, (&20, &200)); 727/// 728/// // Removing the last element in the tree returns [`None`]. 729/// assert!(cursor.remove_current().0.is_none()); 730/// 731/// # Ok::<(), Error>(()) 732/// ``` 733/// 734/// Nodes adjacent to the current node can also be removed. 735/// 736/// ``` 737/// use kernel::{alloc::flags, rbtree::RBTree}; 738/// 739/// // Create a new tree. 740/// let mut tree = RBTree::new(); 741/// 742/// // Insert three elements. 743/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 744/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 745/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 746/// 747/// // Get a cursor to the first element. 748/// let mut cursor = tree.cursor_front_mut().unwrap(); 749/// let mut current = cursor.current(); 750/// assert_eq!(current, (&10, &100)); 751/// 752/// // Calling `remove_prev` from the first element returns [`None`]. 753/// assert!(cursor.remove_prev().is_none()); 754/// 755/// // Get a cursor to the last element. 756/// cursor = tree.cursor_back_mut().unwrap(); 757/// current = cursor.current(); 758/// assert_eq!(current, (&30, &300)); 759/// 760/// // Calling `remove_prev` removes and returns the middle element. 761/// assert_eq!(cursor.remove_prev().unwrap().to_key_value(), (20, 200)); 762/// 763/// // Calling `remove_next` from the last element returns [`None`]. 764/// assert!(cursor.remove_next().is_none()); 765/// 766/// // Move to the first element 767/// cursor = cursor.move_prev().unwrap(); 768/// current = cursor.current(); 769/// assert_eq!(current, (&10, &100)); 770/// 771/// // Calling `remove_next` removes and returns the last element. 772/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300)); 773/// 774/// # Ok::<(), Error>(()) 775/// 776/// ``` 777/// 778/// # Invariants 779/// - `current` points to a node that is in the same [`RBTree`] as `tree`. 780pub struct CursorMut<'a, K, V> { 781 tree: &'a mut RBTree<K, V>, 782 current: NonNull<bindings::rb_node>, 783} 784 785/// A bidirectional immutable cursor over the tree nodes, sorted by key. This is a simpler 786/// variant of [`CursorMut`] that is basically providing read only access. 787/// 788/// # Examples 789/// 790/// In the following example, we obtain a cursor to the first element in the tree. 791/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree. 792/// 793/// ``` 794/// use kernel::{alloc::flags, rbtree::RBTree}; 795/// 796/// // Create a new tree. 797/// let mut tree = RBTree::new(); 798/// 799/// // Insert three elements. 800/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 801/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 802/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 803/// 804/// // Get a cursor to the first element. 805/// let cursor = tree.cursor_front().unwrap(); 806/// let current = cursor.current(); 807/// assert_eq!(current, (&10, &100)); 808/// 809/// # Ok::<(), Error>(()) 810/// ``` 811pub struct Cursor<'a, K, V> { 812 _tree: PhantomData<&'a RBTree<K, V>>, 813 current: NonNull<bindings::rb_node>, 814} 815 816// SAFETY: The immutable cursor gives out shared access to `K` and `V` so if `K` and `V` can be 817// shared across threads, then it's safe to share the cursor. 818unsafe impl<'a, K: Sync, V: Sync> Send for Cursor<'a, K, V> {} 819 820// SAFETY: The immutable cursor gives out shared access to `K` and `V` so if `K` and `V` can be 821// shared across threads, then it's safe to share the cursor. 822unsafe impl<'a, K: Sync, V: Sync> Sync for Cursor<'a, K, V> {} 823 824impl<'a, K, V> Cursor<'a, K, V> { 825 /// The current node 826 pub fn current(&self) -> (&K, &V) { 827 // SAFETY: 828 // - `self.current` is a valid node by the type invariants. 829 // - We have an immutable reference by the function signature. 830 unsafe { Self::to_key_value(self.current) } 831 } 832 833 /// # Safety 834 /// 835 /// - `node` must be a valid pointer to a node in an [`RBTree`]. 836 /// - The caller has immutable access to `node` for the duration of `'b`. 837 unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) { 838 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 839 // point to the links field of `Node<K, V>` objects. 840 let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) }; 841 // SAFETY: The passed `node` is the current node or a non-null neighbor, 842 // thus `this` is valid by the type invariants. 843 let k = unsafe { &(*this).key }; 844 // SAFETY: The passed `node` is the current node or a non-null neighbor, 845 // thus `this` is valid by the type invariants. 846 let v = unsafe { &(*this).value }; 847 (k, v) 848 } 849 850 /// Access the previous node without moving the cursor. 851 pub fn peek_prev(&self) -> Option<(&K, &V)> { 852 self.peek(Direction::Prev) 853 } 854 855 /// Access the next node without moving the cursor. 856 pub fn peek_next(&self) -> Option<(&K, &V)> { 857 self.peek(Direction::Next) 858 } 859 860 fn peek(&self, direction: Direction) -> Option<(&K, &V)> { 861 self.get_neighbor_raw(direction).map(|neighbor| { 862 // SAFETY: 863 // - `neighbor` is a valid tree node. 864 // - By the function signature, we have an immutable reference to `self`. 865 unsafe { Self::to_key_value(neighbor) } 866 }) 867 } 868 869 fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> { 870 // SAFETY: `self.current` is valid by the type invariants. 871 let neighbor = unsafe { 872 match direction { 873 Direction::Prev => bindings::rb_prev(self.current.as_ptr()), 874 Direction::Next => bindings::rb_next(self.current.as_ptr()), 875 } 876 }; 877 878 NonNull::new(neighbor) 879 } 880} 881 882// SAFETY: The [`CursorMut`] has exclusive access to both `K` and `V`, so it is sufficient to 883// require them to be `Send`. 884// The cursor only gives out immutable references to the keys, but since it has exclusive access to 885// those same keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the 886// user. 887unsafe impl<'a, K: Send, V: Send> Send for CursorMut<'a, K, V> {} 888 889// SAFETY: The [`CursorMut`] gives out immutable references to `K` and mutable references to `V`, 890// so it has the same thread safety requirements as mutable references. 891unsafe impl<'a, K: Sync, V: Sync> Sync for CursorMut<'a, K, V> {} 892 893impl<'a, K, V> CursorMut<'a, K, V> { 894 /// The current node. 895 pub fn current(&self) -> (&K, &V) { 896 // SAFETY: 897 // - `self.current` is a valid node by the type invariants. 898 // - We have an immutable reference by the function signature. 899 unsafe { Self::to_key_value(self.current) } 900 } 901 902 /// The current node, with a mutable value 903 pub fn current_mut(&mut self) -> (&K, &mut V) { 904 // SAFETY: 905 // - `self.current` is a valid node by the type invariants. 906 // - We have an mutable reference by the function signature. 907 unsafe { Self::to_key_value_mut(self.current) } 908 } 909 910 /// Remove the current node from the tree. 911 /// 912 /// Returns a tuple where the first element is a cursor to the next node, if it exists, 913 /// else the previous node, else [`None`] (if the tree becomes empty). The second element 914 /// is the removed node. 915 pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) { 916 let prev = self.get_neighbor_raw(Direction::Prev); 917 let next = self.get_neighbor_raw(Direction::Next); 918 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 919 // point to the links field of `Node<K, V>` objects. 920 let this = unsafe { container_of!(self.current.as_ptr(), Node<K, V>, links) }; 921 // SAFETY: `this` is valid by the type invariants as described above. 922 let node = unsafe { KBox::from_raw(this) }; 923 let node = RBTreeNode { node }; 924 // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so 925 // the tree cannot change. By the tree invariant, all nodes are valid. 926 unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.tree.root)) }; 927 928 // INVARIANT: 929 // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`. 930 let cursor = next.or(prev).map(|current| Self { 931 current, 932 tree: self.tree, 933 }); 934 935 (cursor, node) 936 } 937 938 /// Remove the previous node, returning it if it exists. 939 pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> { 940 self.remove_neighbor(Direction::Prev) 941 } 942 943 /// Remove the next node, returning it if it exists. 944 pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> { 945 self.remove_neighbor(Direction::Next) 946 } 947 948 fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> { 949 if let Some(neighbor) = self.get_neighbor_raw(direction) { 950 let neighbor = neighbor.as_ptr(); 951 // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so 952 // the tree cannot change. By the tree invariant, all nodes are valid. 953 unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.root)) }; 954 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 955 // point to the links field of `Node<K, V>` objects. 956 let this = unsafe { container_of!(neighbor, Node<K, V>, links) }; 957 // SAFETY: `this` is valid by the type invariants as described above. 958 let node = unsafe { KBox::from_raw(this) }; 959 return Some(RBTreeNode { node }); 960 } 961 None 962 } 963 964 /// Move the cursor to the previous node, returning [`None`] if it doesn't exist. 965 pub fn move_prev(self) -> Option<Self> { 966 self.mv(Direction::Prev) 967 } 968 969 /// Move the cursor to the next node, returning [`None`] if it doesn't exist. 970 pub fn move_next(self) -> Option<Self> { 971 self.mv(Direction::Next) 972 } 973 974 fn mv(self, direction: Direction) -> Option<Self> { 975 // INVARIANT: 976 // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`. 977 self.get_neighbor_raw(direction).map(|neighbor| Self { 978 tree: self.tree, 979 current: neighbor, 980 }) 981 } 982 983 /// Access the previous node without moving the cursor. 984 pub fn peek_prev(&self) -> Option<(&K, &V)> { 985 self.peek(Direction::Prev) 986 } 987 988 /// Access the next node without moving the cursor. 989 pub fn peek_next(&self) -> Option<(&K, &V)> { 990 self.peek(Direction::Next) 991 } 992 993 fn peek(&self, direction: Direction) -> Option<(&K, &V)> { 994 self.get_neighbor_raw(direction).map(|neighbor| { 995 // SAFETY: 996 // - `neighbor` is a valid tree node. 997 // - By the function signature, we have an immutable reference to `self`. 998 unsafe { Self::to_key_value(neighbor) } 999 }) 1000 } 1001 1002 /// Access the previous node mutably without moving the cursor. 1003 pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> { 1004 self.peek_mut(Direction::Prev) 1005 } 1006 1007 /// Access the next node mutably without moving the cursor. 1008 pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> { 1009 self.peek_mut(Direction::Next) 1010 } 1011 1012 fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> { 1013 self.get_neighbor_raw(direction).map(|neighbor| { 1014 // SAFETY: 1015 // - `neighbor` is a valid tree node. 1016 // - By the function signature, we have a mutable reference to `self`. 1017 unsafe { Self::to_key_value_mut(neighbor) } 1018 }) 1019 } 1020 1021 fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> { 1022 // SAFETY: `self.current` is valid by the type invariants. 1023 let neighbor = unsafe { 1024 match direction { 1025 Direction::Prev => bindings::rb_prev(self.current.as_ptr()), 1026 Direction::Next => bindings::rb_next(self.current.as_ptr()), 1027 } 1028 }; 1029 1030 NonNull::new(neighbor) 1031 } 1032 1033 /// # Safety 1034 /// 1035 /// - `node` must be a valid pointer to a node in an [`RBTree`]. 1036 /// - The caller has immutable access to `node` for the duration of `'b`. 1037 unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) { 1038 // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`. 1039 let (k, v) = unsafe { Self::to_key_value_raw(node) }; 1040 // SAFETY: the caller guarantees immutable access to `node`. 1041 (k, unsafe { &*v }) 1042 } 1043 1044 /// # Safety 1045 /// 1046 /// - `node` must be a valid pointer to a node in an [`RBTree`]. 1047 /// - The caller has mutable access to `node` for the duration of `'b`. 1048 unsafe fn to_key_value_mut<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b mut V) { 1049 // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`. 1050 let (k, v) = unsafe { Self::to_key_value_raw(node) }; 1051 // SAFETY: the caller guarantees mutable access to `node`. 1052 (k, unsafe { &mut *v }) 1053 } 1054 1055 /// # Safety 1056 /// 1057 /// - `node` must be a valid pointer to a node in an [`RBTree`]. 1058 /// - The caller has immutable access to the key for the duration of `'b`. 1059 unsafe fn to_key_value_raw<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, *mut V) { 1060 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 1061 // point to the links field of `Node<K, V>` objects. 1062 let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) }; 1063 // SAFETY: The passed `node` is the current node or a non-null neighbor, 1064 // thus `this` is valid by the type invariants. 1065 let k = unsafe { &(*this).key }; 1066 // SAFETY: The passed `node` is the current node or a non-null neighbor, 1067 // thus `this` is valid by the type invariants. 1068 let v = unsafe { addr_of_mut!((*this).value) }; 1069 (k, v) 1070 } 1071} 1072 1073/// Direction for [`Cursor`] and [`CursorMut`] operations. 1074enum Direction { 1075 /// the node immediately before, in sort order 1076 Prev, 1077 /// the node immediately after, in sort order 1078 Next, 1079} 1080 1081impl<'a, K, V> IntoIterator for &'a RBTree<K, V> { 1082 type Item = (&'a K, &'a V); 1083 type IntoIter = Iter<'a, K, V>; 1084 1085 fn into_iter(self) -> Self::IntoIter { 1086 self.iter() 1087 } 1088} 1089 1090/// An iterator over the nodes of a [`RBTree`]. 1091/// 1092/// Instances are created by calling [`RBTree::iter`]. 1093pub struct Iter<'a, K, V> { 1094 _tree: PhantomData<&'a RBTree<K, V>>, 1095 iter_raw: IterRaw<K, V>, 1096} 1097 1098// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same 1099// thread safety requirements as immutable references. 1100unsafe impl<'a, K: Sync, V: Sync> Send for Iter<'a, K, V> {} 1101 1102// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same 1103// thread safety requirements as immutable references. 1104unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {} 1105 1106impl<'a, K, V> Iterator for Iter<'a, K, V> { 1107 type Item = (&'a K, &'a V); 1108 1109 fn next(&mut self) -> Option<Self::Item> { 1110 // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`. 1111 self.iter_raw.next().map(|(k, v)| unsafe { (&*k, &*v) }) 1112 } 1113} 1114 1115impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> { 1116 type Item = (&'a K, &'a mut V); 1117 type IntoIter = IterMut<'a, K, V>; 1118 1119 fn into_iter(self) -> Self::IntoIter { 1120 self.iter_mut() 1121 } 1122} 1123 1124/// A mutable iterator over the nodes of a [`RBTree`]. 1125/// 1126/// Instances are created by calling [`RBTree::iter_mut`]. 1127pub struct IterMut<'a, K, V> { 1128 _tree: PhantomData<&'a mut RBTree<K, V>>, 1129 iter_raw: IterRaw<K, V>, 1130} 1131 1132// SAFETY: The [`IterMut`] has exclusive access to both `K` and `V`, so it is sufficient to require them to be `Send`. 1133// The iterator only gives out immutable references to the keys, but since the iterator has exclusive access to those same 1134// keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the user. 1135unsafe impl<'a, K: Send, V: Send> Send for IterMut<'a, K, V> {} 1136 1137// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same 1138// thread safety requirements as mutable references. 1139unsafe impl<'a, K: Sync, V: Sync> Sync for IterMut<'a, K, V> {} 1140 1141impl<'a, K, V> Iterator for IterMut<'a, K, V> { 1142 type Item = (&'a K, &'a mut V); 1143 1144 fn next(&mut self) -> Option<Self::Item> { 1145 self.iter_raw.next().map(|(k, v)| 1146 // SAFETY: Due to `&mut self`, we have exclusive access to `k` and `v`, for the lifetime of `'a`. 1147 unsafe { (&*k, &mut *v) }) 1148 } 1149} 1150 1151/// A raw iterator over the nodes of a [`RBTree`]. 1152/// 1153/// # Invariants 1154/// - `self.next` is a valid pointer. 1155/// - `self.next` points to a node stored inside of a valid `RBTree`. 1156struct IterRaw<K, V> { 1157 next: *mut bindings::rb_node, 1158 _phantom: PhantomData<fn() -> (K, V)>, 1159} 1160 1161impl<K, V> Iterator for IterRaw<K, V> { 1162 type Item = (*mut K, *mut V); 1163 1164 fn next(&mut self) -> Option<Self::Item> { 1165 if self.next.is_null() { 1166 return None; 1167 } 1168 1169 // SAFETY: By the type invariant of `IterRaw`, `self.next` is a valid node in an `RBTree`, 1170 // and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects. 1171 let cur = unsafe { container_of!(self.next, Node<K, V>, links) }; 1172 1173 // SAFETY: `self.next` is a valid tree node by the type invariants. 1174 self.next = unsafe { bindings::rb_next(self.next) }; 1175 1176 // SAFETY: By the same reasoning above, it is safe to dereference the node. 1177 Some(unsafe { (addr_of_mut!((*cur).key), addr_of_mut!((*cur).value)) }) 1178 } 1179} 1180 1181/// A memory reservation for a red-black tree node. 1182/// 1183/// 1184/// It contains the memory needed to hold a node that can be inserted into a red-black tree. One 1185/// can be obtained by directly allocating it ([`RBTreeNodeReservation::new`]). 1186pub struct RBTreeNodeReservation<K, V> { 1187 node: KBox<MaybeUninit<Node<K, V>>>, 1188} 1189 1190impl<K, V> RBTreeNodeReservation<K, V> { 1191 /// Allocates memory for a node to be eventually initialised and inserted into the tree via a 1192 /// call to [`RBTree::insert`]. 1193 pub fn new(flags: Flags) -> Result<RBTreeNodeReservation<K, V>> { 1194 Ok(RBTreeNodeReservation { 1195 node: KBox::new_uninit(flags)?, 1196 }) 1197 } 1198} 1199 1200// SAFETY: This doesn't actually contain K or V, and is just a memory allocation. Those can always 1201// be moved across threads. 1202unsafe impl<K, V> Send for RBTreeNodeReservation<K, V> {} 1203 1204// SAFETY: This doesn't actually contain K or V, and is just a memory allocation. 1205unsafe impl<K, V> Sync for RBTreeNodeReservation<K, V> {} 1206 1207impl<K, V> RBTreeNodeReservation<K, V> { 1208 /// Initialises a node reservation. 1209 /// 1210 /// It then becomes an [`RBTreeNode`] that can be inserted into a tree. 1211 pub fn into_node(self, key: K, value: V) -> RBTreeNode<K, V> { 1212 let node = KBox::write( 1213 self.node, 1214 Node { 1215 key, 1216 value, 1217 links: bindings::rb_node::default(), 1218 }, 1219 ); 1220 RBTreeNode { node } 1221 } 1222} 1223 1224/// A red-black tree node. 1225/// 1226/// The node is fully initialised (with key and value) and can be inserted into a tree without any 1227/// extra allocations or failure paths. 1228pub struct RBTreeNode<K, V> { 1229 node: KBox<Node<K, V>>, 1230} 1231 1232impl<K, V> RBTreeNode<K, V> { 1233 /// Allocates and initialises a node that can be inserted into the tree via 1234 /// [`RBTree::insert`]. 1235 pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> { 1236 Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value)) 1237 } 1238 1239 /// Get the key and value from inside the node. 1240 pub fn to_key_value(self) -> (K, V) { 1241 let node = KBox::into_inner(self.node); 1242 1243 (node.key, node.value) 1244 } 1245} 1246 1247// SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across 1248// threads. 1249unsafe impl<K: Send, V: Send> Send for RBTreeNode<K, V> {} 1250 1251// SAFETY: If K and V can be accessed without synchronization, then it's also okay to access 1252// [`RBTreeNode`] without synchronization. 1253unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {} 1254 1255impl<K, V> RBTreeNode<K, V> { 1256 /// Drop the key and value, but keep the allocation. 1257 /// 1258 /// It then becomes a reservation that can be re-initialised into a different node (i.e., with 1259 /// a different key and/or value). 1260 /// 1261 /// The existing key and value are dropped in-place as part of this operation, that is, memory 1262 /// may be freed (but only for the key/value; memory for the node itself is kept for reuse). 1263 pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> { 1264 RBTreeNodeReservation { 1265 node: KBox::drop_contents(self.node), 1266 } 1267 } 1268} 1269 1270/// A view into a single entry in a map, which may either be vacant or occupied. 1271/// 1272/// This enum is constructed from the [`RBTree::entry`]. 1273/// 1274/// [`entry`]: fn@RBTree::entry 1275pub enum Entry<'a, K, V> { 1276 /// This [`RBTree`] does not have a node with this key. 1277 Vacant(VacantEntry<'a, K, V>), 1278 /// This [`RBTree`] already has a node with this key. 1279 Occupied(OccupiedEntry<'a, K, V>), 1280} 1281 1282/// Like [`Entry`], except that it doesn't have ownership of the key. 1283enum RawEntry<'a, K, V> { 1284 Vacant(RawVacantEntry<'a, K, V>), 1285 Occupied(OccupiedEntry<'a, K, V>), 1286} 1287 1288/// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`] enum. 1289pub struct VacantEntry<'a, K, V> { 1290 key: K, 1291 raw: RawVacantEntry<'a, K, V>, 1292} 1293 1294/// Like [`VacantEntry`], but doesn't hold on to the key. 1295/// 1296/// # Invariants 1297/// - `parent` may be null if the new node becomes the root. 1298/// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `parent` is 1299/// null, it is a pointer to the root of the [`RBTree`]. 1300struct RawVacantEntry<'a, K, V> { 1301 rbtree: *mut RBTree<K, V>, 1302 /// The node that will become the parent of the new node if we insert one. 1303 parent: *mut bindings::rb_node, 1304 /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is 1305 /// null. 1306 child_field_of_parent: *mut *mut bindings::rb_node, 1307 _phantom: PhantomData<&'a mut RBTree<K, V>>, 1308} 1309 1310impl<'a, K, V> RawVacantEntry<'a, K, V> { 1311 /// Inserts the given node into the [`RBTree`] at this entry. 1312 /// 1313 /// The `node` must have a key such that inserting it here does not break the ordering of this 1314 /// [`RBTree`]. 1315 fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V { 1316 let node = KBox::into_raw(node.node); 1317 1318 // SAFETY: `node` is valid at least until we call `KBox::from_raw`, which only happens when 1319 // the node is removed or replaced. 1320 let node_links = unsafe { addr_of_mut!((*node).links) }; 1321 1322 // INVARIANT: We are linking in a new node, which is valid. It remains valid because we 1323 // "forgot" it with `KBox::into_raw`. 1324 // SAFETY: The type invariants of `RawVacantEntry` are exactly the safety requirements of `rb_link_node`. 1325 unsafe { bindings::rb_link_node(node_links, self.parent, self.child_field_of_parent) }; 1326 1327 // SAFETY: All pointers are valid. `node` has just been inserted into the tree. 1328 unsafe { bindings::rb_insert_color(node_links, addr_of_mut!((*self.rbtree).root)) }; 1329 1330 // SAFETY: The node is valid until we remove it from the tree. 1331 unsafe { &mut (*node).value } 1332 } 1333} 1334 1335impl<'a, K, V> VacantEntry<'a, K, V> { 1336 /// Inserts the given node into the [`RBTree`] at this entry. 1337 pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V { 1338 self.raw.insert(reservation.into_node(self.key, value)) 1339 } 1340} 1341 1342/// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum. 1343/// 1344/// # Invariants 1345/// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree` 1346pub struct OccupiedEntry<'a, K, V> { 1347 rbtree: &'a mut RBTree<K, V>, 1348 /// The node that this entry corresponds to. 1349 node_links: *mut bindings::rb_node, 1350} 1351 1352impl<'a, K, V> OccupiedEntry<'a, K, V> { 1353 /// Gets a reference to the value in the entry. 1354 pub fn get(&self) -> &V { 1355 // SAFETY: 1356 // - `self.node_links` is a valid pointer to a node in the tree. 1357 // - We have shared access to the underlying tree, and can thus give out a shared reference. 1358 unsafe { &(*container_of!(self.node_links, Node<K, V>, links)).value } 1359 } 1360 1361 /// Gets a mutable reference to the value in the entry. 1362 pub fn get_mut(&mut self) -> &mut V { 1363 // SAFETY: 1364 // - `self.node_links` is a valid pointer to a node in the tree. 1365 // - We have exclusive access to the underlying tree, and can thus give out a mutable reference. 1366 unsafe { &mut (*(container_of!(self.node_links, Node<K, V>, links))).value } 1367 } 1368 1369 /// Converts the entry into a mutable reference to its value. 1370 /// 1371 /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`]. 1372 pub fn into_mut(self) -> &'a mut V { 1373 // SAFETY: 1374 // - `self.node_links` is a valid pointer to a node in the tree. 1375 // - This consumes the `&'a mut RBTree<K, V>`, therefore it can give out a mutable reference that lives for `'a`. 1376 unsafe { &mut (*(container_of!(self.node_links, Node<K, V>, links))).value } 1377 } 1378 1379 /// Remove this entry from the [`RBTree`]. 1380 pub fn remove_node(self) -> RBTreeNode<K, V> { 1381 // SAFETY: The node is a node in the tree, so it is valid. 1382 unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.root) }; 1383 1384 // INVARIANT: The node is being returned and the caller may free it, however, it was 1385 // removed from the tree. So the invariants still hold. 1386 RBTreeNode { 1387 // SAFETY: The node was a node in the tree, but we removed it, so we can convert it 1388 // back into a box. 1389 node: unsafe { KBox::from_raw(container_of!(self.node_links, Node<K, V>, links)) }, 1390 } 1391 } 1392 1393 /// Takes the value of the entry out of the map, and returns it. 1394 pub fn remove(self) -> V { 1395 let rb_node = self.remove_node(); 1396 let node = KBox::into_inner(rb_node.node); 1397 1398 node.value 1399 } 1400 1401 /// Swap the current node for the provided node. 1402 /// 1403 /// The key of both nodes must be equal. 1404 fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> { 1405 let node = KBox::into_raw(node.node); 1406 1407 // SAFETY: `node` is valid at least until we call `KBox::from_raw`, which only happens when 1408 // the node is removed or replaced. 1409 let new_node_links = unsafe { addr_of_mut!((*node).links) }; 1410 1411 // SAFETY: This updates the pointers so that `new_node_links` is in the tree where 1412 // `self.node_links` used to be. 1413 unsafe { 1414 bindings::rb_replace_node(self.node_links, new_node_links, &mut self.rbtree.root) 1415 }; 1416 1417 // SAFETY: 1418 // - `self.node_ptr` produces a valid pointer to a node in the tree. 1419 // - Now that we removed this entry from the tree, we can convert the node to a box. 1420 let old_node = unsafe { KBox::from_raw(container_of!(self.node_links, Node<K, V>, links)) }; 1421 1422 RBTreeNode { node: old_node } 1423 } 1424} 1425 1426struct Node<K, V> { 1427 links: bindings::rb_node, 1428 key: K, 1429 value: V, 1430}
~/
Read-only file preview.