1use super::*;
2use arrayvec::ArrayVec;
3use std::{cmp::Ordering, mem, sync::Arc};
4use ztracing::instrument;
5
6#[derive(Clone)]
7struct StackEntry<'a, T: Item, D> {
8 tree: &'a SumTree<T>,
9 index: u32,
10 position: D,
11}
12
13impl<'a, T: Item, D> StackEntry<'a, T, D> {
14 #[inline]
15 fn index(&self) -> usize {
16 self.index as usize
17 }
18}
19
20impl<T: Item + fmt::Debug, D: fmt::Debug> fmt::Debug for StackEntry<'_, T, D> {
21 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22 f.debug_struct("StackEntry")
23 .field("index", &self.index)
24 .field("position", &self.position)
25 .finish()
26 }
27}
28
29#[derive(Clone)]
30pub struct Cursor<'a, 'b, T: Item, D> {
31 tree: &'a SumTree<T>,
32 stack: ArrayVec<StackEntry<'a, T, D>, 16>,
33 pub position: D,
34 did_seek: bool,
35 at_end: bool,
36 cx: <T::Summary as Summary>::Context<'b>,
37}
38
39impl<T: Item + fmt::Debug, D: fmt::Debug> fmt::Debug for Cursor<'_, '_, T, D>
40where
41 T::Summary: fmt::Debug,
42{
43 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44 f.debug_struct("Cursor")
45 .field("tree", &self.tree)
46 .field("stack", &self.stack)
47 .field("position", &self.position)
48 .field("did_seek", &self.did_seek)
49 .field("at_end", &self.at_end)
50 .finish()
51 }
52}
53
54pub struct Iter<'a, T: Item> {
55 tree: &'a SumTree<T>,
56 stack: ArrayVec<StackEntry<'a, T, ()>, 16>,
57}
58
59impl<'a, 'b, T, D> Cursor<'a, 'b, T, D>
60where
61 T: Item,
62 D: Dimension<'a, T::Summary>,
63{
64 pub fn new(tree: &'a SumTree<T>, cx: <T::Summary as Summary>::Context<'b>) -> Self {
65 Self {
66 tree,
67 stack: ArrayVec::new(),
68 position: D::zero(cx),
69 did_seek: false,
70 at_end: tree.is_empty(),
71 cx,
72 }
73 }
74
75 pub fn reset(&mut self) {
76 self.did_seek = false;
77 self.at_end = self.tree.is_empty();
78 self.stack.truncate(0);
79 self.position = D::zero(self.cx);
80 }
81
82 pub fn start(&self) -> &D {
83 &self.position
84 }
85
86 #[track_caller]
87 pub fn end(&self) -> D {
88 if let Some(item_summary) = self.item_summary() {
89 let mut end = self.start().clone();
90 end.add_summary(item_summary, self.cx);
91 end
92 } else {
93 self.start().clone()
94 }
95 }
96
97 /// Item is None, when the list is empty, or this cursor is at the end of the list.
98 #[track_caller]
99 pub fn item(&self) -> Option<&'a T> {
100 self.assert_did_seek();
101 if let Some(entry) = self.stack.last() {
102 match *entry.tree.0 {
103 Node::Leaf { ref items, .. } => {
104 if entry.index() == items.len() {
105 None
106 } else {
107 Some(&items[entry.index()])
108 }
109 }
110 _ => unreachable!(),
111 }
112 } else {
113 None
114 }
115 }
116
117 #[track_caller]
118 pub fn item_summary(&self) -> Option<&'a T::Summary> {
119 self.assert_did_seek();
120 if let Some(entry) = self.stack.last() {
121 match *entry.tree.0 {
122 Node::Leaf {
123 ref item_summaries, ..
124 } => {
125 if entry.index() == item_summaries.len() {
126 None
127 } else {
128 Some(&item_summaries[entry.index()])
129 }
130 }
131 _ => unreachable!(),
132 }
133 } else {
134 None
135 }
136 }
137
138 #[track_caller]
139 pub fn next_item(&self) -> Option<&'a T> {
140 self.assert_did_seek();
141 if let Some(entry) = self.stack.last() {
142 if entry.index() == entry.tree.0.items().len() - 1 {
143 if let Some(next_leaf) = self.next_leaf() {
144 Some(next_leaf.0.items().first().unwrap())
145 } else {
146 None
147 }
148 } else {
149 match *entry.tree.0 {
150 Node::Leaf { ref items, .. } => Some(&items[entry.index() + 1]),
151 _ => unreachable!(),
152 }
153 }
154 } else if self.at_end {
155 None
156 } else {
157 self.tree.first()
158 }
159 }
160
161 #[track_caller]
162 fn next_leaf(&self) -> Option<&'a SumTree<T>> {
163 for entry in self.stack.iter().rev().skip(1) {
164 if entry.index() < entry.tree.0.child_trees().len() - 1 {
165 match *entry.tree.0 {
166 Node::Internal {
167 ref child_trees, ..
168 } => return Some(child_trees[entry.index() + 1].leftmost_leaf()),
169 Node::Leaf { .. } => unreachable!(),
170 };
171 }
172 }
173 None
174 }
175
176 #[track_caller]
177 pub fn prev_item(&self) -> Option<&'a T> {
178 self.assert_did_seek();
179 if let Some(entry) = self.stack.last() {
180 if entry.index() == 0 {
181 if let Some(prev_leaf) = self.prev_leaf() {
182 Some(prev_leaf.0.items().last().unwrap())
183 } else {
184 None
185 }
186 } else {
187 match *entry.tree.0 {
188 Node::Leaf { ref items, .. } => Some(&items[entry.index() - 1]),
189 _ => unreachable!(),
190 }
191 }
192 } else if self.at_end {
193 self.tree.last()
194 } else {
195 None
196 }
197 }
198
199 #[track_caller]
200 fn prev_leaf(&self) -> Option<&'a SumTree<T>> {
201 for entry in self.stack.iter().rev().skip(1) {
202 if entry.index() != 0 {
203 match *entry.tree.0 {
204 Node::Internal {
205 ref child_trees, ..
206 } => return Some(child_trees[entry.index() - 1].rightmost_leaf()),
207 Node::Leaf { .. } => unreachable!(),
208 };
209 }
210 }
211 None
212 }
213
214 #[track_caller]
215 #[instrument(skip_all)]
216 pub fn prev(&mut self) {
217 self.search_backward(|_| true)
218 }
219
220 #[track_caller]
221 pub fn search_backward<F>(&mut self, mut filter_node: F)
222 where
223 F: FnMut(&T::Summary) -> bool,
224 {
225 if !self.did_seek {
226 self.did_seek = true;
227 self.at_end = true;
228 }
229
230 if self.at_end {
231 self.position = D::zero(self.cx);
232 self.at_end = self.tree.is_empty();
233 if !self.tree.is_empty() {
234 self.stack.push(StackEntry {
235 tree: self.tree,
236 index: self.tree.0.child_summaries().len() as u32,
237 position: D::from_summary(self.tree.summary(), self.cx),
238 });
239 }
240 }
241
242 let mut descending = false;
243 while !self.stack.is_empty() {
244 if let Some(StackEntry { position, .. }) = self.stack.iter().rev().nth(1) {
245 self.position = position.clone();
246 } else {
247 self.position = D::zero(self.cx);
248 }
249
250 let entry = self.stack.last_mut().unwrap();
251 if !descending {
252 if entry.index() == 0 {
253 self.stack.pop();
254 continue;
255 } else {
256 entry.index -= 1;
257 }
258 }
259
260 for summary in &entry.tree.0.child_summaries()[..entry.index()] {
261 self.position.add_summary(summary, self.cx);
262 }
263 entry.position = self.position.clone();
264
265 descending = filter_node(&entry.tree.0.child_summaries()[entry.index()]);
266 match entry.tree.0.as_ref() {
267 Node::Internal { child_trees, .. } => {
268 if descending {
269 let tree = &child_trees[entry.index()];
270 self.stack.push(StackEntry {
271 position: D::zero(self.cx),
272 tree,
273 index: tree.0.child_summaries().len() as u32 - 1,
274 })
275 }
276 }
277 Node::Leaf { .. } => {
278 if descending {
279 break;
280 }
281 }
282 }
283 }
284 }
285
286 #[track_caller]
287 pub fn next(&mut self) {
288 self.search_forward(|_| true)
289 }
290
291 #[track_caller]
292 pub fn search_forward<F>(&mut self, mut filter_node: F)
293 where
294 F: FnMut(&T::Summary) -> bool,
295 {
296 let mut descend = false;
297
298 if self.stack.is_empty() {
299 if !self.at_end {
300 self.stack.push(StackEntry {
301 tree: self.tree,
302 index: 0,
303 position: D::zero(self.cx),
304 });
305 descend = true;
306 }
307 self.did_seek = true;
308 }
309
310 while !self.stack.is_empty() {
311 let new_subtree = {
312 let entry = self.stack.last_mut().unwrap();
313 match entry.tree.0.as_ref() {
314 Node::Internal {
315 child_trees,
316 child_summaries,
317 ..
318 } => {
319 if !descend {
320 entry.index += 1;
321 entry.position = self.position.clone();
322 }
323
324 while entry.index() < child_summaries.len() {
325 let next_summary = &child_summaries[entry.index()];
326 if filter_node(next_summary) {
327 break;
328 } else {
329 entry.index += 1;
330 entry.position.add_summary(next_summary, self.cx);
331 self.position.add_summary(next_summary, self.cx);
332 }
333 }
334
335 child_trees.get(entry.index())
336 }
337 Node::Leaf { item_summaries, .. } => {
338 if !descend {
339 let item_summary = &item_summaries[entry.index()];
340 entry.index += 1;
341 entry.position.add_summary(item_summary, self.cx);
342 self.position.add_summary(item_summary, self.cx);
343 }
344
345 loop {
346 if let Some(next_item_summary) = item_summaries.get(entry.index()) {
347 if filter_node(next_item_summary) {
348 return;
349 } else {
350 entry.index += 1;
351 entry.position.add_summary(next_item_summary, self.cx);
352 self.position.add_summary(next_item_summary, self.cx);
353 }
354 } else {
355 break None;
356 }
357 }
358 }
359 }
360 };
361
362 if let Some(subtree) = new_subtree {
363 descend = true;
364 self.stack.push(StackEntry {
365 tree: subtree,
366 index: 0,
367 position: self.position.clone(),
368 });
369 } else {
370 descend = false;
371 self.stack.pop();
372 }
373 }
374
375 self.at_end = self.stack.is_empty();
376 debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
377 }
378
379 #[track_caller]
380 fn assert_did_seek(&self) {
381 assert!(
382 self.did_seek,
383 "Must call `seek`, `next` or `prev` before calling this method"
384 );
385 }
386
387 pub fn did_seek(&self) -> bool {
388 self.did_seek
389 }
390}
391
392impl<'a, 'b, T, D> Cursor<'a, 'b, T, D>
393where
394 T: Item,
395 D: Dimension<'a, T::Summary>,
396{
397 /// Returns whether we found the item you were seeking for.
398 #[track_caller]
399 #[instrument(skip_all)]
400 pub fn seek<Target>(&mut self, pos: &Target, bias: Bias) -> bool
401 where
402 Target: SeekTarget<'a, T::Summary, D>,
403 {
404 self.reset();
405 self.seek_internal(pos, bias, &mut ())
406 }
407
408 /// Returns whether we found the item you were seeking for.
409 ///
410 /// # Panics
411 ///
412 /// If we did not seek before, use seek instead in that case.
413 #[track_caller]
414 #[instrument(skip_all)]
415 pub fn seek_forward<Target>(&mut self, pos: &Target, bias: Bias) -> bool
416 where
417 Target: SeekTarget<'a, T::Summary, D>,
418 {
419 self.seek_internal(pos, bias, &mut ())
420 }
421
422 /// Advances the cursor and returns traversed items as a tree.
423 #[track_caller]
424 pub fn slice<Target>(&mut self, end: &Target, bias: Bias) -> SumTree<T>
425 where
426 Target: SeekTarget<'a, T::Summary, D>,
427 {
428 let mut slice = SliceSeekAggregate {
429 tree: SumTree::new(self.cx),
430 leaf_items: ArrayVec::new(),
431 leaf_item_summaries: ArrayVec::new(),
432 leaf_summary: <T::Summary as Summary>::zero(self.cx),
433 };
434 self.seek_internal(end, bias, &mut slice);
435 slice.tree
436 }
437
438 #[track_caller]
439 pub fn suffix(&mut self) -> SumTree<T> {
440 self.slice(&End::new(), Bias::Right)
441 }
442
443 #[track_caller]
444 pub fn summary<Target, Output>(&mut self, end: &Target, bias: Bias) -> Output
445 where
446 Target: SeekTarget<'a, T::Summary, D>,
447 Output: Dimension<'a, T::Summary>,
448 {
449 let mut summary = SummarySeekAggregate(Output::zero(self.cx));
450 self.seek_internal(end, bias, &mut summary);
451 summary.0
452 }
453
454 /// Returns whether we found the item you were seeking for.
455 #[track_caller]
456 #[instrument(skip_all)]
457 fn seek_internal(
458 &mut self,
459 target: &dyn SeekTarget<'a, T::Summary, D>,
460 bias: Bias,
461 aggregate: &mut dyn SeekAggregate<'a, T>,
462 ) -> bool {
463 assert!(
464 target.cmp(&self.position, self.cx).is_ge(),
465 "cannot seek backward",
466 );
467
468 if !self.did_seek {
469 self.did_seek = true;
470 self.stack.push(StackEntry {
471 tree: self.tree,
472 index: 0,
473 position: D::zero(self.cx),
474 });
475 }
476
477 let mut ascending = false;
478 'outer: while let Some(entry) = self.stack.last_mut() {
479 match *entry.tree.0 {
480 Node::Internal {
481 ref child_summaries,
482 ref child_trees,
483 ..
484 } => {
485 if ascending {
486 entry.index += 1;
487 entry.position = self.position.clone();
488 }
489
490 for (child_tree, child_summary) in child_trees[entry.index()..]
491 .iter()
492 .zip(&child_summaries[entry.index()..])
493 {
494 let mut child_end = self.position.clone();
495 child_end.add_summary(child_summary, self.cx);
496
497 let comparison = target.cmp(&child_end, self.cx);
498 if comparison == Ordering::Greater
499 || (comparison == Ordering::Equal && bias == Bias::Right)
500 {
501 self.position = child_end;
502 aggregate.push_tree(child_tree, child_summary, self.cx);
503 entry.index += 1;
504 entry.position = self.position.clone();
505 } else {
506 self.stack.push(StackEntry {
507 tree: child_tree,
508 index: 0,
509 position: self.position.clone(),
510 });
511 ascending = false;
512 continue 'outer;
513 }
514 }
515 }
516 Node::Leaf {
517 ref items,
518 ref item_summaries,
519 ..
520 } => {
521 aggregate.begin_leaf();
522
523 for (item, item_summary) in items[entry.index()..]
524 .iter()
525 .zip(&item_summaries[entry.index()..])
526 {
527 let mut child_end = self.position.clone();
528 child_end.add_summary(item_summary, self.cx);
529
530 let comparison = target.cmp(&child_end, self.cx);
531 if comparison == Ordering::Greater
532 || (comparison == Ordering::Equal && bias == Bias::Right)
533 {
534 self.position = child_end;
535 aggregate.push_item(item, item_summary, self.cx);
536 entry.index += 1;
537 } else {
538 aggregate.end_leaf(self.cx);
539 break 'outer;
540 }
541 }
542
543 aggregate.end_leaf(self.cx);
544 }
545 }
546
547 self.stack.pop();
548 ascending = true;
549 }
550
551 self.at_end = self.stack.is_empty();
552 debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
553
554 let mut end = self.position.clone();
555 if bias == Bias::Left
556 && let Some(summary) = self.item_summary()
557 {
558 end.add_summary(summary, self.cx);
559 }
560
561 target.cmp(&end, self.cx) == Ordering::Equal
562 }
563}
564
565impl<'a, T: Item> Iter<'a, T> {
566 pub(crate) fn new(tree: &'a SumTree<T>) -> Self {
567 Self {
568 tree,
569 stack: Default::default(),
570 }
571 }
572}
573
574impl<'a, T: Item> Iterator for Iter<'a, T> {
575 type Item = &'a T;
576
577 fn next(&mut self) -> Option<Self::Item> {
578 let mut descend = false;
579
580 if self.stack.is_empty() {
581 self.stack.push(StackEntry {
582 tree: self.tree,
583 index: 0,
584 position: (),
585 });
586 descend = true;
587 }
588
589 while let Some(entry) = self.stack.last_mut() {
590 let new_subtree = {
591 match entry.tree.0.as_ref() {
592 Node::Internal { child_trees, .. } => {
593 if !descend {
594 entry.index += 1;
595 }
596 child_trees.get(entry.index())
597 }
598 Node::Leaf { items, .. } => {
599 if !descend {
600 entry.index += 1;
601 }
602
603 if let Some(next_item) = items.get(entry.index()) {
604 return Some(next_item);
605 } else {
606 None
607 }
608 }
609 }
610 };
611
612 if let Some(subtree) = new_subtree {
613 descend = true;
614 self.stack.push(StackEntry {
615 tree: subtree,
616 index: 0,
617 position: (),
618 });
619 } else {
620 descend = false;
621 self.stack.pop();
622 }
623 }
624
625 None
626 }
627
628 fn last(mut self) -> Option<Self::Item> {
629 self.stack.clear();
630 self.tree.rightmost_leaf().last()
631 }
632
633 fn size_hint(&self) -> (usize, Option<usize>) {
634 let lower_bound = match self.stack.last() {
635 Some(top) => top.tree.0.child_summaries().len() - top.index as usize,
636 None => self.tree.0.child_summaries().len(),
637 };
638
639 (lower_bound, None)
640 }
641}
642
643impl<'a, 'b, T: Item, D> Iterator for Cursor<'a, 'b, T, D>
644where
645 D: Dimension<'a, T::Summary>,
646{
647 type Item = &'a T;
648
649 fn next(&mut self) -> Option<Self::Item> {
650 if !self.did_seek {
651 self.next();
652 }
653
654 if let Some(item) = self.item() {
655 self.next();
656 Some(item)
657 } else {
658 None
659 }
660 }
661}
662
663pub struct FilterCursor<'a, 'b, F, T: Item, D> {
664 cursor: Cursor<'a, 'b, T, D>,
665 filter_node: F,
666}
667
668impl<'a, 'b, F, T: Item, D> FilterCursor<'a, 'b, F, T, D>
669where
670 F: FnMut(&T::Summary) -> bool,
671 T: Item,
672 D: Dimension<'a, T::Summary>,
673{
674 pub fn new(
675 tree: &'a SumTree<T>,
676 cx: <T::Summary as Summary>::Context<'b>,
677 filter_node: F,
678 ) -> Self {
679 let cursor = tree.cursor::<D>(cx);
680 Self {
681 cursor,
682 filter_node,
683 }
684 }
685
686 pub fn start(&self) -> &D {
687 self.cursor.start()
688 }
689
690 pub fn end(&self) -> D {
691 self.cursor.end()
692 }
693
694 pub fn item(&self) -> Option<&'a T> {
695 self.cursor.item()
696 }
697
698 pub fn item_summary(&self) -> Option<&'a T::Summary> {
699 self.cursor.item_summary()
700 }
701
702 pub fn next(&mut self) {
703 self.cursor.search_forward(&mut self.filter_node);
704 }
705
706 pub fn prev(&mut self) {
707 self.cursor.search_backward(&mut self.filter_node);
708 }
709}
710
711impl<'a, 'b, F, T: Item, U> Iterator for FilterCursor<'a, 'b, F, T, U>
712where
713 F: FnMut(&T::Summary) -> bool,
714 U: Dimension<'a, T::Summary>,
715{
716 type Item = &'a T;
717
718 fn next(&mut self) -> Option<Self::Item> {
719 if !self.cursor.did_seek {
720 self.next();
721 }
722
723 if let Some(item) = self.item() {
724 self.cursor.search_forward(&mut self.filter_node);
725 Some(item)
726 } else {
727 None
728 }
729 }
730}
731
732trait SeekAggregate<'a, T: Item> {
733 fn begin_leaf(&mut self);
734 fn end_leaf(&mut self, cx: <T::Summary as Summary>::Context<'_>);
735 fn push_item(
736 &mut self,
737 item: &'a T,
738 summary: &'a T::Summary,
739 cx: <T::Summary as Summary>::Context<'_>,
740 );
741 fn push_tree(
742 &mut self,
743 tree: &'a SumTree<T>,
744 summary: &'a T::Summary,
745 cx: <T::Summary as Summary>::Context<'_>,
746 );
747}
748
749struct SliceSeekAggregate<T: Item> {
750 tree: SumTree<T>,
751 leaf_items: ArrayVec<T, { 2 * TREE_BASE }>,
752 leaf_item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
753 leaf_summary: T::Summary,
754}
755
756struct SummarySeekAggregate<D>(D);
757
758impl<T: Item> SeekAggregate<'_, T> for () {
759 fn begin_leaf(&mut self) {}
760 fn end_leaf(&mut self, _: <T::Summary as Summary>::Context<'_>) {}
761 fn push_item(&mut self, _: &T, _: &T::Summary, _: <T::Summary as Summary>::Context<'_>) {}
762 fn push_tree(
763 &mut self,
764 _: &SumTree<T>,
765 _: &T::Summary,
766 _: <T::Summary as Summary>::Context<'_>,
767 ) {
768 }
769}
770
771impl<T: Item> SeekAggregate<'_, T> for SliceSeekAggregate<T> {
772 fn begin_leaf(&mut self) {}
773 fn end_leaf(&mut self, cx: <T::Summary as Summary>::Context<'_>) {
774 self.tree.append(
775 SumTree(Arc::new(Node::Leaf {
776 summary: mem::replace(&mut self.leaf_summary, <T::Summary as Summary>::zero(cx)),
777 items: mem::take(&mut self.leaf_items),
778 item_summaries: mem::take(&mut self.leaf_item_summaries),
779 })),
780 cx,
781 );
782 }
783 fn push_item(
784 &mut self,
785 item: &T,
786 summary: &T::Summary,
787 cx: <T::Summary as Summary>::Context<'_>,
788 ) {
789 self.leaf_items.push(item.clone());
790 self.leaf_item_summaries.push(summary.clone());
791 Summary::add_summary(&mut self.leaf_summary, summary, cx);
792 }
793 fn push_tree(
794 &mut self,
795 tree: &SumTree<T>,
796 _: &T::Summary,
797 cx: <T::Summary as Summary>::Context<'_>,
798 ) {
799 self.tree.append(tree.clone(), cx);
800 }
801}
802
803impl<'a, T: Item, D> SeekAggregate<'a, T> for SummarySeekAggregate<D>
804where
805 D: Dimension<'a, T::Summary>,
806{
807 fn begin_leaf(&mut self) {}
808 fn end_leaf(&mut self, _: <T::Summary as Summary>::Context<'_>) {}
809 fn push_item(
810 &mut self,
811 _: &T,
812 summary: &'a T::Summary,
813 cx: <T::Summary as Summary>::Context<'_>,
814 ) {
815 self.0.add_summary(summary, cx);
816 }
817 fn push_tree(
818 &mut self,
819 _: &SumTree<T>,
820 summary: &'a T::Summary,
821 cx: <T::Summary as Summary>::Context<'_>,
822 ) {
823 self.0.add_summary(summary, cx);
824 }
825}
826
827struct End<D>(PhantomData<D>);
828
829impl<D> End<D> {
830 fn new() -> Self {
831 Self(PhantomData)
832 }
833}
834
835impl<'a, S: Summary, D: Dimension<'a, S>> SeekTarget<'a, S, D> for End<D> {
836 fn cmp(&self, _: &D, _: S::Context<'_>) -> Ordering {
837 Ordering::Greater
838 }
839}
840
841impl<D> fmt::Debug for End<D> {
842 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
843 f.debug_tuple("End").finish()
844 }
845}