treevec.rs
1 use arrayvec::ArrayVec; 2 3 #[cfg(test)] 4 use serde::ser::SerializeSeq; 5 6 /// A container with fast random access and random insertion/deletion. 7 /// 8 /// The implementation is based on a short, wide, more-or-less-balanced tree. 9 /// You get to choose how wide it is: `B` is the maximum size of each node, 10 /// and we ensure that every node but the root has at least `B/2` elements. 11 /// The depth of the tree is therefore of order `log_B (n)`, where `n` 12 /// is the number of elements. 13 #[derive(Clone, Debug)] 14 pub struct TreeVec<T, const B: usize> { 15 root: Box<Node<T, B>>, 16 } 17 18 #[cfg(test)] 19 impl<T: serde::Serialize, const B: usize> serde::Serialize for TreeVec<T, B> { 20 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> 21 where 22 S: serde::Serializer, 23 { 24 let mut seq = serializer.serialize_seq(Some(self.len()))?; 25 for x in self.iter() { 26 seq.serialize_element(x)?; 27 } 28 seq.end() 29 } 30 } 31 32 /// A single node in the tree. 33 #[cfg_attr(test, derive(serde::Serialize))] 34 #[derive(Clone, Debug)] 35 enum Node<T, const B: usize> { 36 Leaf { 37 /// Guaranteed to have at least `B/2` elements (and in particular, 38 /// to be non-empty) *except* if this node happens to be the root 39 /// of the tree. 40 data: ArrayVec<T, B>, 41 }, 42 /// An internal node. The two arrays, `size` and `children`, are guaranteed 43 /// to be non-empty. And unless this node is the root of the tree, 44 /// they will have at least `B/2` elements. 45 Internal { 46 /// For each child, the number of elements in its subtree. 47 /// 48 /// (And so the number of elements in my subtree is the sum of these.) 49 size: ArrayVec<usize, B>, 50 /// The children. Each subtree is guaranteed to have the same height. 51 children: ArrayVec<Box<Node<T, B>>, B>, 52 }, 53 } 54 55 /// The result of inserting an element into a subtree. 56 enum InsertResult<T, const B: usize> { 57 /// It fits! 58 Done, 59 /// It didn't fit, and the subtree had to be split into two subtrees. 60 /// The node in here is the second of the two subtrees. 61 Split(Box<Node<T, B>>), 62 } 63 64 /// The result of removing an element from a subtree. 65 enum RemoveResult<T> { 66 /// It's gone! 67 Done(T), 68 /// It's gone, but the root of the subtree now has too few children. 69 Undersize(T), 70 } 71 72 impl<T> RemoveResult<T> { 73 fn into_inner(self) -> T { 74 match self { 75 RemoveResult::Done(ret) | RemoveResult::Undersize(ret) => ret, 76 } 77 } 78 } 79 80 /// The result of fixing up the sizes of two sibling subtrees. 81 enum MergeResult { 82 /// The two subtrees both fit in a single node: one has been drained and added 83 /// to the other. 84 Absorbed, 85 /// Some elements were moved from one subtree to another, and now they both 86 /// have valid sizes. 87 Rebalanced, 88 } 89 90 /// Which subtree does the given offset belong to? 91 /// 92 /// Returns the index of the subtree containing that offset, and the offset 93 /// relative to the subtree. 94 fn child_idx(sizes: &[usize], mut offset: usize) -> Option<(usize, usize)> { 95 for (idx, &size) in sizes.iter().enumerate() { 96 if size > offset { 97 return Some((idx, offset)); 98 } 99 offset -= size; 100 } 101 None 102 } 103 104 impl<T, const B: usize> Node<T, B> { 105 /// Returns the size of the subtree rooted at this node. 106 fn subtree_size(&self) -> usize { 107 match self { 108 Node::Leaf { data } => data.len(), 109 Node::Internal { size, .. } => size.iter().copied().sum(), 110 } 111 } 112 113 /// Gets the item at offset `offset`, relative to the subtree rooted at this node. 114 fn get(&self, offset: usize) -> Option<&T> { 115 match self { 116 Node::Leaf { data } => data.get(offset), 117 Node::Internal { size, children } => { 118 let (idx, offset) = child_idx(size, offset)?; 119 children[idx].get(offset) 120 } 121 } 122 } 123 124 /// Gets (mutably) the item at offset `offset`, relative to the subtree rooted at this node. 125 fn get_mut(&mut self, offset: usize) -> Option<&mut T> { 126 match self { 127 Node::Leaf { data } => data.get_mut(offset), 128 Node::Internal { size, children } => { 129 let (idx, offset) = child_idx(size, offset)?; 130 children[idx].get_mut(offset) 131 } 132 } 133 } 134 135 /// Returns the first item in the subtree. 136 /// 137 /// `tree.first()` is equivalent to `tree.get(0)`, but faster. 138 fn first(&self) -> Option<&T> { 139 match self { 140 Node::Leaf { data } => data.first(), 141 Node::Internal { children, .. } => children.first().and_then(|c| c.first()), 142 } 143 } 144 145 /// Inserts an element into the subtree rooted at this node. 146 /// 147 /// The offset is relative to this subtree. 148 /// 149 /// If this insertion causes the subtree to overflow, the root node will be split. 150 /// In this case, this node will be modified, and the new sibling will be returned. 151 fn insert(&mut self, offset: usize, element: T) -> InsertResult<T, B> { 152 match self { 153 Node::Leaf { data } => { 154 if data.is_full() { 155 let mut second_half: ArrayVec<T, B> = data.drain(B / 2..).collect(); 156 if offset <= B / 2 { 157 data.insert(offset, element) 158 } else { 159 second_half.insert(offset - B / 2, element) 160 } 161 InsertResult::Split(Box::new(Node::Leaf { data: second_half })) 162 } else { 163 data.insert(offset, element); 164 InsertResult::Done 165 } 166 } 167 Node::Internal { size, children } => { 168 let (idx, offset) = if offset > 0 { 169 // unwrap: if this fails, it's out-of-bounds 170 let (idx, offset) = child_idx(size, offset - 1).unwrap(); 171 (idx, offset + 1) 172 } else { 173 (0, 0) 174 }; 175 match children[idx].insert(offset, element) { 176 InsertResult::Done => { 177 size[idx] += 1; 178 InsertResult::Done 179 } 180 InsertResult::Split(node) => { 181 size[idx] = children[idx].subtree_size(); 182 183 if children.is_full() { 184 let mut second_half_children: ArrayVec<_, B> = 185 children.drain(B / 2..).collect(); 186 let mut second_half_size: ArrayVec<_, B> = 187 size.drain(B / 2..).collect(); 188 if idx < B / 2 { 189 size.insert(idx + 1, node.subtree_size()); 190 children.insert(idx + 1, node); 191 } else { 192 second_half_size.insert(idx + 1 - B / 2, node.subtree_size()); 193 second_half_children.insert(idx + 1 - B / 2, node); 194 } 195 InsertResult::Split(Box::new(Node::Internal { 196 size: second_half_size, 197 children: second_half_children, 198 })) 199 } else { 200 size.insert(idx + 1, node.subtree_size()); 201 children.insert(idx + 1, node); 202 InsertResult::Done 203 } 204 } 205 } 206 } 207 } 208 } 209 210 /// We're undersized (by one; we only support removing a single element at a time) 211 /// and maybe our right_sibling can help by giving us some more elements. 212 fn merge_from_right(&mut self, right_sibling: &mut Node<T, B>) -> MergeResult { 213 match (self, right_sibling) { 214 (Node::Leaf { data: left_data }, Node::Leaf { data: right_data }) => { 215 debug_assert!(right_data.len() >= left_data.len()); 216 if left_data.len() + right_data.len() <= B { 217 left_data.extend(right_data.drain(..)); 218 MergeResult::Absorbed 219 } else { 220 // Since we're only undersized by one, we could just move a single 221 // element over. It's probably more efficient to rebalance more 222 // aggressively, so we can avoid rebalancing in the future. So 223 // here we try to equalize the sizes. 224 let count = (right_data.len() - left_data.len()) / 2; 225 // We're undersized, and since the two of us don't fit in a single 226 // node, right_data must have at least B/2 + 1 elements. Therefore 227 // the difference in sizes is at least 2. 228 debug_assert!(count > 0); 229 left_data.extend(right_data.drain(..count)); 230 MergeResult::Rebalanced 231 } 232 } 233 ( 234 Node::Internal { 235 size: left_size, 236 children: left_children, 237 }, 238 Node::Internal { 239 size: right_size, 240 children: right_children, 241 }, 242 ) => { 243 if left_children.len() + right_children.len() <= B { 244 left_size.extend(right_size.drain(..)); 245 left_children.extend(right_children.drain(..)); 246 MergeResult::Absorbed 247 } else { 248 let count = (right_children.len() - left_children.len()) / 2; 249 debug_assert!(count > 0); 250 left_children.extend(right_children.drain(..count)); 251 left_size.extend(right_size.drain(..count)); 252 MergeResult::Rebalanced 253 } 254 } 255 _ => unreachable!(), 256 } 257 } 258 259 /// We're undersized (by one; we only support removing a single element at a time) 260 /// and maybe our right_sibling can help by giving us some more elements. 261 fn merge_from_left(&mut self, left_sibling: &mut Node<T, B>) -> MergeResult { 262 match (left_sibling, self) { 263 (Node::Leaf { data: left_data }, Node::Leaf { data: right_data }) => { 264 debug_assert!(right_data.len() <= left_data.len()); 265 if left_data.len() + right_data.len() <= B { 266 // There's no efficient and safe way to insert a bunch of 267 // data at the beginning of right, so instead we append to 268 // left and then swap them. 269 left_data.extend(right_data.drain(..)); 270 std::mem::swap(left_data, right_data); 271 MergeResult::Absorbed 272 } else { 273 // Unlike merge_from_right, here we only move a single element 274 // from the left to the right. This is just because safe rust 275 // makes it tricky to efficiently move more; ideally we'd also 276 // be rebalancing here. 277 right_data.insert(0, left_data.pop().unwrap()); 278 MergeResult::Rebalanced 279 } 280 } 281 ( 282 Node::Internal { 283 size: left_size, 284 children: left_children, 285 }, 286 Node::Internal { 287 size: right_size, 288 children: right_children, 289 }, 290 ) => { 291 if left_children.len() + right_children.len() <= B { 292 left_size.extend(right_size.drain(..)); 293 left_children.extend(right_children.drain(..)); 294 std::mem::swap(left_children, right_children); 295 std::mem::swap(left_size, right_size); 296 MergeResult::Absorbed 297 } else { 298 right_children.insert(0, left_children.pop().unwrap()); 299 right_size.insert(0, left_size.pop().unwrap()); 300 MergeResult::Rebalanced 301 } 302 } 303 _ => unreachable!(), 304 } 305 } 306 307 /// Remove an element from a given offset in this subtree. 308 /// 309 /// This may leave the root of the subtree undersized, in which case the 310 /// return value will let you know. 311 fn remove(&mut self, offset: usize) -> RemoveResult<T> { 312 match self { 313 Node::Leaf { data } => { 314 let ret = data.remove(offset); 315 if data.len() < B / 2 { 316 RemoveResult::Undersize(ret) 317 } else { 318 RemoveResult::Done(ret) 319 } 320 } 321 Node::Internal { size, children } => { 322 let (idx, offset) = child_idx(size, offset).unwrap(); 323 size[idx] -= 1; 324 match children[idx].remove(offset) { 325 RemoveResult::Done(ret) => RemoveResult::Done(ret), 326 RemoveResult::Undersize(ret) => { 327 // We now have an undersized node at index i. We try 328 // to merge in more from the right neighbor (because 329 // that's the more efficient merge direction in our 330 // implementation). But if `i` is the last element, we 331 // fall back to merging from the left. 332 if idx + 1 < children.len() { 333 // This little incantation seems to be the easiest 334 // (stable, safe) way to get two &muts to different 335 // elements. 336 let (a, b) = children.split_at_mut(idx + 1); 337 let cur = a.last_mut().unwrap(); 338 let next = b.first_mut().unwrap(); 339 340 match cur.merge_from_right(next) { 341 MergeResult::Absorbed => { 342 size[idx] = cur.subtree_size(); 343 344 children.remove(idx + 1); 345 size.remove(idx + 1); 346 347 if children.len() < B / 2 { 348 RemoveResult::Undersize(ret) 349 } else { 350 RemoveResult::Done(ret) 351 } 352 } 353 MergeResult::Rebalanced => { 354 size[idx] = cur.subtree_size(); 355 size[idx + 1] = next.subtree_size(); 356 RemoveResult::Done(ret) 357 } 358 } 359 } else { 360 // We couldn't merge from the right, since `idx` 361 // is at the end. Since internal nodes can't be 362 // undersized, as long as B >= 4 we must have at 363 // least 2 elements and so there is something to 364 // our left. (Maybe we should have something to 365 // ensure that B >= 4?) 366 debug_assert!(idx > 0); 367 368 let (a, b) = children.split_at_mut(idx); 369 let prev = a.last_mut().unwrap(); 370 let cur = b.first_mut().unwrap(); 371 372 match cur.merge_from_left(prev) { 373 MergeResult::Absorbed => { 374 size[idx] = cur.subtree_size(); 375 376 children.remove(idx - 1); 377 size.remove(idx - 1); 378 if children.len() < B / 2 { 379 RemoveResult::Undersize(ret) 380 } else { 381 RemoveResult::Done(ret) 382 } 383 } 384 MergeResult::Rebalanced => { 385 size[idx - 1] = prev.subtree_size(); 386 size[idx] = cur.subtree_size(); 387 RemoveResult::Done(ret) 388 } 389 } 390 } 391 } 392 } 393 } 394 } 395 } 396 397 /// Returns the offset (in this subtree) of the first element for which `pred` returns false. 398 fn partition_point<P>(&self, mut pred: P) -> usize 399 where 400 P: FnMut(&T) -> bool, 401 { 402 match self { 403 Node::Leaf { data } => data.partition_point(pred), 404 Node::Internal { children, size } => { 405 // It seems to be faster if we first do a binary search over our subtrees, 406 // and then we descend into the interesting subtree. 407 let child_idx = children.partition_point(|child| pred(child.first().unwrap())); 408 if child_idx > 0 { 409 children[child_idx - 1].partition_point(pred) 410 + size[..(child_idx - 1)].iter().copied().sum::<usize>() 411 } else { 412 0 413 } 414 } 415 } 416 } 417 418 /// Assert that this subtree satisfies our invariants, panicking if not. 419 fn check_invariants(&self, is_root: bool) { 420 match self { 421 Node::Leaf { data } => { 422 if !is_root { 423 assert!(data.len() >= B / 2); 424 } 425 } 426 Node::Internal { size, children } => { 427 assert_eq!(size.len(), children.len()); 428 if !is_root { 429 assert!(size.len() >= B / 2); 430 } 431 432 for (child, size) in children.iter().zip(size) { 433 assert_eq!(child.subtree_size(), *size); 434 435 child.check_invariants(false); 436 } 437 } 438 } 439 } 440 } 441 442 impl<T, const B: usize> Default for TreeVec<T, B> { 443 fn default() -> Self { 444 Self { 445 root: Box::new(Node::Leaf { 446 data: ArrayVec::new(), 447 }), 448 } 449 } 450 } 451 452 impl<T, const B: usize> TreeVec<T, B> { 453 /// Constructs a new, empty, `TreeVec`. 454 /// 455 /// This allocates a little; we don't have a special-case non-allocating 456 /// `TreeVec` for empty containers. 457 pub fn new() -> Self { 458 Self::default() 459 } 460 461 /// Returns true if we're empty. 462 pub fn is_empty(&self) -> bool { 463 match &*self.root { 464 Node::Leaf { data } => data.is_empty(), 465 Node::Internal { .. } => false, 466 } 467 } 468 469 /// Returns the number of elements in this container. 470 pub fn len(&self) -> usize { 471 self.root.subtree_size() 472 } 473 474 /// Gets a reference to the element at a specified index, or `None` if the 475 /// index is out-of-bounds. 476 /// 477 /// Runtime is linear in the height of the tree (logarithmic in the length). 478 pub fn get(&self, index: usize) -> Option<&T> { 479 self.root.get(index) 480 } 481 482 /// Gets a mutable reference to the element at a specified index, or `None` 483 /// if the index is out-of-bounds. 484 /// 485 /// Runtime is linear in the height of the tree (logarithmic in the length). 486 pub fn get_mut(&mut self, index: usize) -> Option<&mut T> { 487 self.root.get_mut(index) 488 } 489 490 /// Inserts an element at a specified index, or panics if the index is 491 /// out-of-bounds. 492 /// 493 /// Runtime is linear in the height of the tree (logarithmic in the length), 494 /// and also linear in `B`. 495 pub fn insert(&mut self, index: usize, element: T) { 496 match self.root.insert(index, element) { 497 InsertResult::Done => {} 498 InsertResult::Split(node) => { 499 let mut root = Box::new(Node::Internal { 500 size: ArrayVec::new(), 501 children: ArrayVec::new(), 502 }); 503 std::mem::swap(&mut root, &mut self.root); 504 505 let Node::Internal { size, children } = &mut *self.root else { 506 unreachable!(); 507 }; 508 size.push(root.subtree_size()); 509 size.push(node.subtree_size()); 510 children.push(root); 511 children.push(node); 512 } 513 } 514 } 515 516 /// Inserts an element at a specified index, shifting all later elements 517 /// back by 1. 518 /// 519 /// Runtime is linear in the height of the tree (logarithmic in the length), 520 /// and also linear in `B`. 521 pub fn remove(&mut self, index: usize) -> T { 522 let ret = self.root.remove(index).into_inner(); 523 524 if let Node::Internal { children, .. } = &mut *self.root { 525 if children.len() == 1 { 526 // unwrap: an internal node always has children 527 self.root = children.pop().unwrap() 528 } 529 } 530 ret 531 } 532 533 /// Asserts that this subtree satisfies our invariants, panicking if not. 534 pub fn check_invariants(&self) { 535 self.root.check_invariants(true); 536 } 537 538 /// Returns an iterator over all the elements in this container. 539 pub fn iter(&self) -> Iter<'_, T, B> { 540 let inner = match &*self.root { 541 Node::Leaf { data } => IterInner::Simple(data.iter()), 542 Node::Internal { .. } => { 543 let mut ret = TreeIter { 544 stack: Vec::new(), 545 leaf: [].iter(), 546 remaining: self.len(), 547 }; 548 ret.descend(&*self.root); 549 IterInner::Tree(ret) 550 } 551 }; 552 Iter { inner } 553 } 554 555 /// Returns an iterator over mutable references to all the elements in this 556 /// container. 557 pub fn iter_mut(&mut self) -> IterMut<'_, T, B> { 558 let len = self.len(); 559 let inner = match &mut *self.root { 560 Node::Leaf { data } => IterMutInner::Simple(data.iter_mut()), 561 root @ Node::Internal { .. } => { 562 let mut ret = TreeIterMut { 563 stack: Vec::new(), 564 leaf: [].iter_mut(), 565 remaining: len, 566 }; 567 ret.descend(root); 568 IterMutInner::Tree(ret) 569 } 570 }; 571 IterMut { inner } 572 } 573 574 /// Returns the index of the first element for which `pred` returns false. 575 /// 576 /// The above assumes that we're sorted in the sense that `pred` returns `true` for the first 577 /// bunch of elements and then `false` for the rest. If not, we don't necessarily return 578 /// the first offset for which `pred` if `false`, but we do return some offset where 579 /// `pred` is `false` but `pred` is `true` at `offset - 1`. (Here, corner cases are specified 580 /// by declaring that `pred` is `true` at offset `-1` and `false` at offset `len`. 581 pub fn partition_point<P>(&self, pred: P) -> usize 582 where 583 P: FnMut(&T) -> bool, 584 { 585 self.root.partition_point(pred) 586 } 587 588 /// Turns a generic range bounds into an inclusive start and an exclusive end. 589 /// 590 /// This won't work if the range end is inclusive of `usize::MAX`, so don't do that. 591 fn normalize_bounds(&self, range: impl std::ops::RangeBounds<usize>) -> (usize, usize) { 592 let start = match range.start_bound() { 593 std::ops::Bound::Included(x) => *x, 594 std::ops::Bound::Excluded(x) => *x + 1, 595 std::ops::Bound::Unbounded => 0, 596 }; 597 let end = match range.end_bound() { 598 std::ops::Bound::Included(x) => *x + 1, 599 std::ops::Bound::Excluded(x) => *x, 600 std::ops::Bound::Unbounded => self.len(), 601 }; 602 (start, end) 603 } 604 605 /// Returns an iterator over a sub-range of this container. 606 pub fn range(&self, range: impl std::ops::RangeBounds<usize>) -> Iter<'_, T, B> { 607 let (start, end) = self.normalize_bounds(range); 608 let inner = match &*self.root { 609 Node::Leaf { data } => IterInner::Simple(data[start..end].iter()), 610 Node::Internal { .. } => { 611 if end > self.len() { 612 panic!("out of bounds"); 613 } 614 let mut ret = TreeIter { 615 stack: Vec::new(), 616 leaf: [].iter(), 617 remaining: end - start, 618 }; 619 ret.descend_to(&*self.root, start); 620 IterInner::Tree(ret) 621 } 622 }; 623 Iter { inner } 624 } 625 626 /// Returns an iterator over mutable references to a sub-range of this container. 627 pub fn range_mut(&mut self, range: impl std::ops::RangeBounds<usize>) -> IterMut<'_, T, B> { 628 let (start, end) = self.normalize_bounds(range); 629 let len = self.len(); 630 631 let inner = match &mut *self.root { 632 Node::Leaf { data } => IterMutInner::Simple(data[start..end].iter_mut()), 633 root @ Node::Internal { .. } => { 634 if end > len { 635 panic!("out of bounds"); 636 } 637 let mut ret = TreeIterMut { 638 stack: Vec::new(), 639 leaf: [].iter_mut(), 640 remaining: end - start, 641 }; 642 ret.descend_to(root, start); 643 IterMutInner::Tree(ret) 644 } 645 }; 646 IterMut { inner } 647 } 648 } 649 650 impl<T, const B: usize> FromIterator<T> for TreeVec<T, B> { 651 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self { 652 // TODO: a faster implementation, maybe? This is only 653 // used in tests right now. 654 let mut ret = TreeVec::new(); 655 for (idx, x) in iter.into_iter().enumerate() { 656 ret.insert(idx, x); 657 } 658 ret 659 } 660 } 661 662 impl<T, const B: usize> std::ops::Index<usize> for TreeVec<T, B> { 663 type Output = T; 664 665 fn index(&self, index: usize) -> &T { 666 self.get(index).unwrap() 667 } 668 } 669 670 impl<T, const B: usize> std::ops::IndexMut<usize> for TreeVec<T, B> { 671 fn index_mut(&mut self, index: usize) -> &mut T { 672 self.get_mut(index).unwrap() 673 } 674 } 675 676 /// An iterator over elements of a [`TreeVec`]. 677 pub struct Iter<'a, T, const B: usize> { 678 inner: IterInner<'a, T, B>, 679 } 680 681 // We special-case a fast-path for a tree with height 1. 682 // 683 // This, along with some inline annotations, makes a measurable performance 684 // difference. 685 enum IterInner<'a, T, const B: usize> { 686 Simple(std::slice::Iter<'a, T>), 687 Tree(TreeIter<'a, T, B>), 688 } 689 690 struct TreeIter<'a, T, const B: usize> { 691 // If you imagine how the recursive implementation of a tree iteration 692 // works, the local state at each intermediate call consists of a 693 // partially-finished iteration over the children of that internal node. 694 // This stores the stack of that imaginary recursive implementation. 695 stack: Vec<std::slice::Iter<'a, Box<Node<T, B>>>>, 696 leaf: std::slice::Iter<'a, T>, 697 // The number of remaining elements; used for iteration over ranges. The 698 // full-iteration implementation doesn't need it, but it isn't worth having 699 // two different iterator implementations. 700 remaining: usize, 701 } 702 703 impl<'a, T, const B: usize> TreeIter<'a, T, B> { 704 // Fill out the rest of the stack and the leaf, by descending 705 // to the first descendent of `node`. 706 fn descend(&mut self, mut node: &'a Node<T, B>) { 707 loop { 708 match node { 709 Node::Leaf { data } => { 710 self.leaf = data.iter(); 711 return; 712 } 713 Node::Internal { children, .. } => { 714 let mut children = children.iter(); 715 // unwrap: internal nodes are always non-empty 716 node = children.next().unwrap(); 717 self.stack.push(children); 718 } 719 } 720 } 721 } 722 723 // Fill out the rest of the stack and the leaf, by descending 724 // to the node at offset `offset` in the subtree rooted at `node`. 725 fn descend_to(&mut self, mut node: &'a Node<T, B>, mut offset: usize) { 726 loop { 727 match node { 728 Node::Leaf { data } => { 729 self.leaf = data[offset..].iter(); 730 return; 731 } 732 Node::Internal { children, size } => { 733 let Some((idx, child_offset)) = child_idx(size, offset) else { 734 return; 735 }; 736 offset = child_offset; 737 let mut children = children[idx..].iter(); 738 // unwrap: child_idx always returns a valid index into children 739 node = children.next().unwrap(); 740 self.stack.push(children); 741 } 742 } 743 } 744 } 745 746 fn next(&mut self) -> Option<&'a T> { 747 if self.remaining == 0 { 748 None 749 } else { 750 self.remaining -= 1; 751 if let Some(ret) = self.leaf.next() { 752 Some(ret) 753 } else { 754 // The current leaf is exhausted. Unwind the stack until we find 755 // a non-exhausted internal node, then descend again. 756 loop { 757 let stack_top = self.stack.last_mut()?; 758 759 let Some(next_node) = stack_top.next() else { 760 self.stack.pop(); 761 continue; 762 }; 763 764 self.descend(next_node); 765 return self.leaf.next(); 766 } 767 } 768 } 769 } 770 } 771 772 impl<'a, T, const B: usize> Iterator for Iter<'a, T, B> { 773 type Item = &'a T; 774 775 #[inline(always)] 776 fn next(&mut self) -> Option<Self::Item> { 777 match &mut self.inner { 778 IterInner::Simple(it) => it.next(), 779 IterInner::Tree(it) => it.next(), 780 } 781 } 782 } 783 784 /// An iterator over mutable references elements of a [`TreeVec`]. 785 pub struct IterMut<'a, T, const B: usize> { 786 inner: IterMutInner<'a, T, B>, 787 } 788 789 enum IterMutInner<'a, T, const B: usize> { 790 Simple(std::slice::IterMut<'a, T>), 791 Tree(TreeIterMut<'a, T, B>), 792 } 793 794 // This is basically a copy-pasted of `TreeIter`, but with non-mut things 795 // turned into mut things. Right now there's only the one duplication so 796 // is isn't worth trying to build an abstraction. 797 struct TreeIterMut<'a, T, const B: usize> { 798 stack: Vec<std::slice::IterMut<'a, Box<Node<T, B>>>>, 799 leaf: std::slice::IterMut<'a, T>, 800 remaining: usize, 801 } 802 803 impl<'a, T, const B: usize> TreeIterMut<'a, T, B> { 804 fn descend(&mut self, mut node: &'a mut Node<T, B>) { 805 loop { 806 match node { 807 Node::Leaf { data } => { 808 self.leaf = data.iter_mut(); 809 return; 810 } 811 Node::Internal { children, .. } => { 812 let mut children = children.iter_mut(); 813 // unwrap: internal nodes are always non-empty 814 node = children.next().unwrap(); 815 self.stack.push(children); 816 } 817 } 818 } 819 } 820 821 fn descend_to(&mut self, mut node: &'a mut Node<T, B>, mut offset: usize) { 822 loop { 823 match node { 824 Node::Leaf { data } => { 825 self.leaf = data[offset..].iter_mut(); 826 return; 827 } 828 Node::Internal { children, size } => { 829 let Some((idx, child_offset)) = child_idx(size, offset) else { 830 return; 831 }; 832 offset = child_offset; 833 let mut children = children[idx..].iter_mut(); 834 // unwrap: child_idx always returns a valid index into children 835 node = children.next().unwrap(); 836 self.stack.push(children); 837 } 838 } 839 } 840 } 841 842 fn next(&mut self) -> Option<&'a mut T> { 843 if self.remaining == 0 { 844 None 845 } else { 846 self.remaining -= 1; 847 if let Some(ret) = self.leaf.next() { 848 Some(ret) 849 } else { 850 loop { 851 let stack_top = self.stack.last_mut()?; 852 853 let Some(next_node) = stack_top.next() else { 854 self.stack.pop(); 855 continue; 856 }; 857 858 self.descend(next_node); 859 return self.leaf.next(); 860 } 861 } 862 } 863 } 864 } 865 866 impl<'a, T, const B: usize> Iterator for IterMut<'a, T, B> { 867 type Item = &'a mut T; 868 869 #[inline(always)] 870 fn next(&mut self) -> Option<Self::Item> { 871 match &mut self.inner { 872 IterMutInner::Simple(iter) => iter.next(), 873 IterMutInner::Tree(tree_iter) => tree_iter.next(), 874 } 875 } 876 } 877 878 #[cfg(test)] 879 mod tests { 880 use super::*; 881 882 #[test] 883 fn insert_get() { 884 let mut vec = TreeVec::<i32, 4>::default(); 885 vec.insert(0, 1); 886 vec.insert(0, 2); 887 vec.insert(0, 3); 888 vec.insert(0, 4); 889 assert_eq!(*vec.get(0).unwrap(), 4); 890 assert_eq!(*vec.get(1).unwrap(), 3); 891 assert_eq!(*vec.get(2).unwrap(), 2); 892 assert_eq!(*vec.get(3).unwrap(), 1); 893 vec.check_invariants(); 894 895 vec.insert(0, 1); 896 vec.insert(0, 2); 897 vec.insert(0, 3); 898 vec.insert(0, 4); 899 vec.check_invariants(); 900 assert_eq!(*vec.get(0).unwrap(), 4); 901 assert_eq!(*vec.get(1).unwrap(), 3); 902 assert_eq!(*vec.get(2).unwrap(), 2); 903 assert_eq!(*vec.get(3).unwrap(), 1); 904 assert_eq!(*vec.get(4).unwrap(), 4); 905 assert_eq!(*vec.get(5).unwrap(), 3); 906 assert_eq!(*vec.get(6).unwrap(), 2); 907 assert_eq!(*vec.get(7).unwrap(), 1); 908 } 909 910 #[test] 911 fn insert_remove() { 912 let mut vec = TreeVec::<i32, 4>::default(); 913 vec.insert(0, 1); 914 vec.insert(0, 2); 915 vec.insert(0, 3); 916 vec.insert(0, 4); 917 vec.remove(1); 918 vec.check_invariants(); 919 assert_eq!(*vec.get(0).unwrap(), 4); 920 assert_eq!(*vec.get(1).unwrap(), 2); 921 assert_eq!(*vec.get(2).unwrap(), 1); 922 923 vec.insert(0, 1); 924 vec.insert(0, 2); 925 vec.insert(0, 3); 926 vec.insert(0, 4); 927 vec.remove(5); 928 vec.check_invariants(); 929 assert_eq!(*vec.get(0).unwrap(), 4); 930 assert_eq!(*vec.get(1).unwrap(), 3); 931 assert_eq!(*vec.get(2).unwrap(), 2); 932 assert_eq!(*vec.get(3).unwrap(), 1); 933 assert_eq!(*vec.get(4).unwrap(), 4); 934 assert_eq!(*vec.get(5).unwrap(), 1); 935 } 936 937 #[test] 938 fn iter() { 939 let mut vec = TreeVec::<i32, 4>::default(); 940 assert_eq!(vec.iter().cloned().collect::<Vec<_>>(), Vec::<i32>::new()); 941 vec.insert(0, 1); 942 assert_eq!(vec.iter().cloned().collect::<Vec<_>>(), vec![1]); 943 vec.insert(1, 2); 944 assert_eq!(vec.iter().cloned().collect::<Vec<_>>(), vec![1, 2]); 945 vec.insert(0, 1); 946 vec.insert(0, 1); 947 vec.insert(0, 1); 948 vec.insert(0, 1); 949 vec.insert(0, 1); 950 assert_eq!( 951 vec.iter().cloned().collect::<Vec<_>>(), 952 vec![1, 1, 1, 1, 1, 1, 2] 953 ); 954 } 955 }