sum_tree/
sum_tree.rs

1mod cursor;
2#[cfg(any(test, feature = "test-support"))]
3pub mod property_test;
4mod tree_map;
5
6pub use cursor::{Cursor, FilterCursor, Iter};
7use heapless::Vec as ArrayVec;
8use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator as _};
9use std::marker::PhantomData;
10use std::mem;
11use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
12pub use tree_map::{MapSeekTarget, TreeMap, TreeSet};
13
14#[cfg(test)]
15pub const TREE_BASE: usize = 2;
16#[cfg(not(test))]
17pub const TREE_BASE: usize = 6;
18
19// Helper for when we cannot use ArrayVec::<T>::push().unwrap() as T doesn't impl Debug
20trait CapacityResultExt {
21    fn unwrap_oob(self);
22}
23
24impl<T> CapacityResultExt for Result<(), T> {
25    fn unwrap_oob(self) {
26        self.unwrap_or_else(|_| panic!("item should fit into fixed size ArrayVec"))
27    }
28}
29
30/// An item that can be stored in a [`SumTree`]
31///
32/// Must be summarized by a type that implements [`Summary`]
33pub trait Item: Clone {
34    type Summary: Summary;
35
36    fn summary(&self, cx: <Self::Summary as Summary>::Context<'_>) -> Self::Summary;
37}
38
39/// An [`Item`] whose summary has a specific key that can be used to identify it
40pub trait KeyedItem: Item {
41    type Key: for<'a> Dimension<'a, Self::Summary> + Ord;
42
43    fn key(&self) -> Self::Key;
44}
45
46/// A type that describes the Sum of all [`Item`]s in a subtree of the [`SumTree`]
47///
48/// Each Summary type can have multiple [`Dimension`]s that it measures,
49/// which can be used to navigate the tree
50pub trait Summary: Clone {
51    type Context<'a>: Copy;
52    fn zero<'a>(cx: Self::Context<'a>) -> Self;
53    fn add_summary<'a>(&mut self, summary: &Self, cx: Self::Context<'a>);
54}
55
56pub trait ContextLessSummary: Clone {
57    fn zero() -> Self;
58    fn add_summary(&mut self, summary: &Self);
59}
60
61impl<T: ContextLessSummary> Summary for T {
62    type Context<'a> = ();
63
64    fn zero<'a>((): ()) -> Self {
65        T::zero()
66    }
67
68    fn add_summary<'a>(&mut self, summary: &Self, (): ()) {
69        T::add_summary(self, summary)
70    }
71}
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub struct NoSummary;
75
76/// Catch-all implementation for when you need something that implements [`Summary`] without a specific type.
77/// We implement it on a `NoSummary` instead of re-using `()`, as that avoids blanket impl collisions with `impl<T: Summary> Dimension for T`
78/// (as we also need unit type to be a fill-in dimension)
79impl ContextLessSummary for NoSummary {
80    fn zero() -> Self {
81        NoSummary
82    }
83
84    fn add_summary(&mut self, _: &Self) {}
85}
86
87/// Each [`Summary`] type can have more than one [`Dimension`] type that it measures.
88///
89/// You can use dimensions to seek to a specific location in the [`SumTree`]
90///
91/// # Example:
92/// Zed's rope has a `TextSummary` type that summarizes lines, characters, and bytes.
93/// Each of these are different dimensions we may want to seek to
94pub trait Dimension<'a, S: Summary>: Clone {
95    fn zero(cx: S::Context<'_>) -> Self;
96
97    fn add_summary(&mut self, summary: &'a S, cx: S::Context<'_>);
98    #[must_use]
99    fn with_added_summary(mut self, summary: &'a S, cx: S::Context<'_>) -> Self {
100        self.add_summary(summary, cx);
101        self
102    }
103
104    fn from_summary(summary: &'a S, cx: S::Context<'_>) -> Self {
105        let mut dimension = Self::zero(cx);
106        dimension.add_summary(summary, cx);
107        dimension
108    }
109}
110
111impl<'a, T: Summary> Dimension<'a, T> for T {
112    fn zero(cx: T::Context<'_>) -> Self {
113        Summary::zero(cx)
114    }
115
116    fn add_summary(&mut self, summary: &'a T, cx: T::Context<'_>) {
117        Summary::add_summary(self, summary, cx);
118    }
119}
120
121pub trait SeekTarget<'a, S: Summary, D: Dimension<'a, S>> {
122    fn cmp(&self, cursor_location: &D, cx: S::Context<'_>) -> Ordering;
123}
124
125impl<'a, S: Summary, D: Dimension<'a, S> + Ord> SeekTarget<'a, S, D> for D {
126    fn cmp(&self, cursor_location: &Self, _: S::Context<'_>) -> Ordering {
127        Ord::cmp(self, cursor_location)
128    }
129}
130
131impl<'a, T: Summary> Dimension<'a, T> for () {
132    fn zero(_: T::Context<'_>) -> Self {}
133
134    fn add_summary(&mut self, _: &'a T, _: T::Context<'_>) {}
135}
136
137#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
138pub struct Dimensions<D1, D2, D3 = ()>(pub D1, pub D2, pub D3);
139
140impl<'a, T: Summary, D1: Dimension<'a, T>, D2: Dimension<'a, T>, D3: Dimension<'a, T>>
141    Dimension<'a, T> for Dimensions<D1, D2, D3>
142{
143    fn zero(cx: T::Context<'_>) -> Self {
144        Dimensions(D1::zero(cx), D2::zero(cx), D3::zero(cx))
145    }
146
147    fn add_summary(&mut self, summary: &'a T, cx: T::Context<'_>) {
148        self.0.add_summary(summary, cx);
149        self.1.add_summary(summary, cx);
150        self.2.add_summary(summary, cx);
151    }
152}
153
154impl<'a, S, D1, D2, D3> SeekTarget<'a, S, Dimensions<D1, D2, D3>> for D1
155where
156    S: Summary,
157    D1: SeekTarget<'a, S, D1> + Dimension<'a, S>,
158    D2: Dimension<'a, S>,
159    D3: Dimension<'a, S>,
160{
161    fn cmp(&self, cursor_location: &Dimensions<D1, D2, D3>, cx: S::Context<'_>) -> Ordering {
162        self.cmp(&cursor_location.0, cx)
163    }
164}
165
166/// Bias is used to settle ambiguities when determining positions in an ordered sequence.
167///
168/// The primary use case is for text, where Bias influences
169/// which character an offset or anchor is associated with.
170///
171/// # Examples
172/// Given the buffer `AˇBCD`:
173/// - The offset of the cursor is 1
174/// - [Bias::Left] would attach the cursor to the character `A`
175/// - [Bias::Right] would attach the cursor to the character `B`
176///
177/// Given the buffer `A«BCˇ»D`:
178/// - The offset of the cursor is 3, and the selection is from 1 to 3
179/// - The left anchor of the selection has [Bias::Right], attaching it to the character `B`
180/// - The right anchor of the selection has [Bias::Left], attaching it to the character `C`
181///
182/// Given the buffer `{ˇ<...>`, where `<...>` is a folded region:
183/// - The display offset of the cursor is 1, but the offset in the buffer is determined by the bias
184/// - [Bias::Left] would attach the cursor to the character `{`, with a buffer offset of 1
185/// - [Bias::Right] would attach the cursor to the first character of the folded region,
186///   and the buffer offset would be the offset of the first character of the folded region
187#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Debug, Hash, Default)]
188pub enum Bias {
189    /// Attach to the character on the left
190    #[default]
191    Left,
192    /// Attach to the character on the right
193    Right,
194}
195
196impl Bias {
197    pub fn invert(self) -> Self {
198        match self {
199            Self::Left => Self::Right,
200            Self::Right => Self::Left,
201        }
202    }
203}
204
205/// A B+ tree in which each leaf node contains `Item`s of type `T` and a `Summary`s for each `Item`.
206/// Each internal node contains a `Summary` of the items in its subtree.
207///
208/// The maximum number of items per node is `TREE_BASE * 2`.
209///
210/// Any [`Dimension`] supported by the [`Summary`] type can be used to seek to a specific location in the tree.
211#[derive(Clone)]
212pub struct SumTree<T: Item>(Arc<Node<T>>);
213
214impl<T> fmt::Debug for SumTree<T>
215where
216    T: fmt::Debug + Item,
217    T::Summary: fmt::Debug,
218{
219    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
220        f.debug_tuple("SumTree").field(&self.0).finish()
221    }
222}
223
224impl<T: Item> SumTree<T> {
225    pub fn new(cx: <T::Summary as Summary>::Context<'_>) -> Self {
226        SumTree(Arc::new(Node::Leaf {
227            summary: <T::Summary as Summary>::zero(cx),
228            items: ArrayVec::new(),
229            item_summaries: ArrayVec::new(),
230        }))
231    }
232
233    /// Useful in cases where the item type has a non-trivial context type, but the zero value of the summary type doesn't depend on that context.
234    pub fn from_summary(summary: T::Summary) -> Self {
235        SumTree(Arc::new(Node::Leaf {
236            summary,
237            items: ArrayVec::new(),
238            item_summaries: ArrayVec::new(),
239        }))
240    }
241
242    pub fn from_item(item: T, cx: <T::Summary as Summary>::Context<'_>) -> Self {
243        let mut tree = Self::new(cx);
244        tree.push(item, cx);
245        tree
246    }
247
248    pub fn from_iter<I: IntoIterator<Item = T>>(
249        iter: I,
250        cx: <T::Summary as Summary>::Context<'_>,
251    ) -> Self {
252        let mut nodes = Vec::new();
253
254        let mut iter = iter.into_iter().fuse().peekable();
255        while iter.peek().is_some() {
256            let items: ArrayVec<T, { 2 * TREE_BASE }, u8> =
257                iter.by_ref().take(2 * TREE_BASE).collect();
258            let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }, u8> =
259                items.iter().map(|item| item.summary(cx)).collect();
260
261            let mut summary = item_summaries[0].clone();
262            for item_summary in &item_summaries[1..] {
263                <T::Summary as Summary>::add_summary(&mut summary, item_summary, cx);
264            }
265
266            nodes.push(SumTree(Arc::new(Node::Leaf {
267                summary,
268                items,
269                item_summaries,
270            })));
271        }
272
273        let mut parent_nodes = Vec::new();
274        let mut height = 0;
275        while nodes.len() > 1 {
276            height += 1;
277            let mut current_parent_node = None;
278            for child_node in nodes.drain(..) {
279                let parent_node = current_parent_node.get_or_insert_with(|| {
280                    SumTree(Arc::new(Node::Internal {
281                        summary: <T::Summary as Summary>::zero(cx),
282                        height,
283                        child_summaries: ArrayVec::new(),
284                        child_trees: ArrayVec::new(),
285                    }))
286                });
287                let Node::Internal {
288                    summary,
289                    child_summaries,
290                    child_trees,
291                    ..
292                } = Arc::get_mut(&mut parent_node.0).unwrap()
293                else {
294                    unreachable!()
295                };
296                let child_summary = child_node.summary();
297                <T::Summary as Summary>::add_summary(summary, child_summary, cx);
298                child_summaries.push(child_summary.clone()).unwrap_oob();
299                child_trees.push(child_node.clone()).unwrap_oob();
300
301                if child_trees.len() == 2 * TREE_BASE {
302                    parent_nodes.extend(current_parent_node.take());
303                }
304            }
305            parent_nodes.extend(current_parent_node.take());
306            mem::swap(&mut nodes, &mut parent_nodes);
307        }
308
309        if nodes.is_empty() {
310            Self::new(cx)
311        } else {
312            debug_assert_eq!(nodes.len(), 1);
313            nodes.pop().unwrap()
314        }
315    }
316
317    pub fn from_par_iter<I, Iter>(iter: I, cx: <T::Summary as Summary>::Context<'_>) -> Self
318    where
319        I: IntoParallelIterator<Iter = Iter>,
320        Iter: IndexedParallelIterator<Item = T>,
321        T: Send + Sync,
322        T::Summary: Send + Sync,
323        for<'a> <T::Summary as Summary>::Context<'a>: Sync,
324    {
325        let mut nodes = iter
326            .into_par_iter()
327            .chunks(2 * TREE_BASE)
328            .map(|items| {
329                let items: ArrayVec<T, { 2 * TREE_BASE }, u8> = items.into_iter().collect();
330                let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }, u8> =
331                    items.iter().map(|item| item.summary(cx)).collect();
332                let mut summary = item_summaries[0].clone();
333                for item_summary in &item_summaries[1..] {
334                    <T::Summary as Summary>::add_summary(&mut summary, item_summary, cx);
335                }
336                SumTree(Arc::new(Node::Leaf {
337                    summary,
338                    items,
339                    item_summaries,
340                }))
341            })
342            .collect::<Vec<_>>();
343
344        let mut height = 0;
345        while nodes.len() > 1 {
346            height += 1;
347            nodes = nodes
348                .into_par_iter()
349                .chunks(2 * TREE_BASE)
350                .map(|child_nodes| {
351                    let child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }, u8> =
352                        child_nodes.into_iter().collect();
353                    let child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }, u8> = child_trees
354                        .iter()
355                        .map(|child_tree| child_tree.summary().clone())
356                        .collect();
357                    let mut summary = child_summaries[0].clone();
358                    for child_summary in &child_summaries[1..] {
359                        <T::Summary as Summary>::add_summary(&mut summary, child_summary, cx);
360                    }
361                    SumTree(Arc::new(Node::Internal {
362                        height,
363                        summary,
364                        child_summaries,
365                        child_trees,
366                    }))
367                })
368                .collect::<Vec<_>>();
369        }
370
371        if nodes.is_empty() {
372            Self::new(cx)
373        } else {
374            debug_assert_eq!(nodes.len(), 1);
375            nodes.pop().unwrap()
376        }
377    }
378
379    #[allow(unused)]
380    pub fn items<'a>(&'a self, cx: <T::Summary as Summary>::Context<'a>) -> Vec<T> {
381        let mut items = Vec::new();
382        let mut cursor = self.cursor::<()>(cx);
383        cursor.next();
384        while let Some(item) = cursor.item() {
385            items.push(item.clone());
386            cursor.next();
387        }
388        items
389    }
390
391    pub fn iter(&self) -> Iter<'_, T> {
392        Iter::new(self)
393    }
394
395    /// A more efficient version of `Cursor::new()` + `Cursor::seek()` + `Cursor::item()`.
396    ///
397    /// Only returns the item that exactly has the target match.
398    pub fn find_exact<'a, 'slf, D, Target>(
399        &'slf self,
400        cx: <T::Summary as Summary>::Context<'a>,
401        target: &Target,
402        bias: Bias,
403    ) -> (D, D, Option<&'slf T>)
404    where
405        D: Dimension<'slf, T::Summary>,
406        Target: SeekTarget<'slf, T::Summary, D>,
407    {
408        let tree_end = D::zero(cx).with_added_summary(self.summary(), cx);
409        let comparison = target.cmp(&tree_end, cx);
410        if comparison == Ordering::Greater || (comparison == Ordering::Equal && bias == Bias::Right)
411        {
412            return (tree_end.clone(), tree_end, None);
413        }
414
415        let mut pos = D::zero(cx);
416        return match Self::find_iterate::<_, _, true>(cx, target, bias, &mut pos, self) {
417            Some((item, end)) => (pos, end, Some(item)),
418            None => (pos.clone(), pos, None),
419        };
420    }
421
422    /// A more efficient version of `Cursor::new()` + `Cursor::seek()` + `Cursor::item()`
423    pub fn find<'a, 'slf, D, Target>(
424        &'slf self,
425        cx: <T::Summary as Summary>::Context<'a>,
426        target: &Target,
427        bias: Bias,
428    ) -> (D, D, Option<&'slf T>)
429    where
430        D: Dimension<'slf, T::Summary>,
431        Target: SeekTarget<'slf, T::Summary, D>,
432    {
433        let tree_end = D::zero(cx).with_added_summary(self.summary(), cx);
434        let comparison = target.cmp(&tree_end, cx);
435        if comparison == Ordering::Greater || (comparison == Ordering::Equal && bias == Bias::Right)
436        {
437            return (tree_end.clone(), tree_end, None);
438        }
439
440        let mut pos = D::zero(cx);
441        return match Self::find_iterate::<_, _, false>(cx, target, bias, &mut pos, self) {
442            Some((item, end)) => (pos, end, Some(item)),
443            None => (pos.clone(), pos, None),
444        };
445    }
446
447    fn find_iterate<'tree, 'a, D, Target, const EXACT: bool>(
448        cx: <T::Summary as Summary>::Context<'a>,
449        target: &Target,
450        bias: Bias,
451        position: &mut D,
452        mut this: &'tree SumTree<T>,
453    ) -> Option<(&'tree T, D)>
454    where
455        D: Dimension<'tree, T::Summary>,
456        Target: SeekTarget<'tree, T::Summary, D>,
457    {
458        'iterate: loop {
459            match &*this.0 {
460                Node::Internal {
461                    child_summaries,
462                    child_trees,
463                    ..
464                } => {
465                    for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
466                        let child_end = position.clone().with_added_summary(child_summary, cx);
467
468                        let comparison = target.cmp(&child_end, cx);
469                        let target_in_child = comparison == Ordering::Less
470                            || (comparison == Ordering::Equal && bias == Bias::Left);
471                        if target_in_child {
472                            this = child_tree;
473                            continue 'iterate;
474                        }
475                        *position = child_end;
476                    }
477                }
478                Node::Leaf {
479                    items,
480                    item_summaries,
481                    ..
482                } => {
483                    for (item, item_summary) in items.iter().zip(item_summaries) {
484                        let mut child_end = position.clone();
485                        child_end.add_summary(item_summary, cx);
486
487                        let comparison = target.cmp(&child_end, cx);
488                        let entry_found = if EXACT {
489                            comparison == Ordering::Equal
490                        } else {
491                            comparison == Ordering::Less
492                                || (comparison == Ordering::Equal && bias == Bias::Left)
493                        };
494                        if entry_found {
495                            return Some((item, child_end));
496                        }
497
498                        *position = child_end;
499                    }
500                }
501            }
502            return None;
503        }
504    }
505
506    /// A more efficient version of `Cursor::new()` + `Cursor::seek()` + `Cursor::item()`
507    pub fn find_with_prev<'a, 'slf, D, Target>(
508        &'slf self,
509        cx: <T::Summary as Summary>::Context<'a>,
510        target: &Target,
511        bias: Bias,
512    ) -> (D, D, Option<(Option<&'slf T>, &'slf T)>)
513    where
514        D: Dimension<'slf, T::Summary>,
515        Target: SeekTarget<'slf, T::Summary, D>,
516    {
517        let tree_end = D::zero(cx).with_added_summary(self.summary(), cx);
518        let comparison = target.cmp(&tree_end, cx);
519        if comparison == Ordering::Greater || (comparison == Ordering::Equal && bias == Bias::Right)
520        {
521            return (tree_end.clone(), tree_end, None);
522        }
523
524        let mut pos = D::zero(cx);
525        return match Self::find_with_prev_iterate::<_, _, false>(cx, target, bias, &mut pos, self) {
526            Some((prev, item, end)) => (pos, end, Some((prev, item))),
527            None => (pos.clone(), pos, None),
528        };
529    }
530
531    fn find_with_prev_iterate<'tree, 'a, D, Target, const EXACT: bool>(
532        cx: <T::Summary as Summary>::Context<'a>,
533        target: &Target,
534        bias: Bias,
535        position: &mut D,
536        mut this: &'tree SumTree<T>,
537    ) -> Option<(Option<&'tree T>, &'tree T, D)>
538    where
539        D: Dimension<'tree, T::Summary>,
540        Target: SeekTarget<'tree, T::Summary, D>,
541    {
542        let mut prev = None;
543        'iterate: loop {
544            match &*this.0 {
545                Node::Internal {
546                    child_summaries,
547                    child_trees,
548                    ..
549                } => {
550                    for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
551                        let child_end = position.clone().with_added_summary(child_summary, cx);
552
553                        let comparison = target.cmp(&child_end, cx);
554                        let target_in_child = comparison == Ordering::Less
555                            || (comparison == Ordering::Equal && bias == Bias::Left);
556                        if target_in_child {
557                            this = child_tree;
558                            continue 'iterate;
559                        }
560                        prev = child_tree.last();
561                        *position = child_end;
562                    }
563                }
564                Node::Leaf {
565                    items,
566                    item_summaries,
567                    ..
568                } => {
569                    for (item, item_summary) in items.iter().zip(item_summaries) {
570                        let mut child_end = position.clone();
571                        child_end.add_summary(item_summary, cx);
572
573                        let comparison = target.cmp(&child_end, cx);
574                        let entry_found = if EXACT {
575                            comparison == Ordering::Equal
576                        } else {
577                            comparison == Ordering::Less
578                                || (comparison == Ordering::Equal && bias == Bias::Left)
579                        };
580                        if entry_found {
581                            return Some((prev, item, child_end));
582                        }
583
584                        prev = Some(item);
585                        *position = child_end;
586                    }
587                }
588            }
589            return None;
590        }
591    }
592
593    pub fn cursor<'a, 'b, D>(
594        &'a self,
595        cx: <T::Summary as Summary>::Context<'b>,
596    ) -> Cursor<'a, 'b, T, D>
597    where
598        D: Dimension<'a, T::Summary>,
599    {
600        Cursor::new(self, cx)
601    }
602
603    /// Note: If the summary type requires a non `()` context, then the filter cursor
604    /// that is returned cannot be used with Rust's iterators.
605    pub fn filter<'a, 'b, F, U>(
606        &'a self,
607        cx: <T::Summary as Summary>::Context<'b>,
608        filter_node: F,
609    ) -> FilterCursor<'a, 'b, F, T, U>
610    where
611        F: FnMut(&T::Summary) -> bool,
612        U: Dimension<'a, T::Summary>,
613    {
614        FilterCursor::new(self, cx, filter_node)
615    }
616
617    #[allow(dead_code)]
618    pub fn first(&self) -> Option<&T> {
619        self.leftmost_leaf().0.items().first()
620    }
621
622    pub fn last(&self) -> Option<&T> {
623        self.rightmost_leaf().0.items().last()
624    }
625
626    pub fn last_summary(&self) -> Option<&T::Summary> {
627        self.rightmost_leaf().0.child_summaries().last()
628    }
629
630    pub fn update_last(
631        &mut self,
632        f: impl FnOnce(&mut T),
633        cx: <T::Summary as Summary>::Context<'_>,
634    ) {
635        self.update_last_recursive(f, cx);
636    }
637
638    fn update_last_recursive(
639        &mut self,
640        f: impl FnOnce(&mut T),
641        cx: <T::Summary as Summary>::Context<'_>,
642    ) -> Option<T::Summary> {
643        match Arc::make_mut(&mut self.0) {
644            Node::Internal {
645                summary,
646                child_summaries,
647                child_trees,
648                ..
649            } => {
650                let last_summary = child_summaries.last_mut().unwrap();
651                let last_child = child_trees.last_mut().unwrap();
652                *last_summary = last_child.update_last_recursive(f, cx).unwrap();
653                *summary = sum(child_summaries.iter(), cx);
654                Some(summary.clone())
655            }
656            Node::Leaf {
657                summary,
658                items,
659                item_summaries,
660            } => {
661                if let Some((item, item_summary)) = items.last_mut().zip(item_summaries.last_mut())
662                {
663                    (f)(item);
664                    *item_summary = item.summary(cx);
665                    *summary = sum(item_summaries.iter(), cx);
666                    Some(summary.clone())
667                } else {
668                    None
669                }
670            }
671        }
672    }
673
674    pub fn update_first(
675        &mut self,
676        f: impl FnOnce(&mut T),
677        cx: <T::Summary as Summary>::Context<'_>,
678    ) {
679        self.update_first_recursive(f, cx);
680    }
681
682    fn update_first_recursive(
683        &mut self,
684        f: impl FnOnce(&mut T),
685        cx: <T::Summary as Summary>::Context<'_>,
686    ) -> Option<T::Summary> {
687        match Arc::make_mut(&mut self.0) {
688            Node::Internal {
689                summary,
690                child_summaries,
691                child_trees,
692                ..
693            } => {
694                let first_summary = child_summaries.first_mut().unwrap();
695                let first_child = child_trees.first_mut().unwrap();
696                *first_summary = first_child.update_first_recursive(f, cx).unwrap();
697                *summary = sum(child_summaries.iter(), cx);
698                Some(summary.clone())
699            }
700            Node::Leaf {
701                summary,
702                items,
703                item_summaries,
704            } => {
705                if let Some((item, item_summary)) =
706                    items.first_mut().zip(item_summaries.first_mut())
707                {
708                    (f)(item);
709                    *item_summary = item.summary(cx);
710                    *summary = sum(item_summaries.iter(), cx);
711                    Some(summary.clone())
712                } else {
713                    None
714                }
715            }
716        }
717    }
718
719    pub fn extent<'a, D: Dimension<'a, T::Summary>>(
720        &'a self,
721        cx: <T::Summary as Summary>::Context<'_>,
722    ) -> D {
723        let mut extent = D::zero(cx);
724        match self.0.as_ref() {
725            Node::Internal { summary, .. } | Node::Leaf { summary, .. } => {
726                extent.add_summary(summary, cx);
727            }
728        }
729        extent
730    }
731
732    pub fn summary(&self) -> &T::Summary {
733        match self.0.as_ref() {
734            Node::Internal { summary, .. } => summary,
735            Node::Leaf { summary, .. } => summary,
736        }
737    }
738
739    pub fn is_empty(&self) -> bool {
740        match self.0.as_ref() {
741            Node::Internal { .. } => false,
742            Node::Leaf { items, .. } => items.is_empty(),
743        }
744    }
745
746    pub fn extend<I>(&mut self, iter: I, cx: <T::Summary as Summary>::Context<'_>)
747    where
748        I: IntoIterator<Item = T>,
749    {
750        self.append(Self::from_iter(iter, cx), cx);
751    }
752
753    pub fn par_extend<I, Iter>(&mut self, iter: I, cx: <T::Summary as Summary>::Context<'_>)
754    where
755        I: IntoParallelIterator<Iter = Iter>,
756        Iter: IndexedParallelIterator<Item = T>,
757        T: Send + Sync,
758        T::Summary: Send + Sync,
759        for<'a> <T::Summary as Summary>::Context<'a>: Sync,
760    {
761        self.append(Self::from_par_iter(iter, cx), cx);
762    }
763
764    pub fn push(&mut self, item: T, cx: <T::Summary as Summary>::Context<'_>) {
765        let summary = item.summary(cx);
766        self.append(
767            SumTree(Arc::new(Node::Leaf {
768                summary: summary.clone(),
769                items: ArrayVec::from_iter(Some(item)),
770                item_summaries: ArrayVec::from_iter(Some(summary)),
771            })),
772            cx,
773        );
774    }
775
776    pub fn append(&mut self, mut other: Self, cx: <T::Summary as Summary>::Context<'_>) {
777        if self.is_empty() {
778            *self = other;
779        } else if !other.0.is_leaf() || !other.0.items().is_empty() {
780            if self.0.height() < other.0.height() {
781                if let Some(tree) = Self::append_large(self.clone(), &mut other, cx) {
782                    *self = Self::from_child_trees(tree, other, cx);
783                } else {
784                    *self = other;
785                }
786            } else if let Some(split_tree) = self.push_tree_recursive(other, cx) {
787                *self = Self::from_child_trees(self.clone(), split_tree, cx);
788            }
789        }
790    }
791
792    fn push_tree_recursive(
793        &mut self,
794        other: SumTree<T>,
795        cx: <T::Summary as Summary>::Context<'_>,
796    ) -> Option<SumTree<T>> {
797        match Arc::make_mut(&mut self.0) {
798            Node::Internal {
799                height,
800                summary,
801                child_summaries,
802                child_trees,
803                ..
804            } => {
805                let other_node = other.0.clone();
806                <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
807
808                let height_delta = *height - other_node.height();
809                let mut summaries_to_append = ArrayVec::<T::Summary, { 2 * TREE_BASE }, u8>::new();
810                let mut trees_to_append = ArrayVec::<SumTree<T>, { 2 * TREE_BASE }, u8>::new();
811                if height_delta == 0 {
812                    summaries_to_append.extend(other_node.child_summaries().iter().cloned());
813                    trees_to_append.extend(other_node.child_trees().iter().cloned());
814                } else if height_delta == 1 && !other_node.is_underflowing() {
815                    summaries_to_append
816                        .push(other_node.summary().clone())
817                        .unwrap_oob();
818                    trees_to_append.push(other).unwrap_oob();
819                } else {
820                    let tree_to_append = child_trees
821                        .last_mut()
822                        .unwrap()
823                        .push_tree_recursive(other, cx);
824                    *child_summaries.last_mut().unwrap() =
825                        child_trees.last().unwrap().0.summary().clone();
826
827                    if let Some(split_tree) = tree_to_append {
828                        summaries_to_append
829                            .push(split_tree.0.summary().clone())
830                            .unwrap_oob();
831                        trees_to_append.push(split_tree).unwrap_oob();
832                    }
833                }
834
835                let child_count = child_trees.len() + trees_to_append.len();
836                if child_count > 2 * TREE_BASE {
837                    let left_summaries: ArrayVec<_, { 2 * TREE_BASE }, u8>;
838                    let right_summaries: ArrayVec<_, { 2 * TREE_BASE }, u8>;
839                    let left_trees;
840                    let right_trees;
841
842                    let midpoint = (child_count + child_count % 2) / 2;
843                    {
844                        let mut all_summaries = child_summaries
845                            .iter()
846                            .chain(summaries_to_append.iter())
847                            .cloned();
848                        left_summaries = all_summaries.by_ref().take(midpoint).collect();
849                        right_summaries = all_summaries.collect();
850                        let mut all_trees =
851                            child_trees.iter().chain(trees_to_append.iter()).cloned();
852                        left_trees = all_trees.by_ref().take(midpoint).collect();
853                        right_trees = all_trees.collect();
854                    }
855                    *summary = sum(left_summaries.iter(), cx);
856                    *child_summaries = left_summaries;
857                    *child_trees = left_trees;
858
859                    Some(SumTree(Arc::new(Node::Internal {
860                        height: *height,
861                        summary: sum(right_summaries.iter(), cx),
862                        child_summaries: right_summaries,
863                        child_trees: right_trees,
864                    })))
865                } else {
866                    child_summaries.extend(summaries_to_append);
867                    child_trees.extend(trees_to_append);
868                    None
869                }
870            }
871            Node::Leaf {
872                summary,
873                items,
874                item_summaries,
875            } => {
876                let other_node = other.0;
877
878                let child_count = items.len() + other_node.items().len();
879                if child_count > 2 * TREE_BASE {
880                    let left_items;
881                    let right_items;
882                    let left_summaries;
883                    let right_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }, u8>;
884
885                    let midpoint = (child_count + child_count % 2) / 2;
886                    {
887                        let mut all_items = items.iter().chain(other_node.items().iter()).cloned();
888                        left_items = all_items.by_ref().take(midpoint).collect();
889                        right_items = all_items.collect();
890
891                        let mut all_summaries = item_summaries
892                            .iter()
893                            .chain(other_node.child_summaries())
894                            .cloned();
895                        left_summaries = all_summaries.by_ref().take(midpoint).collect();
896                        right_summaries = all_summaries.collect();
897                    }
898                    *items = left_items;
899                    *item_summaries = left_summaries;
900                    *summary = sum(item_summaries.iter(), cx);
901                    Some(SumTree(Arc::new(Node::Leaf {
902                        items: right_items,
903                        summary: sum(right_summaries.iter(), cx),
904                        item_summaries: right_summaries,
905                    })))
906                } else {
907                    <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
908                    items.extend(other_node.items().iter().cloned());
909                    item_summaries.extend(other_node.child_summaries().iter().cloned());
910                    None
911                }
912            }
913        }
914    }
915
916    // appends the `large` tree to a `small` tree, assumes small.height() <= large.height()
917    fn append_large(
918        small: Self,
919        large: &mut Self,
920        cx: <T::Summary as Summary>::Context<'_>,
921    ) -> Option<Self> {
922        if small.0.height() == large.0.height() {
923            if !small.0.is_underflowing() {
924                Some(small)
925            } else {
926                Self::merge_into_right(small, large, cx)
927            }
928        } else {
929            debug_assert!(small.0.height() < large.0.height());
930            let Node::Internal {
931                height,
932                summary,
933                child_summaries,
934                child_trees,
935            } = Arc::make_mut(&mut large.0)
936            else {
937                unreachable!();
938            };
939            let mut full_summary = small.summary().clone();
940            Summary::add_summary(&mut full_summary, summary, cx);
941            *summary = full_summary;
942
943            let first = child_trees.first_mut().unwrap();
944            let res = Self::append_large(small, first, cx);
945            *child_summaries.first_mut().unwrap() = first.summary().clone();
946            if let Some(tree) = res {
947                if child_trees.len() < 2 * TREE_BASE {
948                    child_summaries
949                        .insert(0, tree.summary().clone())
950                        .unwrap_oob();
951                    child_trees.insert(0, tree).unwrap_oob();
952                    None
953                } else {
954                    let new_child_summaries = {
955                        let mut res = ArrayVec::from_iter([tree.summary().clone()]);
956                        res.extend(child_summaries.drain(..TREE_BASE));
957                        res
958                    };
959                    let tree = SumTree(Arc::new(Node::Internal {
960                        height: *height,
961                        summary: sum(new_child_summaries.iter(), cx),
962                        child_summaries: new_child_summaries,
963                        child_trees: {
964                            let mut res = ArrayVec::from_iter([tree]);
965                            res.extend(child_trees.drain(..TREE_BASE));
966                            res
967                        },
968                    }));
969
970                    *summary = sum(child_summaries.iter(), cx);
971                    Some(tree)
972                }
973            } else {
974                None
975            }
976        }
977    }
978
979    // Merge two nodes into `large`.
980    //
981    // `large` will contain the contents of `small` followed by its own data.
982    // If the combined data exceed the node capacity, returns a new node that
983    // holds the first half of the merged items and `large` is left with the
984    // second half
985    //
986    // The nodes must be on the same height
987    // It only makes sense to call this when `small` is underflowing
988    fn merge_into_right(
989        small: Self,
990        large: &mut Self,
991        cx: <<T as Item>::Summary as Summary>::Context<'_>,
992    ) -> Option<SumTree<T>> {
993        debug_assert_eq!(small.0.height(), large.0.height());
994        match (small.0.as_ref(), Arc::make_mut(&mut large.0)) {
995            (
996                Node::Internal {
997                    summary: small_summary,
998                    child_summaries: small_child_summaries,
999                    child_trees: small_child_trees,
1000                    ..
1001                },
1002                Node::Internal {
1003                    summary,
1004                    child_summaries,
1005                    child_trees,
1006                    height,
1007                },
1008            ) => {
1009                let total_child_count = child_trees.len() + small_child_trees.len();
1010                if total_child_count <= 2 * TREE_BASE {
1011                    let mut all_trees = small_child_trees.clone();
1012                    all_trees.extend(child_trees.drain(..));
1013                    *child_trees = all_trees;
1014
1015                    let mut all_summaries = small_child_summaries.clone();
1016                    all_summaries.extend(child_summaries.drain(..));
1017                    *child_summaries = all_summaries;
1018
1019                    let mut full_summary = small_summary.clone();
1020                    Summary::add_summary(&mut full_summary, summary, cx);
1021                    *summary = full_summary;
1022                    None
1023                } else {
1024                    let midpoint = total_child_count.div_ceil(2);
1025                    let mut all_trees = small_child_trees.iter().chain(child_trees.iter()).cloned();
1026                    let left_trees = all_trees.by_ref().take(midpoint).collect();
1027                    *child_trees = all_trees.collect();
1028
1029                    let mut all_summaries = small_child_summaries
1030                        .iter()
1031                        .chain(child_summaries.iter())
1032                        .cloned();
1033                    let left_summaries: ArrayVec<_, { 2 * TREE_BASE }, u8> =
1034                        all_summaries.by_ref().take(midpoint).collect();
1035                    *child_summaries = all_summaries.collect();
1036
1037                    *summary = sum(child_summaries.iter(), cx);
1038                    Some(SumTree(Arc::new(Node::Internal {
1039                        height: *height,
1040                        summary: sum(left_summaries.iter(), cx),
1041                        child_summaries: left_summaries,
1042                        child_trees: left_trees,
1043                    })))
1044                }
1045            }
1046            (
1047                Node::Leaf {
1048                    summary: small_summary,
1049                    items: small_items,
1050                    item_summaries: small_item_summaries,
1051                },
1052                Node::Leaf {
1053                    summary,
1054                    items,
1055                    item_summaries,
1056                },
1057            ) => {
1058                let total_child_count = small_items.len() + items.len();
1059                if total_child_count <= 2 * TREE_BASE {
1060                    let mut all_items = small_items.clone();
1061                    all_items.extend(items.drain(..));
1062                    *items = all_items;
1063
1064                    let mut all_summaries = small_item_summaries.clone();
1065                    all_summaries.extend(item_summaries.drain(..));
1066                    *item_summaries = all_summaries;
1067
1068                    let mut full_summary = small_summary.clone();
1069                    Summary::add_summary(&mut full_summary, summary, cx);
1070                    *summary = full_summary;
1071                    None
1072                } else {
1073                    let midpoint = total_child_count.div_ceil(2);
1074                    let mut all_items = small_items.iter().chain(items.iter()).cloned();
1075                    let left_items = all_items.by_ref().take(midpoint).collect();
1076                    *items = all_items.collect();
1077
1078                    let mut all_summaries = small_item_summaries
1079                        .iter()
1080                        .chain(item_summaries.iter())
1081                        .cloned();
1082                    let left_summaries: ArrayVec<_, { 2 * TREE_BASE }, u8> =
1083                        all_summaries.by_ref().take(midpoint).collect();
1084                    *item_summaries = all_summaries.collect();
1085
1086                    *summary = sum(item_summaries.iter(), cx);
1087                    Some(SumTree(Arc::new(Node::Leaf {
1088                        items: left_items,
1089                        summary: sum(left_summaries.iter(), cx),
1090                        item_summaries: left_summaries,
1091                    })))
1092                }
1093            }
1094            _ => unreachable!(),
1095        }
1096    }
1097
1098    fn from_child_trees(
1099        left: SumTree<T>,
1100        right: SumTree<T>,
1101        cx: <T::Summary as Summary>::Context<'_>,
1102    ) -> Self {
1103        let height = left.0.height() + 1;
1104        let mut child_summaries = ArrayVec::new();
1105        child_summaries.push(left.0.summary().clone()).unwrap_oob();
1106        child_summaries.push(right.0.summary().clone()).unwrap_oob();
1107        let mut child_trees = ArrayVec::new();
1108        child_trees.push(left).unwrap_oob();
1109        child_trees.push(right).unwrap_oob();
1110        SumTree(Arc::new(Node::Internal {
1111            height,
1112            summary: sum(child_summaries.iter(), cx),
1113            child_summaries,
1114            child_trees,
1115        }))
1116    }
1117
1118    fn leftmost_leaf(&self) -> &Self {
1119        match *self.0 {
1120            Node::Leaf { .. } => self,
1121            Node::Internal {
1122                ref child_trees, ..
1123            } => child_trees.first().unwrap().leftmost_leaf(),
1124        }
1125    }
1126
1127    fn rightmost_leaf(&self) -> &Self {
1128        match *self.0 {
1129            Node::Leaf { .. } => self,
1130            Node::Internal {
1131                ref child_trees, ..
1132            } => child_trees.last().unwrap().rightmost_leaf(),
1133        }
1134    }
1135}
1136
1137impl<T: Item + PartialEq> PartialEq for SumTree<T> {
1138    fn eq(&self, other: &Self) -> bool {
1139        self.iter().eq(other.iter())
1140    }
1141}
1142
1143impl<T: Item + Eq> Eq for SumTree<T> {}
1144
1145impl<T: KeyedItem> SumTree<T> {
1146    pub fn insert_or_replace<'a, 'b>(
1147        &'a mut self,
1148        item: T,
1149        cx: <T::Summary as Summary>::Context<'b>,
1150    ) -> Option<T> {
1151        let mut replaced = None;
1152        {
1153            let mut cursor = self.cursor::<T::Key>(cx);
1154            let mut new_tree = cursor.slice(&item.key(), Bias::Left);
1155            if let Some(cursor_item) = cursor.item()
1156                && cursor_item.key() == item.key()
1157            {
1158                replaced = Some(cursor_item.clone());
1159                cursor.next();
1160            }
1161            new_tree.push(item, cx);
1162            new_tree.append(cursor.suffix(), cx);
1163            drop(cursor);
1164            *self = new_tree
1165        };
1166        replaced
1167    }
1168
1169    pub fn remove(&mut self, key: &T::Key, cx: <T::Summary as Summary>::Context<'_>) -> Option<T> {
1170        let mut removed = None;
1171        *self = {
1172            let mut cursor = self.cursor::<T::Key>(cx);
1173            let mut new_tree = cursor.slice(key, Bias::Left);
1174            if let Some(item) = cursor.item()
1175                && item.key() == *key
1176            {
1177                removed = Some(item.clone());
1178                cursor.next();
1179            }
1180            new_tree.append(cursor.suffix(), cx);
1181            new_tree
1182        };
1183        removed
1184    }
1185
1186    pub fn edit(
1187        &mut self,
1188        mut edits: Vec<Edit<T>>,
1189        cx: <T::Summary as Summary>::Context<'_>,
1190    ) -> Vec<T> {
1191        if edits.is_empty() {
1192            return Vec::new();
1193        }
1194
1195        let mut removed = Vec::new();
1196        edits.sort_unstable_by_key(|item| item.key());
1197
1198        *self = {
1199            let mut cursor = self.cursor::<T::Key>(cx);
1200            let mut new_tree = SumTree::new(cx);
1201            let mut buffered_items = Vec::new();
1202
1203            cursor.seek(&T::Key::zero(cx), Bias::Left);
1204            for edit in edits {
1205                let new_key = edit.key();
1206                let mut old_item = cursor.item();
1207
1208                if old_item
1209                    .as_ref()
1210                    .is_some_and(|old_item| old_item.key() < new_key)
1211                {
1212                    new_tree.extend(buffered_items.drain(..), cx);
1213                    let slice = cursor.slice(&new_key, Bias::Left);
1214                    new_tree.append(slice, cx);
1215                    old_item = cursor.item();
1216                }
1217
1218                if let Some(old_item) = old_item
1219                    && old_item.key() == new_key
1220                {
1221                    removed.push(old_item.clone());
1222                    cursor.next();
1223                }
1224
1225                match edit {
1226                    Edit::Insert(item) => {
1227                        buffered_items.push(item);
1228                    }
1229                    Edit::Remove(_) => {}
1230                }
1231            }
1232
1233            new_tree.extend(buffered_items, cx);
1234            new_tree.append(cursor.suffix(), cx);
1235            new_tree
1236        };
1237
1238        removed
1239    }
1240
1241    pub fn get<'a>(
1242        &'a self,
1243        key: &T::Key,
1244        cx: <T::Summary as Summary>::Context<'a>,
1245    ) -> Option<&'a T> {
1246        if let (_, _, Some(item)) = self.find_exact::<T::Key, _>(cx, key, Bias::Left) {
1247            Some(item)
1248        } else {
1249            None
1250        }
1251    }
1252}
1253
1254impl<T, S> Default for SumTree<T>
1255where
1256    T: Item<Summary = S>,
1257    S: for<'a> Summary<Context<'a> = ()>,
1258{
1259    fn default() -> Self {
1260        Self::new(())
1261    }
1262}
1263
1264#[derive(Clone)]
1265pub enum Node<T: Item> {
1266    Internal {
1267        height: u8,
1268        summary: T::Summary,
1269        child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }, u8>,
1270        child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }, u8>,
1271    },
1272    Leaf {
1273        summary: T::Summary,
1274        items: ArrayVec<T, { 2 * TREE_BASE }, u8>,
1275        item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }, u8>,
1276    },
1277}
1278
1279impl<T> fmt::Debug for Node<T>
1280where
1281    T: Item + fmt::Debug,
1282    T::Summary: fmt::Debug,
1283{
1284    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1285        match self {
1286            Node::Internal {
1287                height,
1288                summary,
1289                child_summaries,
1290                child_trees,
1291            } => f
1292                .debug_struct("Internal")
1293                .field("height", height)
1294                .field("summary", summary)
1295                .field("child_summaries", child_summaries)
1296                .field("child_trees", child_trees)
1297                .finish(),
1298            Node::Leaf {
1299                summary,
1300                items,
1301                item_summaries,
1302            } => f
1303                .debug_struct("Leaf")
1304                .field("summary", summary)
1305                .field("items", items)
1306                .field("item_summaries", item_summaries)
1307                .finish(),
1308        }
1309    }
1310}
1311
1312impl<T: Item> Node<T> {
1313    fn is_leaf(&self) -> bool {
1314        matches!(self, Node::Leaf { .. })
1315    }
1316
1317    fn height(&self) -> u8 {
1318        match self {
1319            Node::Internal { height, .. } => *height,
1320            Node::Leaf { .. } => 0,
1321        }
1322    }
1323
1324    fn summary(&self) -> &T::Summary {
1325        match self {
1326            Node::Internal { summary, .. } => summary,
1327            Node::Leaf { summary, .. } => summary,
1328        }
1329    }
1330
1331    fn child_summaries(&self) -> &[T::Summary] {
1332        match self {
1333            Node::Internal {
1334                child_summaries, ..
1335            } => child_summaries.as_slice(),
1336            Node::Leaf { item_summaries, .. } => item_summaries.as_slice(),
1337        }
1338    }
1339
1340    fn child_trees(&self) -> &ArrayVec<SumTree<T>, { 2 * TREE_BASE }, u8> {
1341        match self {
1342            Node::Internal { child_trees, .. } => child_trees,
1343            Node::Leaf { .. } => panic!("Leaf nodes have no child trees"),
1344        }
1345    }
1346
1347    fn items(&self) -> &ArrayVec<T, { 2 * TREE_BASE }, u8> {
1348        match self {
1349            Node::Leaf { items, .. } => items,
1350            Node::Internal { .. } => panic!("Internal nodes have no items"),
1351        }
1352    }
1353
1354    fn is_underflowing(&self) -> bool {
1355        match self {
1356            Node::Internal { child_trees, .. } => child_trees.len() < TREE_BASE,
1357            Node::Leaf { items, .. } => items.len() < TREE_BASE,
1358        }
1359    }
1360}
1361
1362#[derive(Debug)]
1363pub enum Edit<T: KeyedItem> {
1364    Insert(T),
1365    Remove(T::Key),
1366}
1367
1368impl<T: KeyedItem> Edit<T> {
1369    fn key(&self) -> T::Key {
1370        match self {
1371            Edit::Insert(item) => item.key(),
1372            Edit::Remove(key) => key.clone(),
1373        }
1374    }
1375}
1376
1377fn sum<'a, T, I>(iter: I, cx: T::Context<'_>) -> T
1378where
1379    T: 'a + Summary,
1380    I: Iterator<Item = &'a T>,
1381{
1382    let mut sum = T::zero(cx);
1383    for value in iter {
1384        sum.add_summary(value, cx);
1385    }
1386    sum
1387}
1388
1389#[cfg(test)]
1390mod tests {
1391    use super::*;
1392    use rand::{distr::StandardUniform, prelude::*};
1393    use std::cmp;
1394
1395    #[ctor::ctor]
1396    fn init_logger() {
1397        zlog::init_test();
1398    }
1399
1400    #[test]
1401    fn test_extend_and_push_tree() {
1402        let mut tree1 = SumTree::default();
1403        tree1.extend(0..20, ());
1404
1405        let mut tree2 = SumTree::default();
1406        tree2.extend(50..100, ());
1407
1408        tree1.append(tree2, ());
1409        assert_eq!(tree1.items(()), (0..20).chain(50..100).collect::<Vec<u8>>());
1410    }
1411
1412    #[test]
1413    fn test_random() {
1414        let mut starting_seed = 0;
1415        if let Ok(value) = std::env::var("SEED") {
1416            starting_seed = value.parse().expect("invalid SEED variable");
1417        }
1418        let mut num_iterations = 100;
1419        if let Ok(value) = std::env::var("ITERATIONS") {
1420            num_iterations = value.parse().expect("invalid ITERATIONS variable");
1421        }
1422        let num_operations = std::env::var("OPERATIONS")
1423            .map_or(5, |o| o.parse().expect("invalid OPERATIONS variable"));
1424
1425        for seed in starting_seed..(starting_seed + num_iterations) {
1426            eprintln!("seed = {}", seed);
1427            let mut rng = StdRng::seed_from_u64(seed);
1428
1429            let rng = &mut rng;
1430            let mut tree = SumTree::<u8>::default();
1431            let count = rng.random_range(0..10);
1432            if rng.random() {
1433                tree.extend(rng.sample_iter(StandardUniform).take(count), ());
1434            } else {
1435                let items = rng
1436                    .sample_iter(StandardUniform)
1437                    .take(count)
1438                    .collect::<Vec<_>>();
1439                tree.par_extend(items, ());
1440            }
1441
1442            for _ in 0..num_operations {
1443                let splice_end = rng.random_range(0..tree.extent::<Count>(()).0 + 1);
1444                let splice_start = rng.random_range(0..splice_end + 1);
1445                let count = rng.random_range(0..10);
1446                let tree_end = tree.extent::<Count>(());
1447                let new_items = rng
1448                    .sample_iter(StandardUniform)
1449                    .take(count)
1450                    .collect::<Vec<u8>>();
1451
1452                let mut reference_items = tree.items(());
1453                reference_items.splice(splice_start..splice_end, new_items.clone());
1454
1455                tree = {
1456                    let mut cursor = tree.cursor::<Count>(());
1457                    let mut new_tree = cursor.slice(&Count(splice_start), Bias::Right);
1458                    if rng.random() {
1459                        new_tree.extend(new_items, ());
1460                    } else {
1461                        new_tree.par_extend(new_items, ());
1462                    }
1463                    cursor.seek(&Count(splice_end), Bias::Right);
1464                    new_tree.append(cursor.slice(&tree_end, Bias::Right), ());
1465                    new_tree
1466                };
1467
1468                assert_eq!(tree.items(()), reference_items);
1469                assert_eq!(
1470                    tree.iter().collect::<Vec<_>>(),
1471                    tree.cursor::<()>(()).collect::<Vec<_>>()
1472                );
1473
1474                log::info!("tree items: {:?}", tree.items(()));
1475
1476                let mut filter_cursor =
1477                    tree.filter::<_, Count>((), |summary| summary.contains_even);
1478                let expected_filtered_items = tree
1479                    .items(())
1480                    .into_iter()
1481                    .enumerate()
1482                    .filter(|(_, item)| (item & 1) == 0)
1483                    .collect::<Vec<_>>();
1484
1485                let mut item_ix = if rng.random() {
1486                    filter_cursor.next();
1487                    0
1488                } else {
1489                    filter_cursor.prev();
1490                    expected_filtered_items.len().saturating_sub(1)
1491                };
1492                while item_ix < expected_filtered_items.len() {
1493                    log::info!("filter_cursor, item_ix: {}", item_ix);
1494                    let actual_item = filter_cursor.item().unwrap();
1495                    let (reference_index, reference_item) = expected_filtered_items[item_ix];
1496                    assert_eq!(actual_item, &reference_item);
1497                    assert_eq!(filter_cursor.start().0, reference_index);
1498                    log::info!("next");
1499                    filter_cursor.next();
1500                    item_ix += 1;
1501
1502                    while item_ix > 0 && rng.random_bool(0.2) {
1503                        log::info!("prev");
1504                        filter_cursor.prev();
1505                        item_ix -= 1;
1506
1507                        if item_ix == 0 && rng.random_bool(0.2) {
1508                            filter_cursor.prev();
1509                            assert_eq!(filter_cursor.item(), None);
1510                            assert_eq!(filter_cursor.start().0, 0);
1511                            filter_cursor.next();
1512                        }
1513                    }
1514                }
1515                assert_eq!(filter_cursor.item(), None);
1516
1517                let mut before_start = false;
1518                let mut cursor = tree.cursor::<Count>(());
1519                let start_pos = rng.random_range(0..=reference_items.len());
1520                cursor.seek(&Count(start_pos), Bias::Right);
1521                let mut pos = rng.random_range(start_pos..=reference_items.len());
1522                cursor.seek_forward(&Count(pos), Bias::Right);
1523
1524                for i in 0..10 {
1525                    assert_eq!(cursor.start().0, pos);
1526
1527                    if pos > 0 {
1528                        assert_eq!(cursor.prev_item().unwrap(), &reference_items[pos - 1]);
1529                    } else {
1530                        assert_eq!(cursor.prev_item(), None);
1531                    }
1532
1533                    if pos < reference_items.len() && !before_start {
1534                        assert_eq!(cursor.item().unwrap(), &reference_items[pos]);
1535                    } else {
1536                        assert_eq!(cursor.item(), None);
1537                    }
1538
1539                    if before_start {
1540                        assert_eq!(cursor.next_item(), reference_items.first());
1541                    } else if pos + 1 < reference_items.len() {
1542                        assert_eq!(cursor.next_item().unwrap(), &reference_items[pos + 1]);
1543                    } else {
1544                        assert_eq!(cursor.next_item(), None);
1545                    }
1546
1547                    if i < 5 {
1548                        cursor.next();
1549                        if pos < reference_items.len() {
1550                            pos += 1;
1551                            before_start = false;
1552                        }
1553                    } else {
1554                        cursor.prev();
1555                        if pos == 0 {
1556                            before_start = true;
1557                        }
1558                        pos = pos.saturating_sub(1);
1559                    }
1560                }
1561            }
1562
1563            for _ in 0..10 {
1564                let end = rng.random_range(0..tree.extent::<Count>(()).0 + 1);
1565                let start = rng.random_range(0..end + 1);
1566                let start_bias = if rng.random() {
1567                    Bias::Left
1568                } else {
1569                    Bias::Right
1570                };
1571                let end_bias = if rng.random() {
1572                    Bias::Left
1573                } else {
1574                    Bias::Right
1575                };
1576
1577                let mut cursor = tree.cursor::<Count>(());
1578                cursor.seek(&Count(start), start_bias);
1579                let slice = cursor.slice(&Count(end), end_bias);
1580
1581                cursor.seek(&Count(start), start_bias);
1582                let summary = cursor.summary::<_, Sum>(&Count(end), end_bias);
1583
1584                assert_eq!(summary.0, slice.summary().sum);
1585            }
1586        }
1587    }
1588
1589    #[test]
1590    fn test_cursor() {
1591        // Empty tree
1592        let tree = SumTree::<u8>::default();
1593        let mut cursor = tree.cursor::<IntegersSummary>(());
1594        assert_eq!(
1595            cursor.slice(&Count(0), Bias::Right).items(()),
1596            Vec::<u8>::new()
1597        );
1598        assert_eq!(cursor.item(), None);
1599        assert_eq!(cursor.prev_item(), None);
1600        assert_eq!(cursor.next_item(), None);
1601        assert_eq!(cursor.start().sum, 0);
1602        cursor.prev();
1603        assert_eq!(cursor.item(), None);
1604        assert_eq!(cursor.prev_item(), None);
1605        assert_eq!(cursor.next_item(), None);
1606        assert_eq!(cursor.start().sum, 0);
1607        cursor.next();
1608        assert_eq!(cursor.item(), None);
1609        assert_eq!(cursor.prev_item(), None);
1610        assert_eq!(cursor.next_item(), None);
1611        assert_eq!(cursor.start().sum, 0);
1612
1613        // Single-element tree
1614        let mut tree = SumTree::<u8>::default();
1615        tree.extend(vec![1], ());
1616        let mut cursor = tree.cursor::<IntegersSummary>(());
1617        assert_eq!(
1618            cursor.slice(&Count(0), Bias::Right).items(()),
1619            Vec::<u8>::new()
1620        );
1621        assert_eq!(cursor.item(), Some(&1));
1622        assert_eq!(cursor.prev_item(), None);
1623        assert_eq!(cursor.next_item(), None);
1624        assert_eq!(cursor.start().sum, 0);
1625
1626        cursor.next();
1627        assert_eq!(cursor.item(), None);
1628        assert_eq!(cursor.prev_item(), Some(&1));
1629        assert_eq!(cursor.next_item(), None);
1630        assert_eq!(cursor.start().sum, 1);
1631
1632        cursor.prev();
1633        assert_eq!(cursor.item(), Some(&1));
1634        assert_eq!(cursor.prev_item(), None);
1635        assert_eq!(cursor.next_item(), None);
1636        assert_eq!(cursor.start().sum, 0);
1637
1638        let mut cursor = tree.cursor::<IntegersSummary>(());
1639        assert_eq!(cursor.slice(&Count(1), Bias::Right).items(()), [1]);
1640        assert_eq!(cursor.item(), None);
1641        assert_eq!(cursor.prev_item(), Some(&1));
1642        assert_eq!(cursor.next_item(), None);
1643        assert_eq!(cursor.start().sum, 1);
1644
1645        cursor.seek(&Count(0), Bias::Right);
1646        assert_eq!(
1647            cursor
1648                .slice(&tree.extent::<Count>(()), Bias::Right)
1649                .items(()),
1650            [1]
1651        );
1652        assert_eq!(cursor.item(), None);
1653        assert_eq!(cursor.prev_item(), Some(&1));
1654        assert_eq!(cursor.next_item(), None);
1655        assert_eq!(cursor.start().sum, 1);
1656
1657        // Multiple-element tree
1658        let mut tree = SumTree::default();
1659        tree.extend(vec![1, 2, 3, 4, 5, 6], ());
1660        let mut cursor = tree.cursor::<IntegersSummary>(());
1661
1662        assert_eq!(cursor.slice(&Count(2), Bias::Right).items(()), [1, 2]);
1663        assert_eq!(cursor.item(), Some(&3));
1664        assert_eq!(cursor.prev_item(), Some(&2));
1665        assert_eq!(cursor.next_item(), Some(&4));
1666        assert_eq!(cursor.start().sum, 3);
1667
1668        cursor.next();
1669        assert_eq!(cursor.item(), Some(&4));
1670        assert_eq!(cursor.prev_item(), Some(&3));
1671        assert_eq!(cursor.next_item(), Some(&5));
1672        assert_eq!(cursor.start().sum, 6);
1673
1674        cursor.next();
1675        assert_eq!(cursor.item(), Some(&5));
1676        assert_eq!(cursor.prev_item(), Some(&4));
1677        assert_eq!(cursor.next_item(), Some(&6));
1678        assert_eq!(cursor.start().sum, 10);
1679
1680        cursor.next();
1681        assert_eq!(cursor.item(), Some(&6));
1682        assert_eq!(cursor.prev_item(), Some(&5));
1683        assert_eq!(cursor.next_item(), None);
1684        assert_eq!(cursor.start().sum, 15);
1685
1686        cursor.next();
1687        cursor.next();
1688        assert_eq!(cursor.item(), None);
1689        assert_eq!(cursor.prev_item(), Some(&6));
1690        assert_eq!(cursor.next_item(), None);
1691        assert_eq!(cursor.start().sum, 21);
1692
1693        cursor.prev();
1694        assert_eq!(cursor.item(), Some(&6));
1695        assert_eq!(cursor.prev_item(), Some(&5));
1696        assert_eq!(cursor.next_item(), None);
1697        assert_eq!(cursor.start().sum, 15);
1698
1699        cursor.prev();
1700        assert_eq!(cursor.item(), Some(&5));
1701        assert_eq!(cursor.prev_item(), Some(&4));
1702        assert_eq!(cursor.next_item(), Some(&6));
1703        assert_eq!(cursor.start().sum, 10);
1704
1705        cursor.prev();
1706        assert_eq!(cursor.item(), Some(&4));
1707        assert_eq!(cursor.prev_item(), Some(&3));
1708        assert_eq!(cursor.next_item(), Some(&5));
1709        assert_eq!(cursor.start().sum, 6);
1710
1711        cursor.prev();
1712        assert_eq!(cursor.item(), Some(&3));
1713        assert_eq!(cursor.prev_item(), Some(&2));
1714        assert_eq!(cursor.next_item(), Some(&4));
1715        assert_eq!(cursor.start().sum, 3);
1716
1717        cursor.prev();
1718        assert_eq!(cursor.item(), Some(&2));
1719        assert_eq!(cursor.prev_item(), Some(&1));
1720        assert_eq!(cursor.next_item(), Some(&3));
1721        assert_eq!(cursor.start().sum, 1);
1722
1723        cursor.prev();
1724        assert_eq!(cursor.item(), Some(&1));
1725        assert_eq!(cursor.prev_item(), None);
1726        assert_eq!(cursor.next_item(), Some(&2));
1727        assert_eq!(cursor.start().sum, 0);
1728
1729        cursor.prev();
1730        assert_eq!(cursor.item(), None);
1731        assert_eq!(cursor.prev_item(), None);
1732        assert_eq!(cursor.next_item(), Some(&1));
1733        assert_eq!(cursor.start().sum, 0);
1734
1735        cursor.next();
1736        assert_eq!(cursor.item(), Some(&1));
1737        assert_eq!(cursor.prev_item(), None);
1738        assert_eq!(cursor.next_item(), Some(&2));
1739        assert_eq!(cursor.start().sum, 0);
1740
1741        let mut cursor = tree.cursor::<IntegersSummary>(());
1742        assert_eq!(
1743            cursor
1744                .slice(&tree.extent::<Count>(()), Bias::Right)
1745                .items(()),
1746            tree.items(())
1747        );
1748        assert_eq!(cursor.item(), None);
1749        assert_eq!(cursor.prev_item(), Some(&6));
1750        assert_eq!(cursor.next_item(), None);
1751        assert_eq!(cursor.start().sum, 21);
1752
1753        cursor.seek(&Count(3), Bias::Right);
1754        assert_eq!(
1755            cursor
1756                .slice(&tree.extent::<Count>(()), Bias::Right)
1757                .items(()),
1758            [4, 5, 6]
1759        );
1760        assert_eq!(cursor.item(), None);
1761        assert_eq!(cursor.prev_item(), Some(&6));
1762        assert_eq!(cursor.next_item(), None);
1763        assert_eq!(cursor.start().sum, 21);
1764
1765        // Seeking can bias left or right
1766        cursor.seek(&Count(1), Bias::Left);
1767        assert_eq!(cursor.item(), Some(&1));
1768        cursor.seek(&Count(1), Bias::Right);
1769        assert_eq!(cursor.item(), Some(&2));
1770
1771        // Slicing without resetting starts from where the cursor is parked at.
1772        cursor.seek(&Count(1), Bias::Right);
1773        assert_eq!(cursor.slice(&Count(3), Bias::Right).items(()), vec![2, 3]);
1774        assert_eq!(cursor.slice(&Count(6), Bias::Left).items(()), vec![4, 5]);
1775        assert_eq!(cursor.slice(&Count(6), Bias::Right).items(()), vec![6]);
1776    }
1777
1778    #[test]
1779    fn test_edit() {
1780        let mut tree = SumTree::<u8>::default();
1781
1782        let removed = tree.edit(vec![Edit::Insert(1), Edit::Insert(2), Edit::Insert(0)], ());
1783        assert_eq!(tree.items(()), vec![0, 1, 2]);
1784        assert_eq!(removed, Vec::<u8>::new());
1785        assert_eq!(tree.get(&0, ()), Some(&0));
1786        assert_eq!(tree.get(&1, ()), Some(&1));
1787        assert_eq!(tree.get(&2, ()), Some(&2));
1788        assert_eq!(tree.get(&4, ()), None);
1789
1790        let removed = tree.edit(vec![Edit::Insert(2), Edit::Insert(4), Edit::Remove(0)], ());
1791        assert_eq!(tree.items(()), vec![1, 2, 4]);
1792        assert_eq!(removed, vec![0, 2]);
1793        assert_eq!(tree.get(&0, ()), None);
1794        assert_eq!(tree.get(&1, ()), Some(&1));
1795        assert_eq!(tree.get(&2, ()), Some(&2));
1796        assert_eq!(tree.get(&4, ()), Some(&4));
1797    }
1798
1799    #[test]
1800    fn test_from_iter() {
1801        assert_eq!(
1802            SumTree::from_iter(0..100, ()).items(()),
1803            (0..100).collect::<Vec<_>>()
1804        );
1805
1806        // Ensure `from_iter` works correctly when the given iterator restarts
1807        // after calling `next` if `None` was already returned.
1808        let mut ix = 0;
1809        let iterator = std::iter::from_fn(|| {
1810            ix = (ix + 1) % 2;
1811            if ix == 1 { Some(1) } else { None }
1812        });
1813        assert_eq!(SumTree::from_iter(iterator, ()).items(()), vec![1]);
1814    }
1815
1816    #[derive(Clone, Default, Debug)]
1817    pub struct IntegersSummary {
1818        count: usize,
1819        sum: usize,
1820        contains_even: bool,
1821        max: u8,
1822    }
1823
1824    #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1825    struct Count(usize);
1826
1827    #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1828    struct Sum(usize);
1829
1830    impl Item for u8 {
1831        type Summary = IntegersSummary;
1832
1833        fn summary(&self, _cx: ()) -> Self::Summary {
1834            IntegersSummary {
1835                count: 1,
1836                sum: *self as usize,
1837                contains_even: (*self & 1) == 0,
1838                max: *self,
1839            }
1840        }
1841    }
1842
1843    impl KeyedItem for u8 {
1844        type Key = u8;
1845
1846        fn key(&self) -> Self::Key {
1847            *self
1848        }
1849    }
1850
1851    impl ContextLessSummary for IntegersSummary {
1852        fn zero() -> Self {
1853            Default::default()
1854        }
1855
1856        fn add_summary(&mut self, other: &Self) {
1857            self.count += other.count;
1858            self.sum += other.sum;
1859            self.contains_even |= other.contains_even;
1860            self.max = cmp::max(self.max, other.max);
1861        }
1862    }
1863
1864    impl Dimension<'_, IntegersSummary> for u8 {
1865        fn zero(_cx: ()) -> Self {
1866            Default::default()
1867        }
1868
1869        fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1870            *self = summary.max;
1871        }
1872    }
1873
1874    impl Dimension<'_, IntegersSummary> for Count {
1875        fn zero(_cx: ()) -> Self {
1876            Default::default()
1877        }
1878
1879        fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1880            self.0 += summary.count;
1881        }
1882    }
1883
1884    impl SeekTarget<'_, IntegersSummary, IntegersSummary> for Count {
1885        fn cmp(&self, cursor_location: &IntegersSummary, _: ()) -> Ordering {
1886            self.0.cmp(&cursor_location.count)
1887        }
1888    }
1889
1890    impl Dimension<'_, IntegersSummary> for Sum {
1891        fn zero(_cx: ()) -> Self {
1892            Default::default()
1893        }
1894
1895        fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1896            self.0 += summary.sum;
1897        }
1898    }
1899}
sum_tree/sum_tree.rs

sum_tree/
sum_tree.rs
