rust: rbtree: add immutable cursor · tjh.dev/kernel@f56b131

+2 -2

drivers/android/binder/freeze.rs

··· 321 321 KVVec::with_capacity(8, GFP_KERNEL).unwrap_or_else(|_err| KVVec::new()); 322 322 323 323 let mut inner = self.lock_with_nodes(); 324 - let mut curr = inner.nodes.cursor_front(); 324 + let mut curr = inner.nodes.cursor_front_mut(); 325 325 while let Some(cursor) = curr { 326 326 let (key, node) = cursor.current(); 327 327 let key = *key; ··· 335 335 // Find the node we were looking at and try again. If the set of nodes was changed, 336 336 // then just proceed to the next node. This is ok because we don't guarantee the 337 337 // inclusion of nodes that are added or removed in parallel with this operation. 338 - curr = inner.nodes.cursor_lower_bound(&key); 338 + curr = inner.nodes.cursor_lower_bound_mut(&key); 339 339 continue; 340 340 } 341 341

+1 -1

drivers/android/binder/process.rs

··· 1293 1293 { 1294 1294 while let Some(node) = { 1295 1295 let mut lock = self.inner.lock(); 1296 - lock.nodes.cursor_front().map(|c| c.remove_current().1) 1296 + lock.nodes.cursor_front_mut().map(|c| c.remove_current().1) 1297 1297 } { 1298 1298 node.to_key_value().1.release(); 1299 1299 }

+1 -1

drivers/android/binder/range_alloc/tree.rs

··· 207 207 } 208 208 209 209 pub(crate) fn reservation_abort(&mut self, offset: usize) -> Result<FreedRange> { 210 - let mut cursor = self.tree.cursor_lower_bound(&offset).ok_or_else(|| { 210 + let mut cursor = self.tree.cursor_lower_bound_mut(&offset).ok_or_else(|| { 211 211 pr_warn!( 212 212 "EINVAL from range_alloc.reservation_abort - offset: {}", 213 213 offset

+197 -47

rust/kernel/rbtree.rs

··· 243 243 } 244 244 245 245 /// Returns a cursor over the tree nodes, starting with the smallest key. 246 - pub fn cursor_front(&mut self) -> Option<Cursor<'_, K, V>> { 246 + pub fn cursor_front_mut(&mut self) -> Option<CursorMut<'_, K, V>> { 247 247 let root = addr_of_mut!(self.root); 248 - // SAFETY: `self.root` is always a valid root node 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. 249 264 let current = unsafe { bindings::rb_first(root) }; 250 265 NonNull::new(current).map(|current| { 251 266 // INVARIANT: 252 267 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 253 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 { 254 284 current, 255 285 tree: self, 256 286 } ··· 288 258 } 289 259 290 260 /// Returns a cursor over the tree nodes, starting with the largest key. 291 - pub fn cursor_back(&mut self) -> Option<Cursor<'_, K, V>> { 292 - let root = addr_of_mut!(self.root); 293 - // SAFETY: `self.root` is always a valid root node 261 + pub fn cursor_back(&self) -> Option<Cursor<'_, K, V>> { 262 + let root = &raw const self.root; 263 + // SAFETY: `self.root` is always a valid root node. 294 264 let current = unsafe { bindings::rb_last(root) }; 295 265 NonNull::new(current).map(|current| { 296 266 // INVARIANT: 297 267 // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 298 268 Cursor { 299 269 current, 300 - tree: self, 270 + _tree: PhantomData, 301 271 } 302 272 }) 303 273 } ··· 451 421 /// If the given key exists, the cursor starts there. 452 422 /// Otherwise it starts with the first larger key in sort order. 453 423 /// If there is no larger key, it returns [`None`]. 454 - pub fn cursor_lower_bound(&mut self, key: &K) -> Option<Cursor<'_, K, V>> 424 + pub fn cursor_lower_bound_mut(&mut self, key: &K) -> Option<CursorMut<'_, K, V>> 455 425 where 456 426 K: Ord, 457 427 { 428 + let best = self.find_best_match(key)?; 429 + 430 + NonNull::new(best.as_ptr()).map(|current| { 431 + // INVARIANT: 432 + // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 433 + CursorMut { 434 + current, 435 + tree: self, 436 + } 437 + }) 438 + } 439 + 440 + /// Returns a cursor over the tree nodes based on the given key. 441 + /// 442 + /// If the given key exists, the cursor starts there. 443 + /// Otherwise it starts with the first larger key in sort order. 444 + /// If there is no larger key, it returns [`None`]. 445 + pub fn cursor_lower_bound(&self, key: &K) -> Option<Cursor<'_, K, V>> 446 + where 447 + K: Ord, 448 + { 449 + let best = self.find_best_match(key)?; 450 + 451 + NonNull::new(best.as_ptr()).map(|current| { 452 + // INVARIANT: 453 + // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 454 + Cursor { 455 + current, 456 + _tree: PhantomData, 457 + } 458 + }) 459 + } 460 + 461 + fn find_best_match(&self, key: &K) -> Option<NonNull<bindings::rb_node>> { 458 462 let mut node = self.root.rb_node; 459 - let mut best_match: Option<NonNull<Node<K, V>>> = None; 463 + let mut best_key: Option<&K> = None; 464 + let mut best_links: Option<NonNull<bindings::rb_node>> = None; 460 465 while !node.is_null() { 461 466 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 462 467 // point to the links field of `Node<K, V>` objects. ··· 504 439 let right_child = unsafe { (*node).rb_right }; 505 440 match key.cmp(this_key) { 506 441 Ordering::Equal => { 507 - best_match = NonNull::new(this); 442 + // SAFETY: `this` is a non-null node so it is valid by the type invariants. 443 + best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) }); 508 444 break; 509 445 } 510 446 Ordering::Greater => { 511 447 node = right_child; 512 448 } 513 449 Ordering::Less => { 514 - let is_better_match = match best_match { 450 + let is_better_match = match best_key { 515 451 None => true, 516 - Some(best) => { 517 - // SAFETY: `best` is a non-null node so it is valid by the type invariants. 518 - let best_key = unsafe { &(*best.as_ptr()).key }; 519 - best_key > this_key 520 - } 452 + Some(best) => best > this_key, 521 453 }; 522 454 if is_better_match { 523 - best_match = NonNull::new(this); 455 + best_key = Some(this_key); 456 + // SAFETY: `this` is a non-null node so it is valid by the type invariants. 457 + best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) }); 524 458 } 525 459 node = left_child; 526 460 } 527 461 }; 528 462 } 529 - 530 - let best = best_match?; 531 - 532 - // SAFETY: `best` is a non-null node so it is valid by the type invariants. 533 - let links = unsafe { addr_of_mut!((*best.as_ptr()).links) }; 534 - 535 - NonNull::new(links).map(|current| { 536 - // INVARIANT: 537 - // - `current` is a valid node in the [`RBTree`] pointed to by `self`. 538 - Cursor { 539 - current, 540 - tree: self, 541 - } 542 - }) 463 + best_links 543 464 } 544 465 } 545 466 ··· 558 507 } 559 508 } 560 509 561 - /// A bidirectional cursor over the tree nodes, sorted by key. 510 + /// A bidirectional mutable cursor over the tree nodes, sorted by key. 562 511 /// 563 512 /// # Examples 564 513 /// ··· 577 526 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 578 527 /// 579 528 /// // Get a cursor to the first element. 580 - /// let mut cursor = tree.cursor_front().unwrap(); 529 + /// let mut cursor = tree.cursor_front_mut().unwrap(); 581 530 /// let mut current = cursor.current(); 582 531 /// assert_eq!(current, (&10, &100)); 583 532 /// ··· 615 564 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 616 565 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 617 566 /// 618 - /// let mut cursor = tree.cursor_back().unwrap(); 567 + /// let mut cursor = tree.cursor_back_mut().unwrap(); 619 568 /// let current = cursor.current(); 620 569 /// assert_eq!(current, (&30, &300)); 621 570 /// ··· 628 577 /// use kernel::rbtree::RBTree; 629 578 /// 630 579 /// let mut tree: RBTree<u16, u16> = RBTree::new(); 631 - /// assert!(tree.cursor_front().is_none()); 580 + /// assert!(tree.cursor_front_mut().is_none()); 632 581 /// 633 582 /// # Ok::<(), Error>(()) 634 583 /// ``` ··· 679 628 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 680 629 /// 681 630 /// // Retrieve a cursor. 682 - /// let mut cursor = tree.cursor_front().unwrap(); 631 + /// let mut cursor = tree.cursor_front_mut().unwrap(); 683 632 /// 684 633 /// // Get a mutable reference to the current value. 685 634 /// let (k, v) = cursor.current_mut(); ··· 706 655 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 707 656 /// 708 657 /// // Remove the first element. 709 - /// let mut cursor = tree.cursor_front().unwrap(); 658 + /// let mut cursor = tree.cursor_front_mut().unwrap(); 710 659 /// let mut current = cursor.current(); 711 660 /// assert_eq!(current, (&10, &100)); 712 661 /// cursor = cursor.remove_current().0.unwrap(); ··· 716 665 /// assert_eq!(current, (&20, &200)); 717 666 /// 718 667 /// // Get a cursor to the last element, and remove it. 719 - /// cursor = tree.cursor_back().unwrap(); 668 + /// cursor = tree.cursor_back_mut().unwrap(); 720 669 /// current = cursor.current(); 721 670 /// assert_eq!(current, (&30, &300)); 722 671 /// ··· 745 694 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 746 695 /// 747 696 /// // Get a cursor to the first element. 748 - /// let mut cursor = tree.cursor_front().unwrap(); 697 + /// let mut cursor = tree.cursor_front_mut().unwrap(); 749 698 /// let mut current = cursor.current(); 750 699 /// assert_eq!(current, (&10, &100)); 751 700 /// ··· 753 702 /// assert!(cursor.remove_prev().is_none()); 754 703 /// 755 704 /// // Get a cursor to the last element. 756 - /// cursor = tree.cursor_back().unwrap(); 705 + /// cursor = tree.cursor_back_mut().unwrap(); 757 706 /// current = cursor.current(); 758 707 /// assert_eq!(current, (&30, &300)); 759 708 /// ··· 777 726 /// 778 727 /// # Invariants 779 728 /// - `current` points to a node that is in the same [`RBTree`] as `tree`. 780 - pub struct Cursor<'a, K, V> { 729 + pub struct CursorMut<'a, K, V> { 781 730 tree: &'a mut RBTree<K, V>, 782 731 current: NonNull<bindings::rb_node>, 783 732 } 784 733 785 - // SAFETY: The [`Cursor`] has exclusive access to both `K` and `V`, so it is sufficient to require them to be `Send`. 786 - // The cursor only gives out immutable references to the keys, but since it has excusive access to those same 787 - // keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the user. 788 - unsafe impl<'a, K: Send, V: Send> Send for Cursor<'a, K, V> {} 734 + /// A bidirectional immutable cursor over the tree nodes, sorted by key. This is a simpler 735 + /// variant of [`CursorMut`] that is basically providing read only access. 736 + /// 737 + /// # Examples 738 + /// 739 + /// In the following example, we obtain a cursor to the first element in the tree. 740 + /// The cursor allows us to iterate bidirectionally over key/value pairs in the tree. 741 + /// 742 + /// ``` 743 + /// use kernel::{alloc::flags, rbtree::RBTree}; 744 + /// 745 + /// // Create a new tree. 746 + /// let mut tree = RBTree::new(); 747 + /// 748 + /// // Insert three elements. 749 + /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; 750 + /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; 751 + /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; 752 + /// 753 + /// // Get a cursor to the first element. 754 + /// let cursor = tree.cursor_front().unwrap(); 755 + /// let current = cursor.current(); 756 + /// assert_eq!(current, (&10, &100)); 757 + /// 758 + /// # Ok::<(), Error>(()) 759 + /// ``` 760 + pub struct Cursor<'a, K, V> { 761 + _tree: PhantomData<&'a RBTree<K, V>>, 762 + current: NonNull<bindings::rb_node>, 763 + } 789 764 790 - // SAFETY: The [`Cursor`] gives out immutable references to K and mutable references to V, 791 - // so it has the same thread safety requirements as mutable references. 765 + // SAFETY: The immutable cursor gives out shared access to `K` and `V` so if `K` and `V` can be 766 + // shared across threads, then it's safe to share the cursor. 767 + unsafe impl<'a, K: Sync, V: Sync> Send for Cursor<'a, K, V> {} 768 + 769 + // SAFETY: The immutable cursor gives out shared access to `K` and `V` so if `K` and `V` can be 770 + // shared across threads, then it's safe to share the cursor. 792 771 unsafe impl<'a, K: Sync, V: Sync> Sync for Cursor<'a, K, V> {} 793 772 794 773 impl<'a, K, V> Cursor<'a, K, V> { 795 774 /// The current node 775 + pub fn current(&self) -> (&K, &V) { 776 + // SAFETY: 777 + // - `self.current` is a valid node by the type invariants. 778 + // - We have an immutable reference by the function signature. 779 + unsafe { Self::to_key_value(self.current) } 780 + } 781 + 782 + /// # Safety 783 + /// 784 + /// - `node` must be a valid pointer to a node in an [`RBTree`]. 785 + /// - The caller has immutable access to `node` for the duration of `'b`. 786 + unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) { 787 + // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 788 + // point to the links field of `Node<K, V>` objects. 789 + let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) }; 790 + // SAFETY: The passed `node` is the current node or a non-null neighbor, 791 + // thus `this` is valid by the type invariants. 792 + let k = unsafe { &(*this).key }; 793 + // SAFETY: The passed `node` is the current node or a non-null neighbor, 794 + // thus `this` is valid by the type invariants. 795 + let v = unsafe { &(*this).value }; 796 + (k, v) 797 + } 798 + 799 + /// Access the previous node without moving the cursor. 800 + pub fn peek_prev(&self) -> Option<(&K, &V)> { 801 + self.peek(Direction::Prev) 802 + } 803 + 804 + /// Access the next node without moving the cursor. 805 + pub fn peek_next(&self) -> Option<(&K, &V)> { 806 + self.peek(Direction::Next) 807 + } 808 + 809 + fn peek(&self, direction: Direction) -> Option<(&K, &V)> { 810 + self.get_neighbor_raw(direction).map(|neighbor| { 811 + // SAFETY: 812 + // - `neighbor` is a valid tree node. 813 + // - By the function signature, we have an immutable reference to `self`. 814 + unsafe { Self::to_key_value(neighbor) } 815 + }) 816 + } 817 + 818 + fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> { 819 + // SAFETY: `self.current` is valid by the type invariants. 820 + let neighbor = unsafe { 821 + match direction { 822 + Direction::Prev => bindings::rb_prev(self.current.as_ptr()), 823 + Direction::Next => bindings::rb_next(self.current.as_ptr()), 824 + } 825 + }; 826 + 827 + NonNull::new(neighbor) 828 + } 829 + } 830 + 831 + // SAFETY: The [`CursorMut`] has exclusive access to both `K` and `V`, so it is sufficient to 832 + // require them to be `Send`. 833 + // The cursor only gives out immutable references to the keys, but since it has exclusive access to 834 + // those same keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the 835 + // user. 836 + unsafe impl<'a, K: Send, V: Send> Send for CursorMut<'a, K, V> {} 837 + 838 + // SAFETY: The [`CursorMut`] gives out immutable references to `K` and mutable references to `V`, 839 + // so it has the same thread safety requirements as mutable references. 840 + unsafe impl<'a, K: Sync, V: Sync> Sync for CursorMut<'a, K, V> {} 841 + 842 + impl<'a, K, V> CursorMut<'a, K, V> { 843 + /// The current node. 796 844 pub fn current(&self) -> (&K, &V) { 797 845 // SAFETY: 798 846 // - `self.current` is a valid node by the type invariants. ··· 1070 920 } 1071 921 } 1072 922 1073 - /// Direction for [`Cursor`] operations. 923 + /// Direction for [`Cursor`] and [`CursorMut`] operations. 1074 924 enum Direction { 1075 925 /// the node immediately before, in sort order 1076 926 Prev,

Configure Feed

Configure Feed