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