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 #[track_caller]
410 pub fn slice<Target>(
411 &mut self,
412 end: &Target,
413 bias: Bias,
414 cx: &<T::Summary as Summary>::Context,
415 ) -> SumTree<T>
416 where
417 Target: SeekTarget<'a, T::Summary, D>,
418 {
419 let mut slice = SliceSeekAggregate {
420 tree: SumTree::new(cx),
421 leaf_items: ArrayVec::new(),
422 leaf_item_summaries: ArrayVec::new(),
423 leaf_summary: <T::Summary as Summary>::zero(cx),
424 };
425 self.seek_internal(end, bias, &mut slice, cx);
426 slice.tree
427 }
428
429 #[track_caller]
430 pub fn suffix(&mut self, cx: &<T::Summary as Summary>::Context) -> SumTree<T> {
431 self.slice(&End::new(), Bias::Right, cx)
432 }
433
434 #[track_caller]
435 pub fn summary<Target, Output>(
436 &mut self,
437 end: &Target,
438 bias: Bias,
439 cx: &<T::Summary as Summary>::Context,
440 ) -> Output
441 where
442 Target: SeekTarget<'a, T::Summary, D>,
443 Output: Dimension<'a, T::Summary>,
444 {
445 let mut summary = SummarySeekAggregate(Output::zero(cx));
446 self.seek_internal(end, bias, &mut summary, cx);
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 cx: &<T::Summary as Summary>::Context,
458 ) -> bool {
459 assert!(
460 target.cmp(&self.position, cx) >= Ordering::Equal,
461 "cannot seek backward",
462 );
463
464 if !self.did_seek {
465 self.did_seek = true;
466 self.stack.push(StackEntry {
467 tree: self.tree,
468 index: 0,
469 position: D::zero(cx),
470 });
471 }
472
473 let mut ascending = false;
474 'outer: while let Some(entry) = self.stack.last_mut() {
475 match *entry.tree.0 {
476 Node::Internal {
477 ref child_summaries,
478 ref child_trees,
479 ..
480 } => {
481 if ascending {
482 entry.index += 1;
483 entry.position = self.position.clone();
484 }
485
486 for (child_tree, child_summary) in child_trees[entry.index..]
487 .iter()
488 .zip(&child_summaries[entry.index..])
489 {
490 let mut child_end = self.position.clone();
491 child_end.add_summary(child_summary, cx);
492
493 let comparison = target.cmp(&child_end, cx);
494 if comparison == Ordering::Greater
495 || (comparison == Ordering::Equal && bias == Bias::Right)
496 {
497 self.position = child_end;
498 aggregate.push_tree(child_tree, child_summary, cx);
499 entry.index += 1;
500 entry.position = self.position.clone();
501 } else {
502 self.stack.push(StackEntry {
503 tree: child_tree,
504 index: 0,
505 position: self.position.clone(),
506 });
507 ascending = false;
508 continue 'outer;
509 }
510 }
511 }
512 Node::Leaf {
513 ref items,
514 ref item_summaries,
515 ..
516 } => {
517 aggregate.begin_leaf();
518
519 for (item, item_summary) in items[entry.index..]
520 .iter()
521 .zip(&item_summaries[entry.index..])
522 {
523 let mut child_end = self.position.clone();
524 child_end.add_summary(item_summary, cx);
525
526 let comparison = target.cmp(&child_end, cx);
527 if comparison == Ordering::Greater
528 || (comparison == Ordering::Equal && bias == Bias::Right)
529 {
530 self.position = child_end;
531 aggregate.push_item(item, item_summary, cx);
532 entry.index += 1;
533 } else {
534 aggregate.end_leaf(cx);
535 break 'outer;
536 }
537 }
538
539 aggregate.end_leaf(cx);
540 }
541 }
542
543 self.stack.pop();
544 ascending = true;
545 }
546
547 self.at_end = self.stack.is_empty();
548 debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
549
550 let mut end = self.position.clone();
551 if bias == Bias::Left {
552 if let Some(summary) = self.item_summary() {
553 end.add_summary(summary, cx);
554 }
555 }
556
557 target.cmp(&end, cx) == Ordering::Equal
558 }
559}
560
561impl<'a, T: Item> Iter<'a, T> {
562 pub(crate) fn new(tree: &'a SumTree<T>) -> Self {
563 Self {
564 tree,
565 stack: Default::default(),
566 }
567 }
568}
569
570impl<'a, T: Item> Iterator for Iter<'a, T> {
571 type Item = &'a T;
572
573 fn next(&mut self) -> Option<Self::Item> {
574 let mut descend = false;
575
576 if self.stack.is_empty() {
577 self.stack.push(StackEntry {
578 tree: self.tree,
579 index: 0,
580 position: (),
581 });
582 descend = true;
583 }
584
585 while !self.stack.is_empty() {
586 let new_subtree = {
587 let entry = self.stack.last_mut().unwrap();
588 match entry.tree.0.as_ref() {
589 Node::Internal { child_trees, .. } => {
590 if !descend {
591 entry.index += 1;
592 }
593 child_trees.get(entry.index)
594 }
595 Node::Leaf { items, .. } => {
596 if !descend {
597 entry.index += 1;
598 }
599
600 if let Some(next_item) = items.get(entry.index) {
601 return Some(next_item);
602 } else {
603 None
604 }
605 }
606 }
607 };
608
609 if let Some(subtree) = new_subtree {
610 descend = true;
611 self.stack.push(StackEntry {
612 tree: subtree,
613 index: 0,
614 position: (),
615 });
616 } else {
617 descend = false;
618 self.stack.pop();
619 }
620 }
621
622 None
623 }
624}
625
626impl<'a, T, S, D> Iterator for Cursor<'a, T, D>
627where
628 T: Item<Summary = S>,
629 S: Summary<Context = ()>,
630 D: Dimension<'a, T::Summary>,
631{
632 type Item = &'a T;
633
634 fn next(&mut self) -> Option<Self::Item> {
635 if !self.did_seek {
636 self.next(&());
637 }
638
639 if let Some(item) = self.item() {
640 self.next(&());
641 Some(item)
642 } else {
643 None
644 }
645 }
646}
647
648pub struct FilterCursor<'a, F, T: Item, D> {
649 cursor: Cursor<'a, T, D>,
650 filter_node: F,
651}
652
653impl<'a, F, T, D> FilterCursor<'a, F, T, D>
654where
655 F: FnMut(&T::Summary) -> bool,
656 T: Item,
657 D: Dimension<'a, T::Summary>,
658{
659 pub fn new(
660 tree: &'a SumTree<T>,
661 cx: &<T::Summary as Summary>::Context,
662 filter_node: F,
663 ) -> Self {
664 let cursor = tree.cursor::<D>(cx);
665 Self {
666 cursor,
667 filter_node,
668 }
669 }
670
671 pub fn start(&self) -> &D {
672 self.cursor.start()
673 }
674
675 pub fn end(&self, cx: &<T::Summary as Summary>::Context) -> D {
676 self.cursor.end(cx)
677 }
678
679 pub fn item(&self) -> Option<&'a T> {
680 self.cursor.item()
681 }
682
683 pub fn item_summary(&self) -> Option<&'a T::Summary> {
684 self.cursor.item_summary()
685 }
686
687 pub fn next(&mut self, cx: &<T::Summary as Summary>::Context) {
688 self.cursor.search_forward(&mut self.filter_node, cx);
689 }
690
691 pub fn prev(&mut self, cx: &<T::Summary as Summary>::Context) {
692 self.cursor.search_backward(&mut self.filter_node, cx);
693 }
694}
695
696impl<'a, F, T, S, U> Iterator for FilterCursor<'a, F, T, U>
697where
698 F: FnMut(&T::Summary) -> bool,
699 T: Item<Summary = S>,
700 S: Summary<Context = ()>, //Context for the summary must be unit type, as .next() doesn't take arguments
701 U: Dimension<'a, T::Summary>,
702{
703 type Item = &'a T;
704
705 fn next(&mut self) -> Option<Self::Item> {
706 if !self.cursor.did_seek {
707 self.next(&());
708 }
709
710 if let Some(item) = self.item() {
711 self.cursor.search_forward(&mut self.filter_node, &());
712 Some(item)
713 } else {
714 None
715 }
716 }
717}
718
719trait SeekAggregate<'a, T: Item> {
720 fn begin_leaf(&mut self);
721 fn end_leaf(&mut self, cx: &<T::Summary as Summary>::Context);
722 fn push_item(
723 &mut self,
724 item: &'a T,
725 summary: &'a T::Summary,
726 cx: &<T::Summary as Summary>::Context,
727 );
728 fn push_tree(
729 &mut self,
730 tree: &'a SumTree<T>,
731 summary: &'a T::Summary,
732 cx: &<T::Summary as Summary>::Context,
733 );
734}
735
736struct SliceSeekAggregate<T: Item> {
737 tree: SumTree<T>,
738 leaf_items: ArrayVec<T, { 2 * TREE_BASE }>,
739 leaf_item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
740 leaf_summary: T::Summary,
741}
742
743struct SummarySeekAggregate<D>(D);
744
745impl<T: Item> SeekAggregate<'_, T> for () {
746 fn begin_leaf(&mut self) {}
747 fn end_leaf(&mut self, _: &<T::Summary as Summary>::Context) {}
748 fn push_item(&mut self, _: &T, _: &T::Summary, _: &<T::Summary as Summary>::Context) {}
749 fn push_tree(&mut self, _: &SumTree<T>, _: &T::Summary, _: &<T::Summary as Summary>::Context) {}
750}
751
752impl<T: Item> SeekAggregate<'_, T> for SliceSeekAggregate<T> {
753 fn begin_leaf(&mut self) {}
754 fn end_leaf(&mut self, cx: &<T::Summary as Summary>::Context) {
755 self.tree.append(
756 SumTree(Arc::new(Node::Leaf {
757 summary: mem::replace(&mut self.leaf_summary, <T::Summary as Summary>::zero(cx)),
758 items: mem::take(&mut self.leaf_items),
759 item_summaries: mem::take(&mut self.leaf_item_summaries),
760 })),
761 cx,
762 );
763 }
764 fn push_item(&mut self, item: &T, summary: &T::Summary, cx: &<T::Summary as Summary>::Context) {
765 self.leaf_items.push(item.clone());
766 self.leaf_item_summaries.push(summary.clone());
767 Summary::add_summary(&mut self.leaf_summary, summary, cx);
768 }
769 fn push_tree(
770 &mut self,
771 tree: &SumTree<T>,
772 _: &T::Summary,
773 cx: &<T::Summary as Summary>::Context,
774 ) {
775 self.tree.append(tree.clone(), cx);
776 }
777}
778
779impl<'a, T: Item, D> SeekAggregate<'a, T> for SummarySeekAggregate<D>
780where
781 D: Dimension<'a, T::Summary>,
782{
783 fn begin_leaf(&mut self) {}
784 fn end_leaf(&mut self, _: &<T::Summary as Summary>::Context) {}
785 fn push_item(&mut self, _: &T, summary: &'a T::Summary, cx: &<T::Summary as Summary>::Context) {
786 self.0.add_summary(summary, cx);
787 }
788 fn push_tree(
789 &mut self,
790 _: &SumTree<T>,
791 summary: &'a T::Summary,
792 cx: &<T::Summary as Summary>::Context,
793 ) {
794 self.0.add_summary(summary, cx);
795 }
796}