rust: rbtree: add `RBTree::entry` · tjh.dev/kernel@a335e95

+231 -76

1 changed file

expand all

rust

kernel

rbtree.rs

+231 -76

rust/kernel/rbtree.rs

··· 293 293 /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. 294 294 /// 295 295 /// This function always succeeds. 296 - pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> { 297 - let node = Box::into_raw(node); 298 - // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when 299 - // the node is removed or replaced. 300 - let node_links = unsafe { addr_of_mut!((*node).links) }; 296 + pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> { 297 + match self.raw_entry(&node.node.key) { 298 + RawEntry::Occupied(entry) => Some(entry.replace(node)), 299 + RawEntry::Vacant(entry) => { 300 + entry.insert(node); 301 + None 302 + } 303 + } 304 + } 301 305 306 + fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> { 307 + let raw_self: *mut RBTree<K, V> = self; 308 + // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`. 302 309 // The parameters of `bindings::rb_link_node` are as follows: 303 310 // - `node`: A pointer to an uninitialized node being inserted. 304 311 // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be ··· 324 317 // in the subtree of `parent` that `child_field_of_parent` points at. Once 325 318 // we find an empty subtree, we can insert the new node using `rb_link_node`. 326 319 let mut parent = core::ptr::null_mut(); 327 - let mut child_field_of_parent: &mut *mut bindings::rb_node = &mut self.root.rb_node; 328 - while !child_field_of_parent.is_null() { 329 - parent = *child_field_of_parent; 320 + let mut child_field_of_parent: &mut *mut bindings::rb_node = 321 + // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above). 322 + unsafe { &mut (*raw_self).root.rb_node }; 323 + while !(*child_field_of_parent).is_null() { 324 + let curr = *child_field_of_parent; 325 + // SAFETY: All links fields we create are in a `Node<K, V>`. 326 + let node = unsafe { container_of!(curr, Node<K, V>, links) }; 330 327 331 - // We need to determine whether `node` should be the left or right child of `parent`, 332 - // so we will compare with the `key` field of `parent` a.k.a. `this` below. 333 - // 334 - // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` 335 - // point to the links field of `Node<K, V>` objects. 336 - let this = unsafe { container_of!(parent, Node<K, V>, links) }; 337 - 338 - // SAFETY: `this` is a non-null node so it is valid by the type invariants. `node` is 339 - // valid until the node is removed. 340 - match unsafe { (*node).key.cmp(&(*this).key) } { 341 - // We would like `node` to be the left child of `parent`. Move to this child to check 342 - // whether we can use it, or continue searching, at the next iteration. 343 - // 344 - // SAFETY: `parent` is a non-null node so it is valid by the type invariants. 345 - Ordering::Less => child_field_of_parent = unsafe { &mut (*parent).rb_left }, 346 - // We would like `node` to be the right child of `parent`. Move to this child to check 347 - // whether we can use it, or continue searching, at the next iteration. 348 - // 349 - // SAFETY: `parent` is a non-null node so it is valid by the type invariants. 350 - Ordering::Greater => child_field_of_parent = unsafe { &mut (*parent).rb_right }, 328 + // SAFETY: `node` is a non-null node so it is valid by the type invariants. 329 + match key.cmp(unsafe { &(*node).key }) { 330 + // SAFETY: `curr` is a non-null node so it is valid by the type invariants. 331 + Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left }, 332 + // SAFETY: `curr` is a non-null node so it is valid by the type invariants. 333 + Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).rb_right }, 351 334 Ordering::Equal => { 352 - // There is an existing node in the tree with this key, and that node is 353 - // `parent`. Thus, we are replacing parent with a new node. 354 - // 355 - // INVARIANT: We are replacing an existing node with a new one, which is valid. 356 - // It remains valid because we "forgot" it with `Box::into_raw`. 357 - // SAFETY: All pointers are non-null and valid. 358 - unsafe { bindings::rb_replace_node(parent, node_links, &mut self.root) }; 359 - 360 - // INVARIANT: The node is being returned and the caller may free it, however, 361 - // it was removed from the tree. So the invariants still hold. 362 - return Some(RBTreeNode { 363 - // SAFETY: `this` was a node in the tree, so it is valid. 364 - node: unsafe { Box::from_raw(this.cast_mut()) }, 365 - }); 335 + return RawEntry::Occupied(OccupiedEntry { 336 + rbtree: self, 337 + node_links: curr, 338 + }) 366 339 } 367 340 } 341 + parent = curr; 368 342 } 369 343 370 - // INVARIANT: We are linking in a new node, which is valid. It remains valid because we 371 - // "forgot" it with `Box::into_raw`. 372 - // SAFETY: All pointers are non-null and valid (`*child_field_of_parent` is null, but `child_field_of_parent` is a 373 - // mutable reference). 374 - unsafe { bindings::rb_link_node(node_links, parent, child_field_of_parent) }; 375 - 376 - // SAFETY: All pointers are valid. `node` has just been inserted into the tree. 377 - unsafe { bindings::rb_insert_color(node_links, &mut self.root) }; 378 - None 344 + RawEntry::Vacant(RawVacantEntry { 345 + rbtree: raw_self, 346 + parent, 347 + child_field_of_parent, 348 + _phantom: PhantomData, 349 + }) 379 350 } 380 351 381 - /// Returns a node with the given key, if one exists. 382 - fn find(&self, key: &K) -> Option<NonNull<Node<K, V>>> { 352 + /// Gets the given key's corresponding entry in the map for in-place manipulation. 353 + pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { 354 + match self.raw_entry(&key) { 355 + RawEntry::Occupied(entry) => Entry::Occupied(entry), 356 + RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }), 357 + } 358 + } 359 + 360 + /// Used for accessing the given node, if it exists. 361 + pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> { 362 + match self.raw_entry(key) { 363 + RawEntry::Occupied(entry) => Some(entry), 364 + RawEntry::Vacant(_entry) => None, 365 + } 366 + } 367 + 368 + /// Returns a reference to the value corresponding to the key. 369 + pub fn get(&self, key: &K) -> Option<&V> { 383 370 let mut node = self.root.rb_node; 384 371 while !node.is_null() { 385 372 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` ··· 385 384 Ordering::Less => unsafe { (*node).rb_left }, 386 385 // SAFETY: `node` is a non-null node so it is valid by the type invariants. 387 386 Ordering::Greater => unsafe { (*node).rb_right }, 388 - Ordering::Equal => return NonNull::new(this.cast_mut()), 387 + // SAFETY: `node` is a non-null node so it is valid by the type invariants. 388 + Ordering::Equal => return Some(unsafe { &(*this).value }), 389 389 } 390 390 } 391 391 None 392 392 } 393 393 394 - /// Returns a reference to the value corresponding to the key. 395 - pub fn get(&self, key: &K) -> Option<&V> { 396 - // SAFETY: The `find` return value is a node in the tree, so it is valid. 397 - self.find(key).map(|node| unsafe { &node.as_ref().value }) 398 - } 399 - 400 394 /// Returns a mutable reference to the value corresponding to the key. 401 395 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { 402 - // SAFETY: The `find` return value is a node in the tree, so it is valid. 403 - self.find(key) 404 - .map(|mut node| unsafe { &mut node.as_mut().value }) 396 + self.find_mut(key).map(|node| node.into_mut()) 405 397 } 406 398 407 399 /// Removes the node with the given key from the tree. 408 400 /// 409 401 /// It returns the node that was removed if one exists, or [`None`] otherwise. 410 - fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> { 411 - let mut node = self.find(key)?; 412 - 413 - // SAFETY: The `find` return value is a node in the tree, so it is valid. 414 - unsafe { bindings::rb_erase(&mut node.as_mut().links, &mut self.root) }; 415 - 416 - // INVARIANT: The node is being returned and the caller may free it, however, it was 417 - // removed from the tree. So the invariants still hold. 418 - Some(RBTreeNode { 419 - // SAFETY: The `find` return value was a node in the tree, so it is valid. 420 - node: unsafe { Box::from_raw(node.as_ptr()) }, 421 - }) 402 + pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> { 403 + self.find_mut(key).map(OccupiedEntry::remove_node) 422 404 } 423 405 424 406 /// Removes the node with the given key from the tree. 425 407 /// 426 408 /// It returns the value that was removed if one exists, or [`None`] otherwise. 427 409 pub fn remove(&mut self, key: &K) -> Option<V> { 428 - self.remove_node(key).map(|node| node.node.value) 410 + self.find_mut(key).map(OccupiedEntry::remove) 429 411 } 430 412 431 413 /// Returns a cursor over the tree nodes based on the given key. ··· 1100 1116 // SAFETY: If K and V can be accessed without synchronization, then it's also okay to access 1101 1117 // [`RBTreeNode`] without synchronization. 1102 1118 unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {} 1119 + 1120 + impl<K, V> RBTreeNode<K, V> { 1121 + /// Drop the key and value, but keep the allocation. 1122 + /// 1123 + /// It then becomes a reservation that can be re-initialised into a different node (i.e., with 1124 + /// a different key and/or value). 1125 + /// 1126 + /// The existing key and value are dropped in-place as part of this operation, that is, memory 1127 + /// may be freed (but only for the key/value; memory for the node itself is kept for reuse). 1128 + pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> { 1129 + RBTreeNodeReservation { 1130 + node: Box::drop_contents(self.node), 1131 + } 1132 + } 1133 + } 1134 + 1135 + /// A view into a single entry in a map, which may either be vacant or occupied. 1136 + /// 1137 + /// This enum is constructed from the [`RBTree::entry`]. 1138 + /// 1139 + /// [`entry`]: fn@RBTree::entry 1140 + pub enum Entry<'a, K, V> { 1141 + /// This [`RBTree`] does not have a node with this key. 1142 + Vacant(VacantEntry<'a, K, V>), 1143 + /// This [`RBTree`] already has a node with this key. 1144 + Occupied(OccupiedEntry<'a, K, V>), 1145 + } 1146 + 1147 + /// Like [`Entry`], except that it doesn't have ownership of the key. 1148 + enum RawEntry<'a, K, V> { 1149 + Vacant(RawVacantEntry<'a, K, V>), 1150 + Occupied(OccupiedEntry<'a, K, V>), 1151 + } 1152 + 1153 + /// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`] enum. 1154 + pub struct VacantEntry<'a, K, V> { 1155 + key: K, 1156 + raw: RawVacantEntry<'a, K, V>, 1157 + } 1158 + 1159 + /// Like [`VacantEntry`], but doesn't hold on to the key. 1160 + /// 1161 + /// # Invariants 1162 + /// - `parent` may be null if the new node becomes the root. 1163 + /// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `parent` is 1164 + /// null, it is a pointer to the root of the [`RBTree`]. 1165 + struct RawVacantEntry<'a, K, V> { 1166 + rbtree: *mut RBTree<K, V>, 1167 + /// The node that will become the parent of the new node if we insert one. 1168 + parent: *mut bindings::rb_node, 1169 + /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is 1170 + /// null. 1171 + child_field_of_parent: *mut *mut bindings::rb_node, 1172 + _phantom: PhantomData<&'a mut RBTree<K, V>>, 1173 + } 1174 + 1175 + impl<'a, K, V> RawVacantEntry<'a, K, V> { 1176 + /// Inserts the given node into the [`RBTree`] at this entry. 1177 + /// 1178 + /// The `node` must have a key such that inserting it here does not break the ordering of this 1179 + /// [`RBTree`]. 1180 + fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V { 1181 + let node = Box::into_raw(node.node); 1182 + 1183 + // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when 1184 + // the node is removed or replaced. 1185 + let node_links = unsafe { addr_of_mut!((*node).links) }; 1186 + 1187 + // INVARIANT: We are linking in a new node, which is valid. It remains valid because we 1188 + // "forgot" it with `Box::into_raw`. 1189 + // SAFETY: The type invariants of `RawVacantEntry` are exactly the safety requirements of `rb_link_node`. 1190 + unsafe { bindings::rb_link_node(node_links, self.parent, self.child_field_of_parent) }; 1191 + 1192 + // SAFETY: All pointers are valid. `node` has just been inserted into the tree. 1193 + unsafe { bindings::rb_insert_color(node_links, addr_of_mut!((*self.rbtree).root)) }; 1194 + 1195 + // SAFETY: The node is valid until we remove it from the tree. 1196 + unsafe { &mut (*node).value } 1197 + } 1198 + } 1199 + 1200 + impl<'a, K, V> VacantEntry<'a, K, V> { 1201 + /// Inserts the given node into the [`RBTree`] at this entry. 1202 + pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V { 1203 + self.raw.insert(reservation.into_node(self.key, value)) 1204 + } 1205 + } 1206 + 1207 + /// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum. 1208 + /// 1209 + /// # Invariants 1210 + /// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree` 1211 + pub struct OccupiedEntry<'a, K, V> { 1212 + rbtree: &'a mut RBTree<K, V>, 1213 + /// The node that this entry corresponds to. 1214 + node_links: *mut bindings::rb_node, 1215 + } 1216 + 1217 + impl<'a, K, V> OccupiedEntry<'a, K, V> { 1218 + /// Gets a reference to the value in the entry. 1219 + pub fn get(&self) -> &V { 1220 + // SAFETY: 1221 + // - `self.node_links` is a valid pointer to a node in the tree. 1222 + // - We have shared access to the underlying tree, and can thus give out a shared reference. 1223 + unsafe { &(*container_of!(self.node_links, Node<K, V>, links)).value } 1224 + } 1225 + 1226 + /// Gets a mutable reference to the value in the entry. 1227 + pub fn get_mut(&mut self) -> &mut V { 1228 + // SAFETY: 1229 + // - `self.node_links` is a valid pointer to a node in the tree. 1230 + // - We have exclusive access to the underlying tree, and can thus give out a mutable reference. 1231 + unsafe { &mut (*(container_of!(self.node_links, Node<K, V>, links).cast_mut())).value } 1232 + } 1233 + 1234 + /// Converts the entry into a mutable reference to its value. 1235 + /// 1236 + /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`]. 1237 + pub fn into_mut(self) -> &'a mut V { 1238 + // SAFETY: 1239 + // - `self.node_links` is a valid pointer to a node in the tree. 1240 + // - This consumes the `&'a mut RBTree<K, V>`, therefore it can give out a mutable reference that lives for `'a`. 1241 + unsafe { &mut (*(container_of!(self.node_links, Node<K, V>, links).cast_mut())).value } 1242 + } 1243 + 1244 + /// Remove this entry from the [`RBTree`]. 1245 + pub fn remove_node(self) -> RBTreeNode<K, V> { 1246 + // SAFETY: The node is a node in the tree, so it is valid. 1247 + unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.root) }; 1248 + 1249 + // INVARIANT: The node is being returned and the caller may free it, however, it was 1250 + // removed from the tree. So the invariants still hold. 1251 + RBTreeNode { 1252 + // SAFETY: The node was a node in the tree, but we removed it, so we can convert it 1253 + // back into a box. 1254 + node: unsafe { 1255 + Box::from_raw(container_of!(self.node_links, Node<K, V>, links).cast_mut()) 1256 + }, 1257 + } 1258 + } 1259 + 1260 + /// Takes the value of the entry out of the map, and returns it. 1261 + pub fn remove(self) -> V { 1262 + self.remove_node().node.value 1263 + } 1264 + 1265 + /// Swap the current node for the provided node. 1266 + /// 1267 + /// The key of both nodes must be equal. 1268 + fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> { 1269 + let node = Box::into_raw(node.node); 1270 + 1271 + // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when 1272 + // the node is removed or replaced. 1273 + let new_node_links = unsafe { addr_of_mut!((*node).links) }; 1274 + 1275 + // SAFETY: This updates the pointers so that `new_node_links` is in the tree where 1276 + // `self.node_links` used to be. 1277 + unsafe { 1278 + bindings::rb_replace_node(self.node_links, new_node_links, &mut self.rbtree.root) 1279 + }; 1280 + 1281 + // SAFETY: 1282 + // - `self.node_ptr` produces a valid pointer to a node in the tree. 1283 + // - Now that we removed this entry from the tree, we can convert the node to a box. 1284 + let old_node = 1285 + unsafe { Box::from_raw(container_of!(self.node_links, Node<K, V>, links).cast_mut()) }; 1286 + 1287 + RBTreeNode { node: old_node } 1288 + } 1289 + } 1103 1290 1104 1291 struct Node<K, V> { 1105 1292 links: bindings::rb_node,

Configure Feed

Configure Feed