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