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