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