1 // SPDX-License-Identifier: GPL-2.0 2 3 //! Red-black trees. 4 //! 5 //! C header: [`include/linux/rbtree.h`](srctr 6 //! 7 //! Reference: <https://docs.kernel.org/core-a 8 9 use crate::{alloc::Flags, bindings, container_ 10 use alloc::boxed::Box; 11 use core::{ 12 cmp::{Ord, Ordering}, 13 marker::PhantomData, 14 mem::MaybeUninit, 15 ptr::{addr_of_mut, from_mut, NonNull}, 16 }; 17 18 /// A red-black tree with owned nodes. 19 /// 20 /// It is backed by the kernel C red-black tre 21 /// 22 /// # Examples 23 /// 24 /// In the example below we do several operati 25 /// the system is out of memory. 26 /// 27 /// ``` 28 /// use kernel::{alloc::flags, rbtree::{RBTree 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: 35 /// tree.try_create_and_insert(10, 100, flags: 36 /// tree.try_create_and_insert(30, 300, flags: 37 /// 38 /// // Check the nodes we just inserted. 39 /// { 40 /// assert_eq!(tree.get(&10).unwrap(), &10 41 /// assert_eq!(tree.get(&20).unwrap(), &20 42 /// assert_eq!(tree.get(&30).unwrap(), &30 43 /// } 44 /// 45 /// // Iterate over the nodes we just inserted 46 /// { 47 /// let mut iter = tree.iter(); 48 /// assert_eq!(iter.next().unwrap(), (&10, 49 /// assert_eq!(iter.next().unwrap(), (&20, 50 /// assert_eq!(iter.next().unwrap(), (&30, 51 /// assert!(iter.next().is_none()); 52 /// } 53 /// 54 /// // Print all elements. 55 /// for (key, value) in &tree { 56 /// pr_info!("{} = {}\n", key, value); 57 /// } 58 /// 59 /// // Replace one of the elements. 60 /// tree.try_create_and_insert(10, 1000, flags 61 /// 62 /// // Check that the tree reflects the replac 63 /// { 64 /// let mut iter = tree.iter(); 65 /// assert_eq!(iter.next().unwrap(), (&10, 66 /// assert_eq!(iter.next().unwrap(), (&20, 67 /// assert_eq!(iter.next().unwrap(), (&30, 68 /// assert!(iter.next().is_none()); 69 /// } 70 /// 71 /// // Change the value of one of the elements 72 /// *tree.get_mut(&30).unwrap() = 3000; 73 /// 74 /// // Check that the tree reflects the update 75 /// { 76 /// let mut iter = tree.iter(); 77 /// assert_eq!(iter.next().unwrap(), (&10, 78 /// assert_eq!(iter.next().unwrap(), (&20, 79 /// assert_eq!(iter.next().unwrap(), (&30, 80 /// assert!(iter.next().is_none()); 81 /// } 82 /// 83 /// // Remove an element. 84 /// tree.remove(&10); 85 /// 86 /// // Check that the tree reflects the remova 87 /// { 88 /// let mut iter = tree.iter(); 89 /// assert_eq!(iter.next().unwrap(), (&20, 90 /// assert_eq!(iter.next().unwrap(), (&30, 91 /// assert!(iter.next().is_none()); 92 /// } 93 /// 94 /// # Ok::<(), Error>(()) 95 /// ``` 96 /// 97 /// In the example below, we first allocate a 98 /// the tree. This is useful when the insertio 99 /// holding a spinlock. 100 /// 101 /// ``` 102 /// use kernel::{alloc::flags, rbtree::{RBTree 103 /// 104 /// fn insert_test(tree: &SpinLock<RBTree<u32, 105 /// // Pre-allocate node. This may fail (a 106 /// let node = RBTreeNode::new(10, 100, fl 107 /// 108 /// // Insert node while holding the lock. 109 /// // attempts. 110 /// let mut guard = tree.lock(); 111 /// guard.insert(node); 112 /// Ok(()) 113 /// } 114 /// ``` 115 /// 116 /// In the example below, we reuse an existing 117 /// 118 /// ``` 119 /// use kernel::{alloc::flags, rbtree::{RBTree 120 /// 121 /// // Create a new tree. 122 /// let mut tree = RBTree::new(); 123 /// 124 /// // Insert three elements. 125 /// tree.try_create_and_insert(20, 200, flags: 126 /// tree.try_create_and_insert(10, 100, flags: 127 /// tree.try_create_and_insert(30, 300, flags: 128 /// 129 /// // Check the nodes we just inserted. 130 /// { 131 /// let mut iter = tree.iter(); 132 /// assert_eq!(iter.next().unwrap(), (&10, 133 /// assert_eq!(iter.next().unwrap(), (&20, 134 /// assert_eq!(iter.next().unwrap(), (&30, 135 /// assert!(iter.next().is_none()); 136 /// } 137 /// 138 /// // Remove a node, getting back ownership o 139 /// let existing = tree.remove(&30).unwrap(); 140 /// 141 /// // Check that the tree reflects the remova 142 /// { 143 /// let mut iter = tree.iter(); 144 /// assert_eq!(iter.next().unwrap(), (&10, 145 /// assert_eq!(iter.next().unwrap(), (&20, 146 /// assert!(iter.next().is_none()); 147 /// } 148 /// 149 /// // Create a preallocated reservation that 150 /// let reservation = RBTreeNodeReservation::n 151 /// 152 /// // Insert a new node into the tree, reusin 153 /// // succeed (no memory allocations). 154 /// tree.insert(reservation.into_node(15, 150) 155 /// 156 /// // Check that the tree reflect the new ins 157 /// { 158 /// let mut iter = tree.iter(); 159 /// assert_eq!(iter.next().unwrap(), (&10, 160 /// assert_eq!(iter.next().unwrap(), (&15, 161 /// assert_eq!(iter.next().unwrap(), (&20, 162 /// assert!(iter.next().is_none()); 163 /// } 164 /// 165 /// # Ok::<(), Error>(()) 166 /// ``` 167 /// 168 /// # Invariants 169 /// 170 /// Non-null parent/children pointers stored i 171 /// valid, and pointing to a field of our inte 172 pub struct RBTree<K, V> { 173 root: bindings::rb_root, 174 _p: PhantomData<Node<K, V>>, 175 } 176 177 // SAFETY: An [`RBTree`] allows the same kinds 178 // fields, so we use the same Send condition a 179 unsafe impl<K: Send, V: Send> Send for RBTree< 180 181 // SAFETY: An [`RBTree`] allows the same kinds 182 // fields, so we use the same Sync condition a 183 unsafe impl<K: Sync, V: Sync> Sync for RBTree< 184 185 impl<K, V> RBTree<K, V> { 186 /// Creates a new and empty tree. 187 pub fn new() -> Self { 188 Self { 189 // INVARIANT: There are no nodes i 190 root: bindings::rb_root::default() 191 _p: PhantomData, 192 } 193 } 194 195 /// Returns an iterator over the tree node 196 pub fn iter(&self) -> Iter<'_, K, V> { 197 Iter { 198 _tree: PhantomData, 199 // INVARIANT: 200 // - `self.root` is a valid poin 201 // - `bindings::rb_first` produc 202 iter_raw: IterRaw { 203 // SAFETY: by the invariants, 204 next: unsafe { bindings::rb_fi 205 _phantom: PhantomData, 206 }, 207 } 208 } 209 210 /// Returns a mutable iterator over the tr 211 pub fn iter_mut(&mut self) -> IterMut<'_, 212 IterMut { 213 _tree: PhantomData, 214 // INVARIANT: 215 // - `self.root` is a valid poin 216 // - `bindings::rb_first` produc 217 iter_raw: IterRaw { 218 // SAFETY: by the invariants, 219 next: unsafe { bindings::rb_fi 220 _phantom: PhantomData, 221 }, 222 } 223 } 224 225 /// Returns an iterator over the keys of t 226 pub fn keys(&self) -> impl Iterator<Item = 227 self.iter().map(|(k, _)| k) 228 } 229 230 /// Returns an iterator over the values of 231 pub fn values(&self) -> impl Iterator<Item 232 self.iter().map(|(_, v)| v) 233 } 234 235 /// Returns a mutable iterator over the va 236 pub fn values_mut(&mut self) -> impl Itera 237 self.iter_mut().map(|(_, v)| v) 238 } 239 240 /// Returns a cursor over the tree nodes, 241 pub fn cursor_front(&mut self) -> Option<C 242 let root = addr_of_mut!(self.root); 243 // SAFETY: `self.root` is always a val 244 let current = unsafe { bindings::rb_fi 245 NonNull::new(current).map(|current| { 246 // INVARIANT: 247 // - `current` is a valid node in 248 Cursor { 249 current, 250 tree: self, 251 } 252 }) 253 } 254 255 /// Returns a cursor over the tree nodes, 256 pub fn cursor_back(&mut self) -> Option<Cu 257 let root = addr_of_mut!(self.root); 258 // SAFETY: `self.root` is always a val 259 let current = unsafe { bindings::rb_la 260 NonNull::new(current).map(|current| { 261 // INVARIANT: 262 // - `current` is a valid node in 263 Cursor { 264 current, 265 tree: self, 266 } 267 }) 268 } 269 } 270 271 impl<K, V> RBTree<K, V> 272 where 273 K: Ord, 274 { 275 /// Tries to insert a new value into the t 276 /// 277 /// It overwrites a node if one already ex 278 /// key/value pair). Returns [`None`] if a 279 /// 280 /// Returns an error if it cannot allocate 281 pub fn try_create_and_insert( 282 &mut self, 283 key: K, 284 value: V, 285 flags: Flags, 286 ) -> Result<Option<RBTreeNode<K, V>>> { 287 Ok(self.insert(RBTreeNode::new(key, va 288 } 289 290 /// Inserts a new node into the tree. 291 /// 292 /// It overwrites a node if one already ex 293 /// key/value pair). Returns [`None`] if a 294 /// 295 /// This function always succeeds. 296 pub fn insert(&mut self, node: RBTreeNode< 297 match self.raw_entry(&node.node.key) { 298 RawEntry::Occupied(entry) => Some( 299 RawEntry::Vacant(entry) => { 300 entry.insert(node); 301 None 302 } 303 } 304 } 305 306 fn raw_entry(&mut self, key: &K) -> RawEnt 307 let raw_self: *mut RBTree<K, V> = self 308 // The returned `RawEntry` is used to 309 // The parameters of `bindings::rb_lin 310 // - `node`: A pointer to an uninitial 311 // - `parent`: A pointer to an existin 312 // null, and `node` will beco 313 // with a pointer to `node`. 314 // - `rb_link`: A pointer to either th 315 // specifies which child of ` 316 // value of `*rb_link` must b 317 // red/black tree is empty, t 318 // this case, `rb_link` is a 319 // 320 // We will traverse the tree looking f 321 // representing an empty subtree where 322 // that we preserve the ordering of th 323 // we store `parent` and `child_field_ 324 // in the subtree of `parent` that `ch 325 // we find an empty subtree, we can in 326 let mut parent = core::ptr::null_mut() 327 let mut child_field_of_parent: &mut *m 328 // SAFETY: `raw_self` is a valid p 329 unsafe { &mut (*raw_self).root.rb_ 330 while !(*child_field_of_parent).is_nul 331 let curr = *child_field_of_parent; 332 // SAFETY: All links fields we cre 333 let node = unsafe { container_of!( 334 335 // SAFETY: `node` is a non-null no 336 match key.cmp(unsafe { &(*node).ke 337 // SAFETY: `curr` is a non-nul 338 Ordering::Less => child_field_ 339 // SAFETY: `curr` is a non-nul 340 Ordering::Greater => child_fie 341 Ordering::Equal => { 342 return RawEntry::Occupied( 343 rbtree: self, 344 node_links: curr, 345 }) 346 } 347 } 348 parent = curr; 349 } 350 351 RawEntry::Vacant(RawVacantEntry { 352 rbtree: raw_self, 353 parent, 354 child_field_of_parent, 355 _phantom: PhantomData, 356 }) 357 } 358 359 /// Gets the given key's corresponding ent 360 pub fn entry(&mut self, key: K) -> Entry<' 361 match self.raw_entry(&key) { 362 RawEntry::Occupied(entry) => Entry 363 RawEntry::Vacant(entry) => Entry:: 364 } 365 } 366 367 /// Used for accessing the given node, if 368 pub fn find_mut(&mut self, key: &K) -> Opt 369 match self.raw_entry(key) { 370 RawEntry::Occupied(entry) => Some( 371 RawEntry::Vacant(_entry) => None, 372 } 373 } 374 375 /// Returns a reference to the value corre 376 pub fn get(&self, key: &K) -> Option<&V> { 377 let mut node = self.root.rb_node; 378 while !node.is_null() { 379 // SAFETY: By the type invariant o 380 // point to the links field of `No 381 let this = unsafe { container_of!( 382 // SAFETY: `this` is a non-null no 383 node = match key.cmp(unsafe { &(*t 384 // SAFETY: `node` is a non-nul 385 Ordering::Less => unsafe { (*n 386 // SAFETY: `node` is a non-nul 387 Ordering::Greater => unsafe { 388 // SAFETY: `node` is a non-nul 389 Ordering::Equal => return Some 390 } 391 } 392 None 393 } 394 395 /// Returns a mutable reference to the val 396 pub fn get_mut(&mut self, key: &K) -> Opti 397 self.find_mut(key).map(|node| node.int 398 } 399 400 /// Removes the node with the given key fr 401 /// 402 /// It returns the node that was removed i 403 pub fn remove_node(&mut self, key: &K) -> 404 self.find_mut(key).map(OccupiedEntry:: 405 } 406 407 /// Removes the node with the given key fr 408 /// 409 /// It returns the value that was removed 410 pub fn remove(&mut self, key: &K) -> Optio 411 self.find_mut(key).map(OccupiedEntry:: 412 } 413 414 /// Returns a cursor over the tree nodes b 415 /// 416 /// If the given key exists, the cursor st 417 /// Otherwise it starts with the first lar 418 /// If there is no larger key, it returns 419 pub fn cursor_lower_bound(&mut self, key: 420 where 421 K: Ord, 422 { 423 let mut node = self.root.rb_node; 424 let mut best_match: Option<NonNull<Nod 425 while !node.is_null() { 426 // SAFETY: By the type invariant o 427 // point to the links field of `No 428 let this = unsafe { container_of!( 429 // SAFETY: `this` is a non-null no 430 let this_key = unsafe { &(*this).k 431 // SAFETY: `node` is a non-null no 432 let left_child = unsafe { (*node). 433 // SAFETY: `node` is a non-null no 434 let right_child = unsafe { (*node) 435 match key.cmp(this_key) { 436 Ordering::Equal => { 437 best_match = NonNull::new( 438 break; 439 } 440 Ordering::Greater => { 441 node = right_child; 442 } 443 Ordering::Less => { 444 let is_better_match = matc 445 None => true, 446 Some(best) => { 447 // SAFETY: `best` 448 let best_key = uns 449 best_key > this_ke 450 } 451 }; 452 if is_better_match { 453 best_match = NonNull:: 454 } 455 node = left_child; 456 } 457 }; 458 } 459 460 let best = best_match?; 461 462 // SAFETY: `best` is a non-null node s 463 let links = unsafe { addr_of_mut!((*be 464 465 NonNull::new(links).map(|current| { 466 // INVARIANT: 467 // - `current` is a valid node in 468 Cursor { 469 current, 470 tree: self, 471 } 472 }) 473 } 474 } 475 476 impl<K, V> Default for RBTree<K, V> { 477 fn default() -> Self { 478 Self::new() 479 } 480 } 481 482 impl<K, V> Drop for RBTree<K, V> { 483 fn drop(&mut self) { 484 // SAFETY: `root` is valid as it's emb 485 let mut next = unsafe { bindings::rb_f 486 487 // INVARIANT: The loop invariant is th 488 while !next.is_null() { 489 // SAFETY: All links fields we cre 490 let this = unsafe { container_of!( 491 492 // Find out what the next node is 493 // SAFETY: `next` and all nodes in 494 next = unsafe { bindings::rb_next_ 495 496 // INVARIANT: This is the destruct 497 // but it is not observable. The l 498 499 // SAFETY: `this` is valid per the 500 unsafe { drop(Box::from_raw(this.c 501 } 502 } 503 } 504 505 /// A bidirectional cursor over the tree nodes 506 /// 507 /// # Examples 508 /// 509 /// In the following example, we obtain a curs 510 /// The cursor allows us to iterate bidirectio 511 /// 512 /// ``` 513 /// use kernel::{alloc::flags, rbtree::RBTree} 514 /// 515 /// // Create a new tree. 516 /// let mut tree = RBTree::new(); 517 /// 518 /// // Insert three elements. 519 /// tree.try_create_and_insert(10, 100, flags: 520 /// tree.try_create_and_insert(20, 200, flags: 521 /// tree.try_create_and_insert(30, 300, flags: 522 /// 523 /// // Get a cursor to the first element. 524 /// let mut cursor = tree.cursor_front().unwra 525 /// let mut current = cursor.current(); 526 /// assert_eq!(current, (&10, &100)); 527 /// 528 /// // Move the cursor, updating it to the 2nd 529 /// cursor = cursor.move_next().unwrap(); 530 /// current = cursor.current(); 531 /// assert_eq!(current, (&20, &200)); 532 /// 533 /// // Peek at the next element without impact 534 /// let next = cursor.peek_next().unwrap(); 535 /// assert_eq!(next, (&30, &300)); 536 /// current = cursor.current(); 537 /// assert_eq!(current, (&20, &200)); 538 /// 539 /// // Moving past the last element causes the 540 /// cursor = cursor.move_next().unwrap(); 541 /// current = cursor.current(); 542 /// assert_eq!(current, (&30, &300)); 543 /// let cursor = cursor.move_next(); 544 /// assert!(cursor.is_none()); 545 /// 546 /// # Ok::<(), Error>(()) 547 /// ``` 548 /// 549 /// A cursor can also be obtained at the last 550 /// 551 /// ``` 552 /// use kernel::{alloc::flags, rbtree::RBTree} 553 /// 554 /// // Create a new tree. 555 /// let mut tree = RBTree::new(); 556 /// 557 /// // Insert three elements. 558 /// tree.try_create_and_insert(10, 100, flags: 559 /// tree.try_create_and_insert(20, 200, flags: 560 /// tree.try_create_and_insert(30, 300, flags: 561 /// 562 /// let mut cursor = tree.cursor_back().unwrap 563 /// let current = cursor.current(); 564 /// assert_eq!(current, (&30, &300)); 565 /// 566 /// # Ok::<(), Error>(()) 567 /// ``` 568 /// 569 /// Obtaining a cursor returns [`None`] if the 570 /// 571 /// ``` 572 /// use kernel::rbtree::RBTree; 573 /// 574 /// let mut tree: RBTree<u16, u16> = RBTree::n 575 /// assert!(tree.cursor_front().is_none()); 576 /// 577 /// # Ok::<(), Error>(()) 578 /// ``` 579 /// 580 /// [`RBTree::cursor_lower_bound`] can be used 581 /// 582 /// ``` 583 /// use kernel::{alloc::flags, rbtree::RBTree} 584 /// 585 /// // Create a new tree. 586 /// let mut tree = RBTree::new(); 587 /// 588 /// // Insert five elements. 589 /// tree.try_create_and_insert(10, 100, flags: 590 /// tree.try_create_and_insert(20, 200, flags: 591 /// tree.try_create_and_insert(30, 300, flags: 592 /// tree.try_create_and_insert(40, 400, flags: 593 /// tree.try_create_and_insert(50, 500, flags: 594 /// 595 /// // If the provided key exists, a cursor to 596 /// let cursor = tree.cursor_lower_bound(&20). 597 /// let current = cursor.current(); 598 /// assert_eq!(current, (&20, &200)); 599 /// 600 /// // If the provided key doesn't exist, a cu 601 /// let cursor = tree.cursor_lower_bound(&25). 602 /// let current = cursor.current(); 603 /// assert_eq!(current, (&30, &300)); 604 /// 605 /// // If there is no larger key, [`None`] is 606 /// let cursor = tree.cursor_lower_bound(&55); 607 /// assert!(cursor.is_none()); 608 /// 609 /// # Ok::<(), Error>(()) 610 /// ``` 611 /// 612 /// The cursor allows mutation of values in th 613 /// 614 /// ``` 615 /// use kernel::{alloc::flags, rbtree::RBTree} 616 /// 617 /// // Create a new tree. 618 /// let mut tree = RBTree::new(); 619 /// 620 /// // Insert three elements. 621 /// tree.try_create_and_insert(10, 100, flags: 622 /// tree.try_create_and_insert(20, 200, flags: 623 /// tree.try_create_and_insert(30, 300, flags: 624 /// 625 /// // Retrieve a cursor. 626 /// let mut cursor = tree.cursor_front().unwra 627 /// 628 /// // Get a mutable reference to the current 629 /// let (k, v) = cursor.current_mut(); 630 /// *v = 1000; 631 /// 632 /// // The updated value is reflected in the t 633 /// let updated = tree.get(&10).unwrap(); 634 /// assert_eq!(updated, &1000); 635 /// 636 /// # Ok::<(), Error>(()) 637 /// ``` 638 /// 639 /// It also allows node removal. The following 640 /// 641 /// ``` 642 /// use kernel::{alloc::flags, rbtree::RBTree} 643 /// 644 /// // Create a new tree. 645 /// let mut tree = RBTree::new(); 646 /// 647 /// // Insert three elements. 648 /// tree.try_create_and_insert(10, 100, flags: 649 /// tree.try_create_and_insert(20, 200, flags: 650 /// tree.try_create_and_insert(30, 300, flags: 651 /// 652 /// // Remove the first element. 653 /// let mut cursor = tree.cursor_front().unwra 654 /// let mut current = cursor.current(); 655 /// assert_eq!(current, (&10, &100)); 656 /// cursor = cursor.remove_current().0.unwrap( 657 /// 658 /// // If a node exists after the current elem 659 /// current = cursor.current(); 660 /// assert_eq!(current, (&20, &200)); 661 /// 662 /// // Get a cursor to the last element, and r 663 /// cursor = tree.cursor_back().unwrap(); 664 /// current = cursor.current(); 665 /// assert_eq!(current, (&30, &300)); 666 /// 667 /// // Since there is no next node, the previo 668 /// cursor = cursor.remove_current().0.unwrap( 669 /// current = cursor.current(); 670 /// assert_eq!(current, (&20, &200)); 671 /// 672 /// // Removing the last element in the tree r 673 /// assert!(cursor.remove_current().0.is_none( 674 /// 675 /// # Ok::<(), Error>(()) 676 /// ``` 677 /// 678 /// Nodes adjacent to the current node can als 679 /// 680 /// ``` 681 /// use kernel::{alloc::flags, rbtree::RBTree} 682 /// 683 /// // Create a new tree. 684 /// let mut tree = RBTree::new(); 685 /// 686 /// // Insert three elements. 687 /// tree.try_create_and_insert(10, 100, flags: 688 /// tree.try_create_and_insert(20, 200, flags: 689 /// tree.try_create_and_insert(30, 300, flags: 690 /// 691 /// // Get a cursor to the first element. 692 /// let mut cursor = tree.cursor_front().unwra 693 /// let mut current = cursor.current(); 694 /// assert_eq!(current, (&10, &100)); 695 /// 696 /// // Calling `remove_prev` from the first el 697 /// assert!(cursor.remove_prev().is_none()); 698 /// 699 /// // Get a cursor to the last element. 700 /// cursor = tree.cursor_back().unwrap(); 701 /// current = cursor.current(); 702 /// assert_eq!(current, (&30, &300)); 703 /// 704 /// // Calling `remove_prev` removes and retur 705 /// assert_eq!(cursor.remove_prev().unwrap().t 706 /// 707 /// // Calling `remove_next` from the last ele 708 /// assert!(cursor.remove_next().is_none()); 709 /// 710 /// // Move to the first element 711 /// cursor = cursor.move_prev().unwrap(); 712 /// current = cursor.current(); 713 /// assert_eq!(current, (&10, &100)); 714 /// 715 /// // Calling `remove_next` removes and retur 716 /// assert_eq!(cursor.remove_next().unwrap().t 717 /// 718 /// # Ok::<(), Error>(()) 719 /// 720 /// ``` 721 /// 722 /// # Invariants 723 /// - `current` points to a node that is in th 724 pub struct Cursor<'a, K, V> { 725 tree: &'a mut RBTree<K, V>, 726 current: NonNull<bindings::rb_node>, 727 } 728 729 // SAFETY: The [`Cursor`] has exclusive access 730 // The cursor only gives out immutable referen 731 // keys, `Send` is sufficient. `Sync` would be 732 unsafe impl<'a, K: Send, V: Send> Send for Cur 733 734 // SAFETY: The [`Cursor`] gives out immutable 735 // so it has the same thread safety requiremen 736 unsafe impl<'a, K: Sync, V: Sync> Sync for Cur 737 738 impl<'a, K, V> Cursor<'a, K, V> { 739 /// The current node 740 pub fn current(&self) -> (&K, &V) { 741 // SAFETY: 742 // - `self.current` is a valid node by 743 // - We have an immutable reference by 744 unsafe { Self::to_key_value(self.curre 745 } 746 747 /// The current node, with a mutable value 748 pub fn current_mut(&mut self) -> (&K, &mut 749 // SAFETY: 750 // - `self.current` is a valid node by 751 // - We have an mutable reference by t 752 unsafe { Self::to_key_value_mut(self.c 753 } 754 755 /// Remove the current node from the tree. 756 /// 757 /// Returns a tuple where the first elemen 758 /// else the previous node, else [`None`] 759 /// is the removed node. 760 pub fn remove_current(self) -> (Option<Sel 761 let prev = self.get_neighbor_raw(Direc 762 let next = self.get_neighbor_raw(Direc 763 // SAFETY: By the type invariant of `S 764 // point to the links field of `Node<K 765 let this = unsafe { container_of!(self 766 // SAFETY: `this` is valid by the type 767 let node = unsafe { Box::from_raw(this 768 let node = RBTreeNode { node }; 769 // SAFETY: The reference to the tree u 770 // the tree cannot change. By the tree 771 unsafe { bindings::rb_erase(&mut (*thi 772 773 let current = match (prev, next) { 774 (_, Some(next)) => next, 775 (Some(prev), None) => prev, 776 (None, None) => { 777 return (None, node); 778 } 779 }; 780 781 ( 782 // INVARIANT: 783 // - `current` is a valid node in 784 Some(Self { 785 current, 786 tree: self.tree, 787 }), 788 node, 789 ) 790 } 791 792 /// Remove the previous node, returning it 793 pub fn remove_prev(&mut self) -> Option<RB 794 self.remove_neighbor(Direction::Prev) 795 } 796 797 /// Remove the next node, returning it if 798 pub fn remove_next(&mut self) -> Option<RB 799 self.remove_neighbor(Direction::Next) 800 } 801 802 fn remove_neighbor(&mut self, direction: D 803 if let Some(neighbor) = self.get_neigh 804 let neighbor = neighbor.as_ptr(); 805 // SAFETY: The reference to the tr 806 // the tree cannot change. By the 807 unsafe { bindings::rb_erase(neighb 808 // SAFETY: By the type invariant o 809 // point to the links field of `No 810 let this = unsafe { container_of!( 811 // SAFETY: `this` is valid by the 812 let node = unsafe { Box::from_raw( 813 return Some(RBTreeNode { node }); 814 } 815 None 816 } 817 818 /// Move the cursor to the previous node, 819 pub fn move_prev(self) -> Option<Self> { 820 self.mv(Direction::Prev) 821 } 822 823 /// Move the cursor to the next node, retu 824 pub fn move_next(self) -> Option<Self> { 825 self.mv(Direction::Next) 826 } 827 828 fn mv(self, direction: Direction) -> Optio 829 // INVARIANT: 830 // - `neighbor` is a valid node in the 831 self.get_neighbor_raw(direction).map(| 832 tree: self.tree, 833 current: neighbor, 834 }) 835 } 836 837 /// Access the previous node without movin 838 pub fn peek_prev(&self) -> Option<(&K, &V) 839 self.peek(Direction::Prev) 840 } 841 842 /// Access the previous node without movin 843 pub fn peek_next(&self) -> Option<(&K, &V) 844 self.peek(Direction::Next) 845 } 846 847 fn peek(&self, direction: Direction) -> Op 848 self.get_neighbor_raw(direction).map(| 849 // SAFETY: 850 // - `neighbor` is a valid tree no 851 // - By the function signature, we 852 unsafe { Self::to_key_value(neighb 853 }) 854 } 855 856 /// Access the previous node mutably witho 857 pub fn peek_prev_mut(&mut self) -> Option< 858 self.peek_mut(Direction::Prev) 859 } 860 861 /// Access the next node mutably without m 862 pub fn peek_next_mut(&mut self) -> Option< 863 self.peek_mut(Direction::Next) 864 } 865 866 fn peek_mut(&mut self, direction: Directio 867 self.get_neighbor_raw(direction).map(| 868 // SAFETY: 869 // - `neighbor` is a valid tree no 870 // - By the function signature, we 871 unsafe { Self::to_key_value_mut(ne 872 }) 873 } 874 875 fn get_neighbor_raw(&self, direction: Dire 876 // SAFETY: `self.current` is valid by 877 let neighbor = unsafe { 878 match direction { 879 Direction::Prev => bindings::r 880 Direction::Next => bindings::r 881 } 882 }; 883 884 NonNull::new(neighbor) 885 } 886 887 /// SAFETY: 888 /// - `node` must be a valid pointer to a 889 /// - The caller has immutable access to ` 890 unsafe fn to_key_value<'b>(node: NonNull<b 891 // SAFETY: the caller guarantees that 892 let (k, v) = unsafe { Self::to_key_val 893 // SAFETY: the caller guarantees immut 894 (k, unsafe { &*v }) 895 } 896 897 /// SAFETY: 898 /// - `node` must be a valid pointer to a 899 /// - The caller has mutable access to `no 900 unsafe fn to_key_value_mut<'b>(node: NonNu 901 // SAFETY: the caller guarantees that 902 let (k, v) = unsafe { Self::to_key_val 903 // SAFETY: the caller guarantees mutab 904 (k, unsafe { &mut *v }) 905 } 906 907 /// SAFETY: 908 /// - `node` must be a valid pointer to a 909 /// - The caller has immutable access to t 910 unsafe fn to_key_value_raw<'b>(node: NonNu 911 // SAFETY: By the type invariant of `S 912 // point to the links field of `Node<K 913 let this = unsafe { container_of!(node 914 // SAFETY: The passed `node` is the cu 915 // thus `this` is valid by the type in 916 let k = unsafe { &(*this).key }; 917 // SAFETY: The passed `node` is the cu 918 // thus `this` is valid by the type in 919 let v = unsafe { addr_of_mut!((*this). 920 (k, v) 921 } 922 } 923 924 /// Direction for [`Cursor`] operations. 925 enum Direction { 926 /// the node immediately before, in sort o 927 Prev, 928 /// the node immediately after, in sort or 929 Next, 930 } 931 932 impl<'a, K, V> IntoIterator for &'a RBTree<K, 933 type Item = (&'a K, &'a V); 934 type IntoIter = Iter<'a, K, V>; 935 936 fn into_iter(self) -> Self::IntoIter { 937 self.iter() 938 } 939 } 940 941 /// An iterator over the nodes of a [`RBTree`] 942 /// 943 /// Instances are created by calling [`RBTree: 944 pub struct Iter<'a, K, V> { 945 _tree: PhantomData<&'a RBTree<K, V>>, 946 iter_raw: IterRaw<K, V>, 947 } 948 949 // SAFETY: The [`Iter`] gives out immutable re 950 // thread safety requirements as immutable ref 951 unsafe impl<'a, K: Sync, V: Sync> Send for Ite 952 953 // SAFETY: The [`Iter`] gives out immutable re 954 // thread safety requirements as immutable ref 955 unsafe impl<'a, K: Sync, V: Sync> Sync for Ite 956 957 impl<'a, K, V> Iterator for Iter<'a, K, V> { 958 type Item = (&'a K, &'a V); 959 960 fn next(&mut self) -> Option<Self::Item> { 961 // SAFETY: Due to `self._tree`, `k` an 962 self.iter_raw.next().map(|(k, v)| unsa 963 } 964 } 965 966 impl<'a, K, V> IntoIterator for &'a mut RBTree 967 type Item = (&'a K, &'a mut V); 968 type IntoIter = IterMut<'a, K, V>; 969 970 fn into_iter(self) -> Self::IntoIter { 971 self.iter_mut() 972 } 973 } 974 975 /// A mutable iterator over the nodes of a [`R 976 /// 977 /// Instances are created by calling [`RBTree: 978 pub struct IterMut<'a, K, V> { 979 _tree: PhantomData<&'a mut RBTree<K, V>>, 980 iter_raw: IterRaw<K, V>, 981 } 982 983 // SAFETY: The [`IterMut`] has exclusive acces 984 // The iterator only gives out immutable refer 985 // keys, `Send` is sufficient. `Sync` would be 986 unsafe impl<'a, K: Send, V: Send> Send for Ite 987 988 // SAFETY: The [`IterMut`] gives out immutable 989 // thread safety requirements as mutable refer 990 unsafe impl<'a, K: Sync, V: Sync> Sync for Ite 991 992 impl<'a, K, V> Iterator for IterMut<'a, K, V> 993 type Item = (&'a K, &'a mut V); 994 995 fn next(&mut self) -> Option<Self::Item> { 996 self.iter_raw.next().map(|(k, v)| 997 // SAFETY: Due to `&mut self`, we 998 unsafe { (&*k, &mut *v) }) 999 } 1000 } 1001 1002 /// A raw iterator over the nodes of a [`RBTr 1003 /// 1004 /// # Invariants 1005 /// - `self.next` is a valid pointer. 1006 /// - `self.next` points to a node stored ins 1007 struct IterRaw<K, V> { 1008 next: *mut bindings::rb_node, 1009 _phantom: PhantomData<fn() -> (K, V)>, 1010 } 1011 1012 impl<K, V> Iterator for IterRaw<K, V> { 1013 type Item = (*mut K, *mut V); 1014 1015 fn next(&mut self) -> Option<Self::Item> 1016 if self.next.is_null() { 1017 return None; 1018 } 1019 1020 // SAFETY: By the type invariant of ` 1021 // and by the type invariant of `RBTr 1022 let cur = unsafe { container_of!(self 1023 1024 // SAFETY: `self.next` is a valid tre 1025 self.next = unsafe { bindings::rb_nex 1026 1027 // SAFETY: By the same reasoning abov 1028 Some(unsafe { (addr_of_mut!((*cur).ke 1029 } 1030 } 1031 1032 /// A memory reservation for a red-black tree 1033 /// 1034 /// 1035 /// It contains the memory needed to hold a n 1036 /// can be obtained by directly allocating it 1037 pub struct RBTreeNodeReservation<K, V> { 1038 node: Box<MaybeUninit<Node<K, V>>>, 1039 } 1040 1041 impl<K, V> RBTreeNodeReservation<K, V> { 1042 /// Allocates memory for a node to be eve 1043 /// call to [`RBTree::insert`]. 1044 pub fn new(flags: Flags) -> Result<RBTree 1045 Ok(RBTreeNodeReservation { 1046 node: <Box<_> as BoxExt<_>>::new_ 1047 }) 1048 } 1049 } 1050 1051 // SAFETY: This doesn't actually contain K or 1052 // be moved across threads. 1053 unsafe impl<K, V> Send for RBTreeNodeReservat 1054 1055 // SAFETY: This doesn't actually contain K or 1056 unsafe impl<K, V> Sync for RBTreeNodeReservat 1057 1058 impl<K, V> RBTreeNodeReservation<K, V> { 1059 /// Initialises a node reservation. 1060 /// 1061 /// It then becomes an [`RBTreeNode`] tha 1062 pub fn into_node(mut self, key: K, value: 1063 self.node.write(Node { 1064 key, 1065 value, 1066 links: bindings::rb_node::default 1067 }); 1068 // SAFETY: We just wrote to it. 1069 let node = unsafe { self.node.assume_ 1070 RBTreeNode { node } 1071 } 1072 } 1073 1074 /// A red-black tree node. 1075 /// 1076 /// The node is fully initialised (with key a 1077 /// extra allocations or failure paths. 1078 pub struct RBTreeNode<K, V> { 1079 node: Box<Node<K, V>>, 1080 } 1081 1082 impl<K, V> RBTreeNode<K, V> { 1083 /// Allocates and initialises a node that 1084 /// [`RBTree::insert`]. 1085 pub fn new(key: K, value: V, flags: Flags 1086 Ok(RBTreeNodeReservation::new(flags)? 1087 } 1088 1089 /// Get the key and value from inside the 1090 pub fn to_key_value(self) -> (K, V) { 1091 (self.node.key, self.node.value) 1092 } 1093 } 1094 1095 // SAFETY: If K and V can be sent across thre 1096 // threads. 1097 unsafe impl<K: Send, V: Send> Send for RBTree 1098 1099 // SAFETY: If K and V can be accessed without 1100 // [`RBTreeNode`] without synchronization. 1101 unsafe impl<K: Sync, V: Sync> Sync for RBTree 1102 1103 impl<K, V> RBTreeNode<K, V> { 1104 /// Drop the key and value, but keep the 1105 /// 1106 /// It then becomes a reservation that ca 1107 /// a different key and/or value). 1108 /// 1109 /// The existing key and value are droppe 1110 /// may be freed (but only for the key/va 1111 pub fn into_reservation(self) -> RBTreeNo 1112 RBTreeNodeReservation { 1113 node: Box::drop_contents(self.nod 1114 } 1115 } 1116 } 1117 1118 /// A view into a single entry in a map, whic 1119 /// 1120 /// This enum is constructed from the [`RBTre 1121 /// 1122 /// [`entry`]: fn@RBTree::entry 1123 pub enum Entry<'a, K, V> { 1124 /// This [`RBTree`] does not have a node 1125 Vacant(VacantEntry<'a, K, V>), 1126 /// This [`RBTree`] already has a node wi 1127 Occupied(OccupiedEntry<'a, K, V>), 1128 } 1129 1130 /// Like [`Entry`], except that it doesn't ha 1131 enum RawEntry<'a, K, V> { 1132 Vacant(RawVacantEntry<'a, K, V>), 1133 Occupied(OccupiedEntry<'a, K, V>), 1134 } 1135 1136 /// A view into a vacant entry in a [`RBTree` 1137 pub struct VacantEntry<'a, K, V> { 1138 key: K, 1139 raw: RawVacantEntry<'a, K, V>, 1140 } 1141 1142 /// Like [`VacantEntry`], but doesn't hold on 1143 /// 1144 /// # Invariants 1145 /// - `parent` may be null if the new node be 1146 /// - `child_field_of_parent` is a valid poin 1147 /// null, it is a pointer to the root of 1148 struct RawVacantEntry<'a, K, V> { 1149 rbtree: *mut RBTree<K, V>, 1150 /// The node that will become the parent 1151 parent: *mut bindings::rb_node, 1152 /// This points to the left-child or righ 1153 /// null. 1154 child_field_of_parent: *mut *mut bindings 1155 _phantom: PhantomData<&'a mut RBTree<K, V 1156 } 1157 1158 impl<'a, K, V> RawVacantEntry<'a, K, V> { 1159 /// Inserts the given node into the [`RBT 1160 /// 1161 /// The `node` must have a key such that 1162 /// [`RBTree`]. 1163 fn insert(self, node: RBTreeNode<K, V>) - 1164 let node = Box::into_raw(node.node); 1165 1166 // SAFETY: `node` is valid at least u 1167 // the node is removed or replaced. 1168 let node_links = unsafe { addr_of_mut 1169 1170 // INVARIANT: We are linking in a new 1171 // "forgot" it with `Box::into_raw`. 1172 // SAFETY: The type invariants of `Ra 1173 unsafe { bindings::rb_link_node(node_ 1174 1175 // SAFETY: All pointers are valid. `n 1176 unsafe { bindings::rb_insert_color(no 1177 1178 // SAFETY: The node is valid until we 1179 unsafe { &mut (*node).value } 1180 } 1181 } 1182 1183 impl<'a, K, V> VacantEntry<'a, K, V> { 1184 /// Inserts the given node into the [`RBT 1185 pub fn insert(self, value: V, reservation 1186 self.raw.insert(reservation.into_node 1187 } 1188 } 1189 1190 /// A view into an occupied entry in a [`RBTr 1191 /// 1192 /// # Invariants 1193 /// - `node_links` is a valid, non-null point 1194 pub struct OccupiedEntry<'a, K, V> { 1195 rbtree: &'a mut RBTree<K, V>, 1196 /// The node that this entry corresponds 1197 node_links: *mut bindings::rb_node, 1198 } 1199 1200 impl<'a, K, V> OccupiedEntry<'a, K, V> { 1201 /// Gets a reference to the value in the 1202 pub fn get(&self) -> &V { 1203 // SAFETY: 1204 // - `self.node_links` is a valid poi 1205 // - We have shared access to the und 1206 unsafe { &(*container_of!(self.node_l 1207 } 1208 1209 /// Gets a mutable reference to the value 1210 pub fn get_mut(&mut self) -> &mut V { 1211 // SAFETY: 1212 // - `self.node_links` is a valid poi 1213 // - We have exclusive access to the 1214 unsafe { &mut (*(container_of!(self.n 1215 } 1216 1217 /// Converts the entry into a mutable ref 1218 /// 1219 /// If you need multiple references to th 1220 pub fn into_mut(self) -> &'a mut V { 1221 // SAFETY: 1222 // - `self.node_links` is a valid poi 1223 // - This consumes the `&'a mut RBTre 1224 unsafe { &mut (*(container_of!(self.n 1225 } 1226 1227 /// Remove this entry from the [`RBTree`] 1228 pub fn remove_node(self) -> RBTreeNode<K, 1229 // SAFETY: The node is a node in the 1230 unsafe { bindings::rb_erase(self.node 1231 1232 // INVARIANT: The node is being retur 1233 // removed from the tree. So the inva 1234 RBTreeNode { 1235 // SAFETY: The node was a node in 1236 // back into a box. 1237 node: unsafe { 1238 Box::from_raw(container_of!(s 1239 }, 1240 } 1241 } 1242 1243 /// Takes the value of the entry out of t 1244 pub fn remove(self) -> V { 1245 self.remove_node().node.value 1246 } 1247 1248 /// Swap the current node for the provide 1249 /// 1250 /// The key of both nodes must be equal. 1251 fn replace(self, node: RBTreeNode<K, V>) 1252 let node = Box::into_raw(node.node); 1253 1254 // SAFETY: `node` is valid at least u 1255 // the node is removed or replaced. 1256 let new_node_links = unsafe { addr_of 1257 1258 // SAFETY: This updates the pointers 1259 // `self.node_links` used to be. 1260 unsafe { 1261 bindings::rb_replace_node(self.no 1262 }; 1263 1264 // SAFETY: 1265 // - `self.node_ptr` produces a valid 1266 // - Now that we removed this entry f 1267 let old_node = 1268 unsafe { Box::from_raw(container_ 1269 1270 RBTreeNode { node: old_node } 1271 } 1272 } 1273 1274 struct Node<K, V> { 1275 links: bindings::rb_node, 1276 key: K, 1277 value: V, 1278 }
Linux® is a registered trademark of Linus Torvalds in the United States and other countries.
TOMOYO® is a registered trademark of NTT DATA CORPORATION.