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