zed_sum_tree/
cursor.rs

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: u32,
9    position: D,
10}
11
12impl<'a, T: Item, D> StackEntry<'a, T, D> {
13    #[inline]
14    fn index(&self) -> usize {
15        self.index as usize
16    }
17}
18
19impl<T: Item + fmt::Debug, D: fmt::Debug> fmt::Debug for StackEntry<'_, T, D> {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        f.debug_struct("StackEntry")
22            .field("index", &self.index)
23            .field("position", &self.position)
24            .finish()
25    }
26}
27
28#[derive(Clone)]
29pub struct Cursor<'a, 'b, T: Item, D> {
30    tree: &'a SumTree<T>,
31    stack: ArrayVec<StackEntry<'a, T, D>, 16>,
32    position: D,
33    did_seek: bool,
34    at_end: bool,
35    cx: <T::Summary as Summary>::Context<'b>,
36}
37
38impl<T: Item + fmt::Debug, D: fmt::Debug> fmt::Debug for Cursor<'_, '_, T, D>
39where
40    T::Summary: fmt::Debug,
41{
42    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43        f.debug_struct("Cursor")
44            .field("tree", &self.tree)
45            .field("stack", &self.stack)
46            .field("position", &self.position)
47            .field("did_seek", &self.did_seek)
48            .field("at_end", &self.at_end)
49            .finish()
50    }
51}
52
53pub struct Iter<'a, T: Item> {
54    tree: &'a SumTree<T>,
55    stack: ArrayVec<StackEntry<'a, T, ()>, 16>,
56}
57
58impl<'a, 'b, T, D> Cursor<'a, 'b, T, D>
59where
60    T: Item,
61    D: Dimension<'a, T::Summary>,
62{
63    pub fn new(tree: &'a SumTree<T>, cx: <T::Summary as Summary>::Context<'b>) -> Self {
64        Self {
65            tree,
66            stack: ArrayVec::new(),
67            position: D::zero(cx),
68            did_seek: false,
69            at_end: tree.is_empty(),
70            cx,
71        }
72    }
73
74    fn reset(&mut self) {
75        self.did_seek = false;
76        self.at_end = self.tree.is_empty();
77        self.stack.truncate(0);
78        self.position = D::zero(self.cx);
79    }
80
81    pub fn start(&self) -> &D {
82        &self.position
83    }
84
85    #[track_caller]
86    pub fn end(&self) -> D {
87        if let Some(item_summary) = self.item_summary() {
88            let mut end = self.start().clone();
89            end.add_summary(item_summary, self.cx);
90            end
91        } else {
92            self.start().clone()
93        }
94    }
95
96    /// Item is None, when the list is empty, or this cursor is at the end of the list.
97    #[track_caller]
98    pub fn item(&self) -> Option<&'a T> {
99        self.assert_did_seek();
100        if let Some(entry) = self.stack.last() {
101            match *entry.tree.0 {
102                Node::Leaf { ref items, .. } => {
103                    if entry.index() == items.len() {
104                        None
105                    } else {
106                        Some(&items[entry.index()])
107                    }
108                }
109                _ => unreachable!(),
110            }
111        } else {
112            None
113        }
114    }
115
116    #[track_caller]
117    pub fn item_summary(&self) -> Option<&'a T::Summary> {
118        self.assert_did_seek();
119        if let Some(entry) = self.stack.last() {
120            match *entry.tree.0 {
121                Node::Leaf {
122                    ref item_summaries, ..
123                } => {
124                    if entry.index() == item_summaries.len() {
125                        None
126                    } else {
127                        Some(&item_summaries[entry.index()])
128                    }
129                }
130                _ => unreachable!(),
131            }
132        } else {
133            None
134        }
135    }
136
137    #[track_caller]
138    pub fn next_item(&self) -> Option<&'a T> {
139        self.assert_did_seek();
140        if let Some(entry) = self.stack.last() {
141            if entry.index() == entry.tree.0.items().len() - 1 {
142                if let Some(next_leaf) = self.next_leaf() {
143                    Some(next_leaf.0.items().first().unwrap())
144                } else {
145                    None
146                }
147            } else {
148                match *entry.tree.0 {
149                    Node::Leaf { ref items, .. } => Some(&items[entry.index() + 1]),
150                    _ => unreachable!(),
151                }
152            }
153        } else if self.at_end {
154            None
155        } else {
156            self.tree.first()
157        }
158    }
159
160    #[track_caller]
161    fn next_leaf(&self) -> Option<&'a SumTree<T>> {
162        for entry in self.stack.iter().rev().skip(1) {
163            if entry.index() < entry.tree.0.child_trees().len() - 1 {
164                match *entry.tree.0 {
165                    Node::Internal {
166                        ref child_trees, ..
167                    } => return Some(child_trees[entry.index() + 1].leftmost_leaf()),
168                    Node::Leaf { .. } => unreachable!(),
169                };
170            }
171        }
172        None
173    }
174
175    #[track_caller]
176    pub fn prev_item(&self) -> Option<&'a T> {
177        self.assert_did_seek();
178        if let Some(entry) = self.stack.last() {
179            if entry.index() == 0 {
180                if let Some(prev_leaf) = self.prev_leaf() {
181                    Some(prev_leaf.0.items().last().unwrap())
182                } else {
183                    None
184                }
185            } else {
186                match *entry.tree.0 {
187                    Node::Leaf { ref items, .. } => Some(&items[entry.index() - 1]),
188                    _ => unreachable!(),
189                }
190            }
191        } else if self.at_end {
192            self.tree.last()
193        } else {
194            None
195        }
196    }
197
198    #[track_caller]
199    fn prev_leaf(&self) -> Option<&'a SumTree<T>> {
200        for entry in self.stack.iter().rev().skip(1) {
201            if entry.index() != 0 {
202                match *entry.tree.0 {
203                    Node::Internal {
204                        ref child_trees, ..
205                    } => return Some(child_trees[entry.index() - 1].rightmost_leaf()),
206                    Node::Leaf { .. } => unreachable!(),
207                };
208            }
209        }
210        None
211    }
212
213    #[track_caller]
214    pub fn prev(&mut self) {
215        self.search_backward(|_| true)
216    }
217
218    #[track_caller]
219    pub fn search_backward<F>(&mut self, mut filter_node: F)
220    where
221        F: FnMut(&T::Summary) -> bool,
222    {
223        if !self.did_seek {
224            self.did_seek = true;
225            self.at_end = true;
226        }
227
228        if self.at_end {
229            self.position = D::zero(self.cx);
230            self.at_end = self.tree.is_empty();
231            if !self.tree.is_empty() {
232                self.stack.push(StackEntry {
233                    tree: self.tree,
234                    index: self.tree.0.child_summaries().len() as u32,
235                    position: D::from_summary(self.tree.summary(), self.cx),
236                });
237            }
238        }
239
240        let mut descending = false;
241        while !self.stack.is_empty() {
242            if let Some(StackEntry { position, .. }) = self.stack.iter().rev().nth(1) {
243                self.position = position.clone();
244            } else {
245                self.position = D::zero(self.cx);
246            }
247
248            let entry = self.stack.last_mut().unwrap();
249            if !descending {
250                if entry.index() == 0 {
251                    self.stack.pop();
252                    continue;
253                } else {
254                    entry.index -= 1;
255                }
256            }
257
258            for summary in &entry.tree.0.child_summaries()[..entry.index()] {
259                self.position.add_summary(summary, self.cx);
260            }
261            entry.position = self.position.clone();
262
263            descending = filter_node(&entry.tree.0.child_summaries()[entry.index()]);
264            match entry.tree.0.as_ref() {
265                Node::Internal { child_trees, .. } => {
266                    if descending {
267                        let tree = &child_trees[entry.index()];
268                        self.stack.push(StackEntry {
269                            position: D::zero(self.cx),
270                            tree,
271                            index: tree.0.child_summaries().len() as u32 - 1,
272                        })
273                    }
274                }
275                Node::Leaf { .. } => {
276                    if descending {
277                        break;
278                    }
279                }
280            }
281        }
282    }
283
284    #[track_caller]
285    pub fn next(&mut self) {
286        self.search_forward(|_| true)
287    }
288
289    #[track_caller]
290    pub fn search_forward<F>(&mut self, mut filter_node: F)
291    where
292        F: FnMut(&T::Summary) -> bool,
293    {
294        let mut descend = false;
295
296        if self.stack.is_empty() {
297            if !self.at_end {
298                self.stack.push(StackEntry {
299                    tree: self.tree,
300                    index: 0,
301                    position: D::zero(self.cx),
302                });
303                descend = true;
304            }
305            self.did_seek = true;
306        }
307
308        while !self.stack.is_empty() {
309            let new_subtree = {
310                let entry = self.stack.last_mut().unwrap();
311                match entry.tree.0.as_ref() {
312                    Node::Internal {
313                        child_trees,
314                        child_summaries,
315                        ..
316                    } => {
317                        if !descend {
318                            entry.index += 1;
319                            entry.position = self.position.clone();
320                        }
321
322                        while entry.index() < child_summaries.len() {
323                            let next_summary = &child_summaries[entry.index()];
324                            if filter_node(next_summary) {
325                                break;
326                            } else {
327                                entry.index += 1;
328                                entry.position.add_summary(next_summary, self.cx);
329                                self.position.add_summary(next_summary, self.cx);
330                            }
331                        }
332
333                        child_trees.get(entry.index())
334                    }
335                    Node::Leaf { item_summaries, .. } => {
336                        if !descend {
337                            let item_summary = &item_summaries[entry.index()];
338                            entry.index += 1;
339                            entry.position.add_summary(item_summary, self.cx);
340                            self.position.add_summary(item_summary, self.cx);
341                        }
342
343                        loop {
344                            if let Some(next_item_summary) = item_summaries.get(entry.index()) {
345                                if filter_node(next_item_summary) {
346                                    return;
347                                } else {
348                                    entry.index += 1;
349                                    entry.position.add_summary(next_item_summary, self.cx);
350                                    self.position.add_summary(next_item_summary, self.cx);
351                                }
352                            } else {
353                                break None;
354                            }
355                        }
356                    }
357                }
358            };
359
360            if let Some(subtree) = new_subtree {
361                descend = true;
362                self.stack.push(StackEntry {
363                    tree: subtree,
364                    index: 0,
365                    position: self.position.clone(),
366                });
367            } else {
368                descend = false;
369                self.stack.pop();
370            }
371        }
372
373        self.at_end = self.stack.is_empty();
374        debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
375    }
376
377    #[track_caller]
378    fn assert_did_seek(&self) {
379        assert!(
380            self.did_seek,
381            "Must call `seek`, `next` or `prev` before calling this method"
382        );
383    }
384}
385
386impl<'a, 'b, T, D> Cursor<'a, 'b, T, D>
387where
388    T: Item,
389    D: Dimension<'a, T::Summary>,
390{
391    #[track_caller]
392    pub fn seek<Target>(&mut self, pos: &Target, bias: Bias) -> bool
393    where
394        Target: SeekTarget<'a, T::Summary, D>,
395    {
396        self.reset();
397        self.seek_internal(pos, bias, &mut ())
398    }
399
400    #[track_caller]
401    pub fn seek_forward<Target>(&mut self, pos: &Target, bias: Bias) -> bool
402    where
403        Target: SeekTarget<'a, T::Summary, D>,
404    {
405        self.seek_internal(pos, bias, &mut ())
406    }
407
408    /// Advances the cursor and returns traversed items as a tree.
409    #[track_caller]
410    pub fn slice<Target>(&mut self, end: &Target, bias: Bias) -> SumTree<T>
411    where
412        Target: SeekTarget<'a, T::Summary, D>,
413    {
414        let mut slice = SliceSeekAggregate {
415            tree: SumTree::new(self.cx),
416            leaf_items: ArrayVec::new(),
417            leaf_item_summaries: ArrayVec::new(),
418            leaf_summary: <T::Summary as Summary>::zero(self.cx),
419        };
420        self.seek_internal(end, bias, &mut slice);
421        slice.tree
422    }
423
424    #[track_caller]
425    pub fn suffix(&mut self) -> SumTree<T> {
426        self.slice(&End::new(), Bias::Right)
427    }
428
429    #[track_caller]
430    pub fn summary<Target, Output>(&mut self, end: &Target, bias: Bias) -> Output
431    where
432        Target: SeekTarget<'a, T::Summary, D>,
433        Output: Dimension<'a, T::Summary>,
434    {
435        let mut summary = SummarySeekAggregate(Output::zero(self.cx));
436        self.seek_internal(end, bias, &mut summary);
437        summary.0
438    }
439
440    /// Returns whether we found the item you were seeking for
441    #[track_caller]
442    fn seek_internal(
443        &mut self,
444        target: &dyn SeekTarget<'a, T::Summary, D>,
445        bias: Bias,
446        aggregate: &mut dyn SeekAggregate<'a, T>,
447    ) -> bool {
448        assert!(
449            target.cmp(&self.position, self.cx) >= Ordering::Equal,
450            "cannot seek backward",
451        );
452
453        if !self.did_seek {
454            self.did_seek = true;
455            self.stack.push(StackEntry {
456                tree: self.tree,
457                index: 0,
458                position: D::zero(self.cx),
459            });
460        }
461
462        let mut ascending = false;
463        'outer: while let Some(entry) = self.stack.last_mut() {
464            match *entry.tree.0 {
465                Node::Internal {
466                    ref child_summaries,
467                    ref child_trees,
468                    ..
469                } => {
470                    if ascending {
471                        entry.index += 1;
472                        entry.position = self.position.clone();
473                    }
474
475                    for (child_tree, child_summary) in child_trees[entry.index()..]
476                        .iter()
477                        .zip(&child_summaries[entry.index()..])
478                    {
479                        let mut child_end = self.position.clone();
480                        child_end.add_summary(child_summary, self.cx);
481
482                        let comparison = target.cmp(&child_end, self.cx);
483                        if comparison == Ordering::Greater
484                            || (comparison == Ordering::Equal && bias == Bias::Right)
485                        {
486                            self.position = child_end;
487                            aggregate.push_tree(child_tree, child_summary, self.cx);
488                            entry.index += 1;
489                            entry.position = self.position.clone();
490                        } else {
491                            self.stack.push(StackEntry {
492                                tree: child_tree,
493                                index: 0,
494                                position: self.position.clone(),
495                            });
496                            ascending = false;
497                            continue 'outer;
498                        }
499                    }
500                }
501                Node::Leaf {
502                    ref items,
503                    ref item_summaries,
504                    ..
505                } => {
506                    aggregate.begin_leaf();
507
508                    for (item, item_summary) in items[entry.index()..]
509                        .iter()
510                        .zip(&item_summaries[entry.index()..])
511                    {
512                        let mut child_end = self.position.clone();
513                        child_end.add_summary(item_summary, self.cx);
514
515                        let comparison = target.cmp(&child_end, self.cx);
516                        if comparison == Ordering::Greater
517                            || (comparison == Ordering::Equal && bias == Bias::Right)
518                        {
519                            self.position = child_end;
520                            aggregate.push_item(item, item_summary, self.cx);
521                            entry.index += 1;
522                        } else {
523                            aggregate.end_leaf(self.cx);
524                            break 'outer;
525                        }
526                    }
527
528                    aggregate.end_leaf(self.cx);
529                }
530            }
531
532            self.stack.pop();
533            ascending = true;
534        }
535
536        self.at_end = self.stack.is_empty();
537        debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
538
539        let mut end = self.position.clone();
540        if bias == Bias::Left
541            && let Some(summary) = self.item_summary()
542        {
543            end.add_summary(summary, self.cx);
544        }
545
546        target.cmp(&end, self.cx) == Ordering::Equal
547    }
548}
549
550impl<'a, T: Item> Iter<'a, T> {
551    pub(crate) fn new(tree: &'a SumTree<T>) -> Self {
552        Self {
553            tree,
554            stack: Default::default(),
555        }
556    }
557}
558
559impl<'a, T: Item> Iterator for Iter<'a, T> {
560    type Item = &'a T;
561
562    fn next(&mut self) -> Option<Self::Item> {
563        let mut descend = false;
564
565        if self.stack.is_empty() {
566            self.stack.push(StackEntry {
567                tree: self.tree,
568                index: 0,
569                position: (),
570            });
571            descend = true;
572        }
573
574        while !self.stack.is_empty() {
575            let new_subtree = {
576                let entry = self.stack.last_mut().unwrap();
577                match entry.tree.0.as_ref() {
578                    Node::Internal { child_trees, .. } => {
579                        if !descend {
580                            entry.index += 1;
581                        }
582                        child_trees.get(entry.index())
583                    }
584                    Node::Leaf { items, .. } => {
585                        if !descend {
586                            entry.index += 1;
587                        }
588
589                        if let Some(next_item) = items.get(entry.index()) {
590                            return Some(next_item);
591                        } else {
592                            None
593                        }
594                    }
595                }
596            };
597
598            if let Some(subtree) = new_subtree {
599                descend = true;
600                self.stack.push(StackEntry {
601                    tree: subtree,
602                    index: 0,
603                    position: (),
604                });
605            } else {
606                descend = false;
607                self.stack.pop();
608            }
609        }
610
611        None
612    }
613}
614
615impl<'a, 'b, T: Item, D> Iterator for Cursor<'a, 'b, T, D>
616where
617    D: Dimension<'a, T::Summary>,
618{
619    type Item = &'a T;
620
621    fn next(&mut self) -> Option<Self::Item> {
622        if !self.did_seek {
623            self.next();
624        }
625
626        if let Some(item) = self.item() {
627            self.next();
628            Some(item)
629        } else {
630            None
631        }
632    }
633}
634
635pub struct FilterCursor<'a, 'b, F, T: Item, D> {
636    cursor: Cursor<'a, 'b, T, D>,
637    filter_node: F,
638}
639
640impl<'a, 'b, F, T: Item, D> FilterCursor<'a, 'b, F, T, D>
641where
642    F: FnMut(&T::Summary) -> bool,
643    T: Item,
644    D: Dimension<'a, T::Summary>,
645{
646    pub fn new(
647        tree: &'a SumTree<T>,
648        cx: <T::Summary as Summary>::Context<'b>,
649        filter_node: F,
650    ) -> Self {
651        let cursor = tree.cursor::<D>(cx);
652        Self {
653            cursor,
654            filter_node,
655        }
656    }
657
658    pub fn start(&self) -> &D {
659        self.cursor.start()
660    }
661
662    pub fn end(&self) -> D {
663        self.cursor.end()
664    }
665
666    pub fn item(&self) -> Option<&'a T> {
667        self.cursor.item()
668    }
669
670    pub fn item_summary(&self) -> Option<&'a T::Summary> {
671        self.cursor.item_summary()
672    }
673
674    pub fn next(&mut self) {
675        self.cursor.search_forward(&mut self.filter_node);
676    }
677
678    pub fn prev(&mut self) {
679        self.cursor.search_backward(&mut self.filter_node);
680    }
681}
682
683impl<'a, 'b, F, T: Item, U> Iterator for FilterCursor<'a, 'b, F, T, U>
684where
685    F: FnMut(&T::Summary) -> bool,
686    U: Dimension<'a, T::Summary>,
687{
688    type Item = &'a T;
689
690    fn next(&mut self) -> Option<Self::Item> {
691        if !self.cursor.did_seek {
692            self.next();
693        }
694
695        if let Some(item) = self.item() {
696            self.cursor.search_forward(&mut self.filter_node);
697            Some(item)
698        } else {
699            None
700        }
701    }
702}
703
704trait SeekAggregate<'a, T: Item> {
705    fn begin_leaf(&mut self);
706    fn end_leaf(&mut self, cx: <T::Summary as Summary>::Context<'_>);
707    fn push_item(
708        &mut self,
709        item: &'a T,
710        summary: &'a T::Summary,
711        cx: <T::Summary as Summary>::Context<'_>,
712    );
713    fn push_tree(
714        &mut self,
715        tree: &'a SumTree<T>,
716        summary: &'a T::Summary,
717        cx: <T::Summary as Summary>::Context<'_>,
718    );
719}
720
721struct SliceSeekAggregate<T: Item> {
722    tree: SumTree<T>,
723    leaf_items: ArrayVec<T, { 2 * TREE_BASE }>,
724    leaf_item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
725    leaf_summary: T::Summary,
726}
727
728struct SummarySeekAggregate<D>(D);
729
730impl<T: Item> SeekAggregate<'_, T> for () {
731    fn begin_leaf(&mut self) {}
732    fn end_leaf(&mut self, _: <T::Summary as Summary>::Context<'_>) {}
733    fn push_item(&mut self, _: &T, _: &T::Summary, _: <T::Summary as Summary>::Context<'_>) {}
734    fn push_tree(
735        &mut self,
736        _: &SumTree<T>,
737        _: &T::Summary,
738        _: <T::Summary as Summary>::Context<'_>,
739    ) {
740    }
741}
742
743impl<T: Item> SeekAggregate<'_, T> for SliceSeekAggregate<T> {
744    fn begin_leaf(&mut self) {}
745    fn end_leaf(&mut self, cx: <T::Summary as Summary>::Context<'_>) {
746        self.tree.append(
747            SumTree(Arc::new(Node::Leaf {
748                summary: mem::replace(&mut self.leaf_summary, <T::Summary as Summary>::zero(cx)),
749                items: mem::take(&mut self.leaf_items),
750                item_summaries: mem::take(&mut self.leaf_item_summaries),
751            })),
752            cx,
753        );
754    }
755    fn push_item(
756        &mut self,
757        item: &T,
758        summary: &T::Summary,
759        cx: <T::Summary as Summary>::Context<'_>,
760    ) {
761        self.leaf_items.push(item.clone());
762        self.leaf_item_summaries.push(summary.clone());
763        Summary::add_summary(&mut self.leaf_summary, summary, cx);
764    }
765    fn push_tree(
766        &mut self,
767        tree: &SumTree<T>,
768        _: &T::Summary,
769        cx: <T::Summary as Summary>::Context<'_>,
770    ) {
771        self.tree.append(tree.clone(), cx);
772    }
773}
774
775impl<'a, T: Item, D> SeekAggregate<'a, T> for SummarySeekAggregate<D>
776where
777    D: Dimension<'a, T::Summary>,
778{
779    fn begin_leaf(&mut self) {}
780    fn end_leaf(&mut self, _: <T::Summary as Summary>::Context<'_>) {}
781    fn push_item(
782        &mut self,
783        _: &T,
784        summary: &'a T::Summary,
785        cx: <T::Summary as Summary>::Context<'_>,
786    ) {
787        self.0.add_summary(summary, cx);
788    }
789    fn push_tree(
790        &mut self,
791        _: &SumTree<T>,
792        summary: &'a T::Summary,
793        cx: <T::Summary as Summary>::Context<'_>,
794    ) {
795        self.0.add_summary(summary, cx);
796    }
797}
798
799struct End<D>(PhantomData<D>);
800
801impl<D> End<D> {
802    fn new() -> Self {
803        Self(PhantomData)
804    }
805}
806
807impl<'a, S: Summary, D: Dimension<'a, S>> SeekTarget<'a, S, D> for End<D> {
808    fn cmp(&self, _: &D, _: S::Context<'_>) -> Ordering {
809        Ordering::Greater
810    }
811}
812
813impl<D> fmt::Debug for End<D> {
814    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
815        f.debug_tuple("End").finish()
816    }
817}
zed_sum_tree/cursor.rs

zed_sum_tree/
cursor.rs
