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