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