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