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