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