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