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