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