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