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