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