1 //! Forest of sets. 2 3 use super::{Comparator, Forest, Node, NodeData, NodePool, Path, SetValue, INNER_SIZE}; 4 use crate::packed_option::PackedOption; 5 #[cfg(test)] 6 use alloc::string::String; 7 #[cfg(test)] 8 use core::fmt; 9 use core::marker::PhantomData; 10 11 /// Tag type defining forest types for a set. 12 struct SetTypes<K>(PhantomData<K>); 13 14 impl<K> Forest for SetTypes<K> 15 where 16 K: Copy, 17 { 18 type Key = K; 19 type Value = SetValue; 20 type LeafKeys = [K; 2 * INNER_SIZE - 1]; 21 type LeafValues = [SetValue; 2 * INNER_SIZE - 1]; 22 splat_key(key: Self::Key) -> Self::LeafKeys23 fn splat_key(key: Self::Key) -> Self::LeafKeys { 24 [key; 2 * INNER_SIZE - 1] 25 } 26 splat_value(value: Self::Value) -> Self::LeafValues27 fn splat_value(value: Self::Value) -> Self::LeafValues { 28 [value; 2 * INNER_SIZE - 1] 29 } 30 } 31 32 /// Memory pool for a forest of `Set` instances. 33 pub struct SetForest<K> 34 where 35 K: Copy, 36 { 37 nodes: NodePool<SetTypes<K>>, 38 } 39 40 impl<K> SetForest<K> 41 where 42 K: Copy, 43 { 44 /// Create a new empty forest. new() -> Self45 pub fn new() -> Self { 46 Self { 47 nodes: NodePool::new(), 48 } 49 } 50 51 /// Clear all sets in the forest. 52 /// 53 /// All `Set` instances belong to this forest are invalidated and should no longer be used. clear(&mut self)54 pub fn clear(&mut self) { 55 self.nodes.clear(); 56 } 57 } 58 59 /// B-tree representing an ordered set of `K`s using `C` for comparing elements. 60 /// 61 /// This is not a general-purpose replacement for `BTreeSet`. See the [module 62 /// documentation](index.html) for more information about design tradeoffs. 63 /// 64 /// Sets can be cloned, but that operation should only be used as part of cloning the whole forest 65 /// they belong to. *Cloning a set does not allocate new memory for the clone*. It creates an alias 66 /// of the same memory. 67 #[derive(Clone)] 68 pub struct Set<K> 69 where 70 K: Copy, 71 { 72 root: PackedOption<Node>, 73 unused: PhantomData<K>, 74 } 75 76 impl<K> Set<K> 77 where 78 K: Copy, 79 { 80 /// Make an empty set. new() -> Self81 pub fn new() -> Self { 82 Self { 83 root: None.into(), 84 unused: PhantomData, 85 } 86 } 87 88 /// Is this an empty set? is_empty(&self) -> bool89 pub fn is_empty(&self) -> bool { 90 self.root.is_none() 91 } 92 93 /// Does the set contain `key`?. contains<C: Comparator<K>>(&self, key: K, forest: &SetForest<K>, comp: &C) -> bool94 pub fn contains<C: Comparator<K>>(&self, key: K, forest: &SetForest<K>, comp: &C) -> bool { 95 self.root 96 .expand() 97 .and_then(|root| Path::default().find(key, root, &forest.nodes, comp)) 98 .is_some() 99 } 100 101 /// Try to insert `key` into the set. 102 /// 103 /// If the set did not contain `key`, insert it and return true. 104 /// 105 /// If `key` is already present, don't change the set and return false. insert<C: Comparator<K>>( &mut self, key: K, forest: &mut SetForest<K>, comp: &C, ) -> bool106 pub fn insert<C: Comparator<K>>( 107 &mut self, 108 key: K, 109 forest: &mut SetForest<K>, 110 comp: &C, 111 ) -> bool { 112 self.cursor(forest, comp).insert(key) 113 } 114 115 /// Remove `key` from the set and return true. 116 /// 117 /// If `key` was not present in the set, return false. remove<C: Comparator<K>>( &mut self, key: K, forest: &mut SetForest<K>, comp: &C, ) -> bool118 pub fn remove<C: Comparator<K>>( 119 &mut self, 120 key: K, 121 forest: &mut SetForest<K>, 122 comp: &C, 123 ) -> bool { 124 let mut c = self.cursor(forest, comp); 125 if c.goto(key) { 126 c.remove(); 127 true 128 } else { 129 false 130 } 131 } 132 133 /// Remove all entries. clear(&mut self, forest: &mut SetForest<K>)134 pub fn clear(&mut self, forest: &mut SetForest<K>) { 135 if let Some(root) = self.root.take() { 136 forest.nodes.free_tree(root); 137 } 138 } 139 140 /// Retains only the elements specified by the predicate. 141 /// 142 /// Remove all elements where the predicate returns false. retain<F>(&mut self, forest: &mut SetForest<K>, mut predicate: F) where F: FnMut(K) -> bool,143 pub fn retain<F>(&mut self, forest: &mut SetForest<K>, mut predicate: F) 144 where 145 F: FnMut(K) -> bool, 146 { 147 let mut path = Path::default(); 148 if let Some(root) = self.root.expand() { 149 path.first(root, &forest.nodes); 150 } 151 while let Some((node, entry)) = path.leaf_pos() { 152 if predicate(forest.nodes[node].unwrap_leaf().0[entry]) { 153 path.next(&forest.nodes); 154 } else { 155 self.root = path.remove(&mut forest.nodes).into(); 156 } 157 } 158 } 159 160 /// Create a cursor for navigating this set. The cursor is initially positioned off the end of 161 /// the set. cursor<'a, C: Comparator<K>>( &'a mut self, forest: &'a mut SetForest<K>, comp: &'a C, ) -> SetCursor<'a, K, C>162 pub fn cursor<'a, C: Comparator<K>>( 163 &'a mut self, 164 forest: &'a mut SetForest<K>, 165 comp: &'a C, 166 ) -> SetCursor<'a, K, C> { 167 SetCursor::new(self, forest, comp) 168 } 169 170 /// Create an iterator traversing this set. The iterator type is `K`. iter<'a>(&'a self, forest: &'a SetForest<K>) -> SetIter<'a, K>171 pub fn iter<'a>(&'a self, forest: &'a SetForest<K>) -> SetIter<'a, K> { 172 SetIter { 173 root: self.root, 174 pool: &forest.nodes, 175 path: Path::default(), 176 } 177 } 178 } 179 180 impl<K> Default for Set<K> 181 where 182 K: Copy, 183 { default() -> Self184 fn default() -> Self { 185 Self::new() 186 } 187 } 188 189 /// A position in a `Set` used to navigate and modify the ordered set. 190 /// 191 /// A cursor always points at an element in the set, or "off the end" which is a position after the 192 /// last element in the set. 193 pub struct SetCursor<'a, K, C> 194 where 195 K: 'a + Copy, 196 C: 'a + Comparator<K>, 197 { 198 root: &'a mut PackedOption<Node>, 199 pool: &'a mut NodePool<SetTypes<K>>, 200 comp: &'a C, 201 path: Path<SetTypes<K>>, 202 } 203 204 impl<'a, K, C> SetCursor<'a, K, C> 205 where 206 K: Copy, 207 C: Comparator<K>, 208 { 209 /// Create a cursor with a default (invalid) location. new(container: &'a mut Set<K>, forest: &'a mut SetForest<K>, comp: &'a C) -> Self210 fn new(container: &'a mut Set<K>, forest: &'a mut SetForest<K>, comp: &'a C) -> Self { 211 Self { 212 root: &mut container.root, 213 pool: &mut forest.nodes, 214 comp, 215 path: Path::default(), 216 } 217 } 218 219 /// Is this cursor pointing to an empty set? is_empty(&self) -> bool220 pub fn is_empty(&self) -> bool { 221 self.root.is_none() 222 } 223 224 /// Move cursor to the next element and return it. 225 /// 226 /// If the cursor reaches the end, return `None` and leave the cursor at the off-the-end 227 /// position. 228 #[cfg_attr(feature = "cargo-clippy", allow(clippy::should_implement_trait))] next(&mut self) -> Option<K>229 pub fn next(&mut self) -> Option<K> { 230 self.path.next(self.pool).map(|(k, _)| k) 231 } 232 233 /// Move cursor to the previous element and return it. 234 /// 235 /// If the cursor is already pointing at the first element, leave it there and return `None`. prev(&mut self) -> Option<K>236 pub fn prev(&mut self) -> Option<K> { 237 self.root 238 .expand() 239 .and_then(|root| self.path.prev(root, self.pool).map(|(k, _)| k)) 240 } 241 242 /// Get the current element, or `None` if the cursor is at the end. elem(&self) -> Option<K>243 pub fn elem(&self) -> Option<K> { 244 self.path 245 .leaf_pos() 246 .and_then(|(node, entry)| self.pool[node].unwrap_leaf().0.get(entry).cloned()) 247 } 248 249 /// Move this cursor to `elem`. 250 /// 251 /// If `elem` is in the set, place the cursor at `elem` and return true. 252 /// 253 /// If `elem` is not in the set, place the cursor at the next larger element (or the end) and 254 /// return false. goto(&mut self, elem: K) -> bool255 pub fn goto(&mut self, elem: K) -> bool { 256 match self.root.expand() { 257 None => false, 258 Some(root) => { 259 if self.path.find(elem, root, self.pool, self.comp).is_some() { 260 true 261 } else { 262 self.path.normalize(self.pool); 263 false 264 } 265 } 266 } 267 } 268 269 /// Move this cursor to the first element. goto_first(&mut self) -> Option<K>270 pub fn goto_first(&mut self) -> Option<K> { 271 self.root.map(|root| self.path.first(root, self.pool).0) 272 } 273 274 /// Try to insert `elem` into the set and leave the cursor at the inserted element. 275 /// 276 /// If the set did not contain `elem`, insert it and return true. 277 /// 278 /// If `elem` is already present, don't change the set, place the cursor at `goto(elem)`, and 279 /// return false. insert(&mut self, elem: K) -> bool280 pub fn insert(&mut self, elem: K) -> bool { 281 match self.root.expand() { 282 None => { 283 let root = self.pool.alloc_node(NodeData::leaf(elem, SetValue())); 284 *self.root = root.into(); 285 self.path.set_root_node(root); 286 true 287 } 288 Some(root) => { 289 // TODO: Optimize the case where `self.path` is already at the correct insert pos. 290 if self.path.find(elem, root, self.pool, self.comp).is_none() { 291 *self.root = self.path.insert(elem, SetValue(), self.pool).into(); 292 true 293 } else { 294 false 295 } 296 } 297 } 298 } 299 300 /// Remove the current element (if any) and return it. 301 /// This advances the cursor to the next element after the removed one. remove(&mut self) -> Option<K>302 pub fn remove(&mut self) -> Option<K> { 303 let elem = self.elem(); 304 if elem.is_some() { 305 *self.root = self.path.remove(self.pool).into(); 306 } 307 elem 308 } 309 } 310 311 #[cfg(test)] 312 impl<'a, K, C> SetCursor<'a, K, C> 313 where 314 K: Copy + fmt::Display, 315 C: Comparator<K>, 316 { verify(&self)317 fn verify(&self) { 318 self.path.verify(self.pool); 319 self.root.map(|root| self.pool.verify_tree(root, self.comp)); 320 } 321 322 /// Get a text version of the path to the current position. tpath(&self) -> String323 fn tpath(&self) -> String { 324 use alloc::string::ToString; 325 self.path.to_string() 326 } 327 } 328 329 /// An iterator visiting the elements of a `Set`. 330 pub struct SetIter<'a, K> 331 where 332 K: 'a + Copy, 333 { 334 root: PackedOption<Node>, 335 pool: &'a NodePool<SetTypes<K>>, 336 path: Path<SetTypes<K>>, 337 } 338 339 impl<'a, K> Iterator for SetIter<'a, K> 340 where 341 K: 'a + Copy, 342 { 343 type Item = K; 344 next(&mut self) -> Option<Self::Item>345 fn next(&mut self) -> Option<Self::Item> { 346 // We use `self.root` to indicate if we need to go to the first element. Reset to `None` 347 // once we've returned the first element. This also works for an empty tree since the 348 // `path.next()` call returns `None` when the path is empty. This also fuses the iterator. 349 match self.root.take() { 350 Some(root) => Some(self.path.first(root, self.pool).0), 351 None => self.path.next(self.pool).map(|(k, _)| k), 352 } 353 } 354 } 355 356 #[cfg(test)] 357 mod tests { 358 use super::super::NodeData; 359 use super::*; 360 use alloc::vec::Vec; 361 use core::mem; 362 363 #[test] node_size()364 fn node_size() { 365 // check that nodes are cache line sized when keys are 32 bits. 366 type F = SetTypes<u32>; 367 assert_eq!(mem::size_of::<NodeData<F>>(), 64); 368 } 369 370 #[test] empty()371 fn empty() { 372 let mut f = SetForest::<u32>::new(); 373 f.clear(); 374 375 let mut s = Set::<u32>::new(); 376 assert!(s.is_empty()); 377 s.clear(&mut f); 378 assert!(!s.contains(7, &f, &())); 379 380 // Iterator for an empty set. 381 assert_eq!(s.iter(&f).next(), None); 382 383 s.retain(&mut f, |_| unreachable!()); 384 385 let mut c = SetCursor::new(&mut s, &mut f, &()); 386 c.verify(); 387 assert_eq!(c.elem(), None); 388 389 assert_eq!(c.goto_first(), None); 390 assert_eq!(c.tpath(), "<empty path>"); 391 } 392 393 #[test] simple_cursor()394 fn simple_cursor() { 395 let mut f = SetForest::<u32>::new(); 396 let mut s = Set::<u32>::new(); 397 let mut c = SetCursor::new(&mut s, &mut f, &()); 398 399 assert!(c.insert(50)); 400 c.verify(); 401 assert_eq!(c.elem(), Some(50)); 402 403 assert!(c.insert(100)); 404 c.verify(); 405 assert_eq!(c.elem(), Some(100)); 406 407 assert!(c.insert(10)); 408 c.verify(); 409 assert_eq!(c.elem(), Some(10)); 410 411 // Basic movement. 412 assert_eq!(c.next(), Some(50)); 413 assert_eq!(c.next(), Some(100)); 414 assert_eq!(c.next(), None); 415 assert_eq!(c.next(), None); 416 assert_eq!(c.prev(), Some(100)); 417 assert_eq!(c.prev(), Some(50)); 418 assert_eq!(c.prev(), Some(10)); 419 assert_eq!(c.prev(), None); 420 assert_eq!(c.prev(), None); 421 422 assert!(c.goto(50)); 423 assert_eq!(c.elem(), Some(50)); 424 assert_eq!(c.remove(), Some(50)); 425 c.verify(); 426 427 assert_eq!(c.elem(), Some(100)); 428 assert_eq!(c.remove(), Some(100)); 429 c.verify(); 430 assert_eq!(c.elem(), None); 431 assert_eq!(c.remove(), None); 432 c.verify(); 433 } 434 435 #[test] two_level_sparse_tree()436 fn two_level_sparse_tree() { 437 let mut f = SetForest::<u32>::new(); 438 let mut s = Set::<u32>::new(); 439 let mut c = SetCursor::new(&mut s, &mut f, &()); 440 441 // Insert enough elements that we get a two-level tree. 442 // Each leaf node holds 8 elements 443 assert!(c.is_empty()); 444 for i in 0..50 { 445 assert!(c.insert(i)); 446 assert_eq!(c.elem(), Some(i)); 447 } 448 assert!(!c.is_empty()); 449 450 assert_eq!(c.goto_first(), Some(0)); 451 assert_eq!(c.tpath(), "node2[0]--node0[0]"); 452 453 assert_eq!(c.prev(), None); 454 for i in 1..50 { 455 assert_eq!(c.next(), Some(i)); 456 } 457 assert_eq!(c.next(), None); 458 for i in (0..50).rev() { 459 assert_eq!(c.prev(), Some(i)); 460 } 461 assert_eq!(c.prev(), None); 462 463 assert!(c.goto(25)); 464 for i in 25..50 { 465 assert_eq!(c.remove(), Some(i)); 466 assert!(!c.is_empty()); 467 c.verify(); 468 } 469 470 for i in (0..25).rev() { 471 assert!(!c.is_empty()); 472 assert_eq!(c.elem(), None); 473 assert_eq!(c.prev(), Some(i)); 474 assert_eq!(c.remove(), Some(i)); 475 c.verify(); 476 } 477 assert_eq!(c.elem(), None); 478 assert!(c.is_empty()); 479 } 480 481 #[test] three_level_sparse_tree()482 fn three_level_sparse_tree() { 483 let mut f = SetForest::<u32>::new(); 484 let mut s = Set::<u32>::new(); 485 let mut c = SetCursor::new(&mut s, &mut f, &()); 486 487 // Insert enough elements that we get a 3-level tree. 488 // Each leaf node holds 8 elements when filled up sequentially. 489 // Inner nodes hold 8 node pointers. 490 assert!(c.is_empty()); 491 for i in 0..150 { 492 assert!(c.insert(i)); 493 assert_eq!(c.elem(), Some(i)); 494 } 495 assert!(!c.is_empty()); 496 497 assert!(c.goto(0)); 498 assert_eq!(c.tpath(), "node11[0]--node2[0]--node0[0]"); 499 500 assert_eq!(c.prev(), None); 501 for i in 1..150 { 502 assert_eq!(c.next(), Some(i)); 503 } 504 assert_eq!(c.next(), None); 505 for i in (0..150).rev() { 506 assert_eq!(c.prev(), Some(i)); 507 } 508 assert_eq!(c.prev(), None); 509 510 assert!(c.goto(125)); 511 for i in 125..150 { 512 assert_eq!(c.remove(), Some(i)); 513 assert!(!c.is_empty()); 514 c.verify(); 515 } 516 517 for i in (0..125).rev() { 518 assert!(!c.is_empty()); 519 assert_eq!(c.elem(), None); 520 assert_eq!(c.prev(), Some(i)); 521 assert_eq!(c.remove(), Some(i)); 522 c.verify(); 523 } 524 assert_eq!(c.elem(), None); 525 assert!(c.is_empty()); 526 } 527 528 // Generate a densely populated 4-level tree. 529 // 530 // Level 1: 1 root 531 // Level 2: 8 inner 532 // Level 3: 64 inner 533 // Level 4: 512 leafs, up to 7680 elements 534 // 535 // A 3-level tree can hold at most 960 elements. dense4l(f: &mut SetForest<i32>) -> Set<i32>536 fn dense4l(f: &mut SetForest<i32>) -> Set<i32> { 537 f.clear(); 538 let mut s = Set::new(); 539 540 // Insert 400 elements in 7 passes over the range to avoid the half-full leaf node pattern 541 // that comes from sequential insertion. This will generate a normal leaf layer. 542 for n in 0..4000 { 543 assert!(s.insert((n * 7) % 4000, f, &())); 544 } 545 s 546 } 547 548 #[test] four_level()549 fn four_level() { 550 let mut f = SetForest::<i32>::new(); 551 let mut s = dense4l(&mut f); 552 553 assert_eq!( 554 s.iter(&f).collect::<Vec<_>>()[0..10], 555 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 556 ); 557 558 let mut c = s.cursor(&mut f, &()); 559 560 c.verify(); 561 562 // Peel off a whole sub-tree of the root by deleting from the front. 563 // The 900 element is near the front of the second sub-tree. 564 assert!(c.goto(900)); 565 assert_eq!(c.tpath(), "node48[1]--node47[0]--node26[0]--node20[4]"); 566 assert!(c.goto(0)); 567 for i in 0..900 { 568 assert!(!c.is_empty()); 569 assert_eq!(c.remove(), Some(i)); 570 } 571 c.verify(); 572 assert_eq!(c.elem(), Some(900)); 573 574 // Delete backwards from somewhere in the middle. 575 assert!(c.goto(3000)); 576 for i in (2000..3000).rev() { 577 assert_eq!(c.prev(), Some(i)); 578 assert_eq!(c.remove(), Some(i)); 579 assert_eq!(c.elem(), Some(3000)); 580 } 581 c.verify(); 582 583 // Remove everything in a scattered manner, triggering many collapsing patterns. 584 for i in 0..4000 { 585 if c.goto((i * 7) % 4000) { 586 c.remove(); 587 } 588 } 589 assert!(c.is_empty()); 590 } 591 592 #[test] four_level_clear()593 fn four_level_clear() { 594 let mut f = SetForest::<i32>::new(); 595 let mut s = dense4l(&mut f); 596 s.clear(&mut f); 597 } 598 } 599