sum_tree.rs

   1mod cursor;
   2mod tree_map;
   3
   4use arrayvec::ArrayVec;
   5pub use cursor::{Cursor, FilterCursor, Iter};
   6use rayon::prelude::*;
   7use std::marker::PhantomData;
   8use std::mem;
   9use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
  10pub use tree_map::{MapSeekTarget, TreeMap, TreeSet};
  11
  12#[cfg(test)]
  13pub const TREE_BASE: usize = 2;
  14#[cfg(not(test))]
  15pub const TREE_BASE: usize = 6;
  16
  17/// An item that can be stored in a [`SumTree`]
  18///
  19/// Must be summarized by a type that implements [`Summary`]
  20pub trait Item: Clone {
  21    type Summary: Summary;
  22
  23    fn summary(&self, cx: <Self::Summary as Summary>::Context<'_>) -> Self::Summary;
  24}
  25
  26/// An [`Item`] whose summary has a specific key that can be used to identify it
  27pub trait KeyedItem: Item {
  28    type Key: for<'a> Dimension<'a, Self::Summary> + Ord;
  29
  30    fn key(&self) -> Self::Key;
  31}
  32
  33/// A type that describes the Sum of all [`Item`]s in a subtree of the [`SumTree`]
  34///
  35/// Each Summary type can have multiple [`Dimension`]s that it measures,
  36/// which can be used to navigate the tree
  37pub trait Summary: Clone {
  38    type Context<'a>: Copy;
  39    fn zero<'a>(cx: Self::Context<'a>) -> Self;
  40    fn add_summary<'a>(&mut self, summary: &Self, cx: Self::Context<'a>);
  41}
  42
  43pub trait ContextLessSummary: Clone {
  44    fn zero() -> Self;
  45    fn add_summary(&mut self, summary: &Self);
  46}
  47
  48impl<T: ContextLessSummary> Summary for T {
  49    type Context<'a> = ();
  50
  51    fn zero<'a>((): ()) -> Self {
  52        T::zero()
  53    }
  54
  55    fn add_summary<'a>(&mut self, summary: &Self, (): ()) {
  56        T::add_summary(self, summary)
  57    }
  58}
  59
  60#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  61pub struct NoSummary;
  62
  63/// Catch-all implementation for when you need something that implements [`Summary`] without a specific type.
  64/// We implement it on a `NoSummary` instead of re-using `()`, as that avoids blanket impl collisions with `impl<T: Summary> Dimension for T`
  65/// (as we also need unit type to be a fill-in dimension)
  66impl ContextLessSummary for NoSummary {
  67    fn zero() -> Self {
  68        NoSummary
  69    }
  70
  71    fn add_summary(&mut self, _: &Self) {}
  72}
  73
  74/// Each [`Summary`] type can have more than one [`Dimension`] type that it measures.
  75///
  76/// You can use dimensions to seek to a specific location in the [`SumTree`]
  77///
  78/// # Example:
  79/// Zed's rope has a `TextSummary` type that summarizes lines, characters, and bytes.
  80/// Each of these are different dimensions we may want to seek to
  81pub trait Dimension<'a, S: Summary>: Clone {
  82    fn zero(cx: S::Context<'_>) -> Self;
  83
  84    fn add_summary(&mut self, summary: &'a S, cx: S::Context<'_>);
  85
  86    fn from_summary(summary: &'a S, cx: S::Context<'_>) -> Self {
  87        let mut dimension = Self::zero(cx);
  88        dimension.add_summary(summary, cx);
  89        dimension
  90    }
  91}
  92
  93impl<'a, T: Summary> Dimension<'a, T> for T {
  94    fn zero(cx: T::Context<'_>) -> Self {
  95        Summary::zero(cx)
  96    }
  97
  98    fn add_summary(&mut self, summary: &'a T, cx: T::Context<'_>) {
  99        Summary::add_summary(self, summary, cx);
 100    }
 101}
 102
 103pub trait SeekTarget<'a, S: Summary, D: Dimension<'a, S>> {
 104    fn cmp(&self, cursor_location: &D, cx: S::Context<'_>) -> Ordering;
 105}
 106
 107impl<'a, S: Summary, D: Dimension<'a, S> + Ord> SeekTarget<'a, S, D> for D {
 108    fn cmp(&self, cursor_location: &Self, _: S::Context<'_>) -> Ordering {
 109        Ord::cmp(self, cursor_location)
 110    }
 111}
 112
 113impl<'a, T: Summary> Dimension<'a, T> for () {
 114    fn zero(_: T::Context<'_>) -> Self {}
 115
 116    fn add_summary(&mut self, _: &'a T, _: T::Context<'_>) {}
 117}
 118
 119#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
 120pub struct Dimensions<D1, D2, D3 = ()>(pub D1, pub D2, pub D3);
 121
 122impl<'a, T: Summary, D1: Dimension<'a, T>, D2: Dimension<'a, T>, D3: Dimension<'a, T>>
 123    Dimension<'a, T> for Dimensions<D1, D2, D3>
 124{
 125    fn zero(cx: T::Context<'_>) -> Self {
 126        Dimensions(D1::zero(cx), D2::zero(cx), D3::zero(cx))
 127    }
 128
 129    fn add_summary(&mut self, summary: &'a T, cx: T::Context<'_>) {
 130        self.0.add_summary(summary, cx);
 131        self.1.add_summary(summary, cx);
 132        self.2.add_summary(summary, cx);
 133    }
 134}
 135
 136impl<'a, S, D1, D2, D3> SeekTarget<'a, S, Dimensions<D1, D2, D3>> for D1
 137where
 138    S: Summary,
 139    D1: SeekTarget<'a, S, D1> + Dimension<'a, S>,
 140    D2: Dimension<'a, S>,
 141    D3: Dimension<'a, S>,
 142{
 143    fn cmp(&self, cursor_location: &Dimensions<D1, D2, D3>, cx: S::Context<'_>) -> Ordering {
 144        self.cmp(&cursor_location.0, cx)
 145    }
 146}
 147
 148/// Bias is used to settle ambiguities when determining positions in an ordered sequence.
 149///
 150/// The primary use case is for text, where Bias influences
 151/// which character an offset or anchor is associated with.
 152///
 153/// # Examples
 154/// Given the buffer `AˇBCD`:
 155/// - The offset of the cursor is 1
 156/// - [Bias::Left] would attach the cursor to the character `A`
 157/// - [Bias::Right] would attach the cursor to the character `B`
 158///
 159/// Given the buffer `A«BCˇ»D`:
 160/// - The offset of the cursor is 3, and the selection is from 1 to 3
 161/// - The left anchor of the selection has [Bias::Right], attaching it to the character `B`
 162/// - The right anchor of the selection has [Bias::Left], attaching it to the character `C`
 163///
 164/// Given the buffer `{ˇ<...>`, where `<...>` is a folded region:
 165/// - The display offset of the cursor is 1, but the offset in the buffer is determined by the bias
 166/// - [Bias::Left] would attach the cursor to the character `{`, with a buffer offset of 1
 167/// - [Bias::Right] would attach the cursor to the first character of the folded region,
 168///   and the buffer offset would be the offset of the first character of the folded region
 169#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Debug, Hash, Default)]
 170pub enum Bias {
 171    /// Attach to the character on the left
 172    #[default]
 173    Left,
 174    /// Attach to the character on the right
 175    Right,
 176}
 177
 178impl Bias {
 179    pub const fn invert(self) -> Self {
 180        match self {
 181            Self::Left => Self::Right,
 182            Self::Right => Self::Left,
 183        }
 184    }
 185}
 186
 187/// A B+ tree in which each leaf node contains `Item`s of type `T` and a `Summary`s for each `Item`.
 188/// Each internal node contains a `Summary` of the items in its subtree.
 189///
 190/// The maximum number of items per node is `TREE_BASE * 2`.
 191///
 192/// Any [`Dimension`] supported by the [`Summary`] type can be used to seek to a specific location in the tree.
 193#[derive(Clone)]
 194pub struct SumTree<T: Item>(Arc<Node<T>>);
 195
 196impl<T> fmt::Debug for SumTree<T>
 197where
 198    T: fmt::Debug + Item,
 199    T::Summary: fmt::Debug,
 200{
 201    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 202        f.debug_tuple("SumTree").field(&self.0).finish()
 203    }
 204}
 205
 206impl<T: Item> SumTree<T> {
 207    pub fn new(cx: <T::Summary as Summary>::Context<'_>) -> Self {
 208        SumTree(Arc::new(Node::Leaf {
 209            summary: <T::Summary as Summary>::zero(cx),
 210            items: ArrayVec::new(),
 211            item_summaries: ArrayVec::new(),
 212        }))
 213    }
 214
 215    /// 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.
 216    pub fn from_summary(summary: T::Summary) -> Self {
 217        SumTree(Arc::new(Node::Leaf {
 218            summary,
 219            items: ArrayVec::new(),
 220            item_summaries: ArrayVec::new(),
 221        }))
 222    }
 223
 224    pub fn from_item(item: T, cx: <T::Summary as Summary>::Context<'_>) -> Self {
 225        let mut tree = Self::new(cx);
 226        tree.push(item, cx);
 227        tree
 228    }
 229
 230    pub fn from_iter<I: IntoIterator<Item = T>>(
 231        iter: I,
 232        cx: <T::Summary as Summary>::Context<'_>,
 233    ) -> Self {
 234        let mut nodes = Vec::new();
 235
 236        let mut iter = iter.into_iter().fuse().peekable();
 237        while iter.peek().is_some() {
 238            let items: ArrayVec<T, { 2 * TREE_BASE }> = iter.by_ref().take(2 * TREE_BASE).collect();
 239            let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> =
 240                items.iter().map(|item| item.summary(cx)).collect();
 241
 242            let mut summary = item_summaries[0].clone();
 243            for item_summary in &item_summaries[1..] {
 244                <T::Summary as Summary>::add_summary(&mut summary, item_summary, cx);
 245            }
 246
 247            nodes.push(Node::Leaf {
 248                summary,
 249                items,
 250                item_summaries,
 251            });
 252        }
 253
 254        let mut parent_nodes = Vec::new();
 255        let mut height = 0;
 256        while nodes.len() > 1 {
 257            height += 1;
 258            let mut current_parent_node = None;
 259            for child_node in nodes.drain(..) {
 260                let parent_node = current_parent_node.get_or_insert_with(|| Node::Internal {
 261                    summary: <T::Summary as Summary>::zero(cx),
 262                    height,
 263                    child_summaries: ArrayVec::new(),
 264                    child_trees: ArrayVec::new(),
 265                });
 266                let Node::Internal {
 267                    summary,
 268                    child_summaries,
 269                    child_trees,
 270                    ..
 271                } = parent_node
 272                else {
 273                    unreachable!()
 274                };
 275                let child_summary = child_node.summary();
 276                <T::Summary as Summary>::add_summary(summary, child_summary, cx);
 277                child_summaries.push(child_summary.clone());
 278                child_trees.push(Self(Arc::new(child_node)));
 279
 280                if child_trees.len() == 2 * TREE_BASE {
 281                    parent_nodes.extend(current_parent_node.take());
 282                }
 283            }
 284            parent_nodes.extend(current_parent_node.take());
 285            mem::swap(&mut nodes, &mut parent_nodes);
 286        }
 287
 288        if nodes.is_empty() {
 289            Self::new(cx)
 290        } else {
 291            debug_assert_eq!(nodes.len(), 1);
 292            Self(Arc::new(nodes.pop().unwrap()))
 293        }
 294    }
 295
 296    pub fn from_par_iter<I, Iter>(iter: I, cx: <T::Summary as Summary>::Context<'_>) -> Self
 297    where
 298        I: IntoParallelIterator<Iter = Iter>,
 299        Iter: IndexedParallelIterator<Item = T>,
 300        T: Send + Sync,
 301        T::Summary: Send + Sync,
 302        for<'a> <T::Summary as Summary>::Context<'a>: Sync,
 303    {
 304        let mut nodes = iter
 305            .into_par_iter()
 306            .chunks(2 * TREE_BASE)
 307            .map(|items| {
 308                let items: ArrayVec<T, { 2 * TREE_BASE }> = items.into_iter().collect();
 309                let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> =
 310                    items.iter().map(|item| item.summary(cx)).collect();
 311                let mut summary = item_summaries[0].clone();
 312                for item_summary in &item_summaries[1..] {
 313                    <T::Summary as Summary>::add_summary(&mut summary, item_summary, cx);
 314                }
 315                SumTree(Arc::new(Node::Leaf {
 316                    summary,
 317                    items,
 318                    item_summaries,
 319                }))
 320            })
 321            .collect::<Vec<_>>();
 322
 323        let mut height = 0;
 324        while nodes.len() > 1 {
 325            height += 1;
 326            nodes = nodes
 327                .into_par_iter()
 328                .chunks(2 * TREE_BASE)
 329                .map(|child_nodes| {
 330                    let child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }> =
 331                        child_nodes.into_iter().collect();
 332                    let child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> = child_trees
 333                        .iter()
 334                        .map(|child_tree| child_tree.summary().clone())
 335                        .collect();
 336                    let mut summary = child_summaries[0].clone();
 337                    for child_summary in &child_summaries[1..] {
 338                        <T::Summary as Summary>::add_summary(&mut summary, child_summary, cx);
 339                    }
 340                    SumTree(Arc::new(Node::Internal {
 341                        height,
 342                        summary,
 343                        child_summaries,
 344                        child_trees,
 345                    }))
 346                })
 347                .collect::<Vec<_>>();
 348        }
 349
 350        if nodes.is_empty() {
 351            Self::new(cx)
 352        } else {
 353            debug_assert_eq!(nodes.len(), 1);
 354            nodes.pop().unwrap()
 355        }
 356    }
 357
 358    #[allow(unused)]
 359    pub fn items<'a>(&'a self, cx: <T::Summary as Summary>::Context<'a>) -> Vec<T> {
 360        let mut items = Vec::new();
 361        let mut cursor = self.cursor::<()>(cx);
 362        cursor.next();
 363        while let Some(item) = cursor.item() {
 364            items.push(item.clone());
 365            cursor.next();
 366        }
 367        items
 368    }
 369
 370    pub fn iter(&self) -> Iter<'_, T> {
 371        Iter::new(self)
 372    }
 373
 374    pub fn cursor<'a, 'b, S>(
 375        &'a self,
 376        cx: <T::Summary as Summary>::Context<'b>,
 377    ) -> Cursor<'a, 'b, T, S>
 378    where
 379        S: Dimension<'a, T::Summary>,
 380    {
 381        Cursor::new(self, cx)
 382    }
 383
 384    /// Note: If the summary type requires a non `()` context, then the filter cursor
 385    /// that is returned cannot be used with Rust's iterators.
 386    pub fn filter<'a, 'b, F, U>(
 387        &'a self,
 388        cx: <T::Summary as Summary>::Context<'b>,
 389        filter_node: F,
 390    ) -> FilterCursor<'a, 'b, F, T, U>
 391    where
 392        F: FnMut(&T::Summary) -> bool,
 393        U: Dimension<'a, T::Summary>,
 394    {
 395        FilterCursor::new(self, cx, filter_node)
 396    }
 397
 398    #[allow(dead_code)]
 399    pub fn first(&self) -> Option<&T> {
 400        self.leftmost_leaf().0.items().first()
 401    }
 402
 403    pub fn last(&self) -> Option<&T> {
 404        self.rightmost_leaf().0.items().last()
 405    }
 406
 407    pub fn update_last(
 408        &mut self,
 409        f: impl FnOnce(&mut T),
 410        cx: <T::Summary as Summary>::Context<'_>,
 411    ) {
 412        self.update_last_recursive(f, cx);
 413    }
 414
 415    fn update_last_recursive(
 416        &mut self,
 417        f: impl FnOnce(&mut T),
 418        cx: <T::Summary as Summary>::Context<'_>,
 419    ) -> Option<T::Summary> {
 420        match Arc::make_mut(&mut self.0) {
 421            Node::Internal {
 422                summary,
 423                child_summaries,
 424                child_trees,
 425                ..
 426            } => {
 427                let last_summary = child_summaries.last_mut().unwrap();
 428                let last_child = child_trees.last_mut().unwrap();
 429                *last_summary = last_child.update_last_recursive(f, cx).unwrap();
 430                *summary = sum(child_summaries.iter(), cx);
 431                Some(summary.clone())
 432            }
 433            Node::Leaf {
 434                summary,
 435                items,
 436                item_summaries,
 437            } => {
 438                if let Some((item, item_summary)) = items.last_mut().zip(item_summaries.last_mut())
 439                {
 440                    (f)(item);
 441                    *item_summary = item.summary(cx);
 442                    *summary = sum(item_summaries.iter(), cx);
 443                    Some(summary.clone())
 444                } else {
 445                    None
 446                }
 447            }
 448        }
 449    }
 450
 451    pub fn extent<'a, D: Dimension<'a, T::Summary>>(
 452        &'a self,
 453        cx: <T::Summary as Summary>::Context<'_>,
 454    ) -> D {
 455        let mut extent = D::zero(cx);
 456        match self.0.as_ref() {
 457            Node::Internal { summary, .. } | Node::Leaf { summary, .. } => {
 458                extent.add_summary(summary, cx);
 459            }
 460        }
 461        extent
 462    }
 463
 464    pub fn summary(&self) -> &T::Summary {
 465        match self.0.as_ref() {
 466            Node::Internal { summary, .. } => summary,
 467            Node::Leaf { summary, .. } => summary,
 468        }
 469    }
 470
 471    pub fn is_empty(&self) -> bool {
 472        match self.0.as_ref() {
 473            Node::Internal { .. } => false,
 474            Node::Leaf { items, .. } => items.is_empty(),
 475        }
 476    }
 477
 478    pub fn extend<I>(&mut self, iter: I, cx: <T::Summary as Summary>::Context<'_>)
 479    where
 480        I: IntoIterator<Item = T>,
 481    {
 482        self.append(Self::from_iter(iter, cx), cx);
 483    }
 484
 485    pub fn par_extend<I, Iter>(&mut self, iter: I, cx: <T::Summary as Summary>::Context<'_>)
 486    where
 487        I: IntoParallelIterator<Iter = Iter>,
 488        Iter: IndexedParallelIterator<Item = T>,
 489        T: Send + Sync,
 490        T::Summary: Send + Sync,
 491        for<'a> <T::Summary as Summary>::Context<'a>: Sync,
 492    {
 493        self.append(Self::from_par_iter(iter, cx), cx);
 494    }
 495
 496    pub fn push(&mut self, item: T, cx: <T::Summary as Summary>::Context<'_>) {
 497        let summary = item.summary(cx);
 498        self.append(
 499            SumTree(Arc::new(Node::Leaf {
 500                summary: summary.clone(),
 501                items: ArrayVec::from_iter(Some(item)),
 502                item_summaries: ArrayVec::from_iter(Some(summary)),
 503            })),
 504            cx,
 505        );
 506    }
 507
 508    pub fn append(&mut self, other: Self, cx: <T::Summary as Summary>::Context<'_>) {
 509        if self.is_empty() {
 510            *self = other;
 511        } else if !other.0.is_leaf() || !other.0.items().is_empty() {
 512            if self.0.height() < other.0.height() {
 513                for tree in other.0.child_trees() {
 514                    self.append(tree.clone(), cx);
 515                }
 516            } else if let Some(split_tree) = self.push_tree_recursive(other, cx) {
 517                *self = Self::from_child_trees(self.clone(), split_tree, cx);
 518            }
 519        }
 520    }
 521
 522    fn push_tree_recursive(
 523        &mut self,
 524        other: SumTree<T>,
 525        cx: <T::Summary as Summary>::Context<'_>,
 526    ) -> Option<SumTree<T>> {
 527        match Arc::make_mut(&mut self.0) {
 528            Node::Internal {
 529                height,
 530                summary,
 531                child_summaries,
 532                child_trees,
 533                ..
 534            } => {
 535                let other_node = other.0.clone();
 536                <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
 537
 538                let height_delta = *height - other_node.height();
 539                let mut summaries_to_append = ArrayVec::<T::Summary, { 2 * TREE_BASE }>::new();
 540                let mut trees_to_append = ArrayVec::<SumTree<T>, { 2 * TREE_BASE }>::new();
 541                if height_delta == 0 {
 542                    summaries_to_append.extend(other_node.child_summaries().iter().cloned());
 543                    trees_to_append.extend(other_node.child_trees().iter().cloned());
 544                } else if height_delta == 1 && !other_node.is_underflowing() {
 545                    summaries_to_append.push(other_node.summary().clone());
 546                    trees_to_append.push(other)
 547                } else {
 548                    let tree_to_append = child_trees
 549                        .last_mut()
 550                        .unwrap()
 551                        .push_tree_recursive(other, cx);
 552                    *child_summaries.last_mut().unwrap() =
 553                        child_trees.last().unwrap().0.summary().clone();
 554
 555                    if let Some(split_tree) = tree_to_append {
 556                        summaries_to_append.push(split_tree.0.summary().clone());
 557                        trees_to_append.push(split_tree);
 558                    }
 559                }
 560
 561                let child_count = child_trees.len() + trees_to_append.len();
 562                if child_count > 2 * TREE_BASE {
 563                    let left_summaries: ArrayVec<_, { 2 * TREE_BASE }>;
 564                    let right_summaries: ArrayVec<_, { 2 * TREE_BASE }>;
 565                    let left_trees;
 566                    let right_trees;
 567
 568                    let midpoint = (child_count + child_count % 2) / 2;
 569                    {
 570                        let mut all_summaries = child_summaries
 571                            .iter()
 572                            .chain(summaries_to_append.iter())
 573                            .cloned();
 574                        left_summaries = all_summaries.by_ref().take(midpoint).collect();
 575                        right_summaries = all_summaries.collect();
 576                        let mut all_trees =
 577                            child_trees.iter().chain(trees_to_append.iter()).cloned();
 578                        left_trees = all_trees.by_ref().take(midpoint).collect();
 579                        right_trees = all_trees.collect();
 580                    }
 581                    *summary = sum(left_summaries.iter(), cx);
 582                    *child_summaries = left_summaries;
 583                    *child_trees = left_trees;
 584
 585                    Some(SumTree(Arc::new(Node::Internal {
 586                        height: *height,
 587                        summary: sum(right_summaries.iter(), cx),
 588                        child_summaries: right_summaries,
 589                        child_trees: right_trees,
 590                    })))
 591                } else {
 592                    child_summaries.extend(summaries_to_append);
 593                    child_trees.extend(trees_to_append);
 594                    None
 595                }
 596            }
 597            Node::Leaf {
 598                summary,
 599                items,
 600                item_summaries,
 601            } => {
 602                let other_node = other.0;
 603
 604                let child_count = items.len() + other_node.items().len();
 605                if child_count > 2 * TREE_BASE {
 606                    let left_items;
 607                    let right_items;
 608                    let left_summaries;
 609                    let right_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>;
 610
 611                    let midpoint = (child_count + child_count % 2) / 2;
 612                    {
 613                        let mut all_items = items.iter().chain(other_node.items().iter()).cloned();
 614                        left_items = all_items.by_ref().take(midpoint).collect();
 615                        right_items = all_items.collect();
 616
 617                        let mut all_summaries = item_summaries
 618                            .iter()
 619                            .chain(other_node.child_summaries())
 620                            .cloned();
 621                        left_summaries = all_summaries.by_ref().take(midpoint).collect();
 622                        right_summaries = all_summaries.collect();
 623                    }
 624                    *items = left_items;
 625                    *item_summaries = left_summaries;
 626                    *summary = sum(item_summaries.iter(), cx);
 627                    Some(SumTree(Arc::new(Node::Leaf {
 628                        items: right_items,
 629                        summary: sum(right_summaries.iter(), cx),
 630                        item_summaries: right_summaries,
 631                    })))
 632                } else {
 633                    <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
 634                    items.extend(other_node.items().iter().cloned());
 635                    item_summaries.extend(other_node.child_summaries().iter().cloned());
 636                    None
 637                }
 638            }
 639        }
 640    }
 641
 642    fn from_child_trees(
 643        left: SumTree<T>,
 644        right: SumTree<T>,
 645        cx: <T::Summary as Summary>::Context<'_>,
 646    ) -> Self {
 647        let height = left.0.height() + 1;
 648        let mut child_summaries = ArrayVec::new();
 649        child_summaries.push(left.0.summary().clone());
 650        child_summaries.push(right.0.summary().clone());
 651        let mut child_trees = ArrayVec::new();
 652        child_trees.push(left);
 653        child_trees.push(right);
 654        SumTree(Arc::new(Node::Internal {
 655            height,
 656            summary: sum(child_summaries.iter(), cx),
 657            child_summaries,
 658            child_trees,
 659        }))
 660    }
 661
 662    fn leftmost_leaf(&self) -> &Self {
 663        match *self.0 {
 664            Node::Leaf { .. } => self,
 665            Node::Internal {
 666                ref child_trees, ..
 667            } => child_trees.first().unwrap().leftmost_leaf(),
 668        }
 669    }
 670
 671    fn rightmost_leaf(&self) -> &Self {
 672        match *self.0 {
 673            Node::Leaf { .. } => self,
 674            Node::Internal {
 675                ref child_trees, ..
 676            } => child_trees.last().unwrap().rightmost_leaf(),
 677        }
 678    }
 679}
 680
 681impl<T: Item + PartialEq> PartialEq for SumTree<T> {
 682    fn eq(&self, other: &Self) -> bool {
 683        self.iter().eq(other.iter())
 684    }
 685}
 686
 687impl<T: Item + Eq> Eq for SumTree<T> {}
 688
 689impl<T: KeyedItem> SumTree<T> {
 690    pub fn insert_or_replace<'a, 'b>(
 691        &'a mut self,
 692        item: T,
 693        cx: <T::Summary as Summary>::Context<'b>,
 694    ) -> Option<T> {
 695        let mut replaced = None;
 696        {
 697            let mut cursor = self.cursor::<T::Key>(cx);
 698            let mut new_tree = cursor.slice(&item.key(), Bias::Left);
 699            if let Some(cursor_item) = cursor.item()
 700                && cursor_item.key() == item.key()
 701            {
 702                replaced = Some(cursor_item.clone());
 703                cursor.next();
 704            }
 705            new_tree.push(item, cx);
 706            new_tree.append(cursor.suffix(), cx);
 707            drop(cursor);
 708            *self = new_tree
 709        };
 710        replaced
 711    }
 712
 713    pub fn remove(&mut self, key: &T::Key, cx: <T::Summary as Summary>::Context<'_>) -> Option<T> {
 714        let mut removed = None;
 715        *self = {
 716            let mut cursor = self.cursor::<T::Key>(cx);
 717            let mut new_tree = cursor.slice(key, Bias::Left);
 718            if let Some(item) = cursor.item()
 719                && item.key() == *key
 720            {
 721                removed = Some(item.clone());
 722                cursor.next();
 723            }
 724            new_tree.append(cursor.suffix(), cx);
 725            new_tree
 726        };
 727        removed
 728    }
 729
 730    pub fn edit(
 731        &mut self,
 732        mut edits: Vec<Edit<T>>,
 733        cx: <T::Summary as Summary>::Context<'_>,
 734    ) -> Vec<T> {
 735        if edits.is_empty() {
 736            return Vec::new();
 737        }
 738
 739        let mut removed = Vec::new();
 740        edits.sort_unstable_by_key(|item| item.key());
 741
 742        *self = {
 743            let mut cursor = self.cursor::<T::Key>(cx);
 744            let mut new_tree = SumTree::new(cx);
 745            let mut buffered_items = Vec::new();
 746
 747            cursor.seek(&T::Key::zero(cx), Bias::Left);
 748            for edit in edits {
 749                let new_key = edit.key();
 750                let mut old_item = cursor.item();
 751
 752                if old_item
 753                    .as_ref()
 754                    .is_some_and(|old_item| old_item.key() < new_key)
 755                {
 756                    new_tree.extend(buffered_items.drain(..), cx);
 757                    let slice = cursor.slice(&new_key, Bias::Left);
 758                    new_tree.append(slice, cx);
 759                    old_item = cursor.item();
 760                }
 761
 762                if let Some(old_item) = old_item
 763                    && old_item.key() == new_key
 764                {
 765                    removed.push(old_item.clone());
 766                    cursor.next();
 767                }
 768
 769                match edit {
 770                    Edit::Insert(item) => {
 771                        buffered_items.push(item);
 772                    }
 773                    Edit::Remove(_) => {}
 774                }
 775            }
 776
 777            new_tree.extend(buffered_items, cx);
 778            new_tree.append(cursor.suffix(), cx);
 779            new_tree
 780        };
 781
 782        removed
 783    }
 784
 785    pub fn get<'a>(
 786        &'a self,
 787        key: &T::Key,
 788        cx: <T::Summary as Summary>::Context<'a>,
 789    ) -> Option<&'a T> {
 790        let mut cursor = self.cursor::<T::Key>(cx);
 791        if cursor.seek(key, Bias::Left) {
 792            cursor.item()
 793        } else {
 794            None
 795        }
 796    }
 797}
 798
 799impl<T, S> Default for SumTree<T>
 800where
 801    T: Item<Summary = S>,
 802    S: for<'a> Summary<Context<'a> = ()>,
 803{
 804    fn default() -> Self {
 805        Self::new(())
 806    }
 807}
 808
 809#[derive(Clone)]
 810pub enum Node<T: Item> {
 811    Internal {
 812        height: u8,
 813        summary: T::Summary,
 814        child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
 815        child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }>,
 816    },
 817    Leaf {
 818        summary: T::Summary,
 819        items: ArrayVec<T, { 2 * TREE_BASE }>,
 820        item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
 821    },
 822}
 823
 824impl<T> fmt::Debug for Node<T>
 825where
 826    T: Item + fmt::Debug,
 827    T::Summary: fmt::Debug,
 828{
 829    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 830        match self {
 831            Node::Internal {
 832                height,
 833                summary,
 834                child_summaries,
 835                child_trees,
 836            } => f
 837                .debug_struct("Internal")
 838                .field("height", height)
 839                .field("summary", summary)
 840                .field("child_summaries", child_summaries)
 841                .field("child_trees", child_trees)
 842                .finish(),
 843            Node::Leaf {
 844                summary,
 845                items,
 846                item_summaries,
 847            } => f
 848                .debug_struct("Leaf")
 849                .field("summary", summary)
 850                .field("items", items)
 851                .field("item_summaries", item_summaries)
 852                .finish(),
 853        }
 854    }
 855}
 856
 857impl<T: Item> Node<T> {
 858    const fn is_leaf(&self) -> bool {
 859        matches!(self, Node::Leaf { .. })
 860    }
 861
 862    const fn height(&self) -> u8 {
 863        match self {
 864            Node::Internal { height, .. } => *height,
 865            Node::Leaf { .. } => 0,
 866        }
 867    }
 868
 869    const fn summary(&self) -> &T::Summary {
 870        match self {
 871            Node::Internal { summary, .. } => summary,
 872            Node::Leaf { summary, .. } => summary,
 873        }
 874    }
 875
 876    fn child_summaries(&self) -> &[T::Summary] {
 877        match self {
 878            Node::Internal {
 879                child_summaries, ..
 880            } => child_summaries.as_slice(),
 881            Node::Leaf { item_summaries, .. } => item_summaries.as_slice(),
 882        }
 883    }
 884
 885    fn child_trees(&self) -> &ArrayVec<SumTree<T>, { 2 * TREE_BASE }> {
 886        match self {
 887            Node::Internal { child_trees, .. } => child_trees,
 888            Node::Leaf { .. } => panic!("Leaf nodes have no child trees"),
 889        }
 890    }
 891
 892    fn items(&self) -> &ArrayVec<T, { 2 * TREE_BASE }> {
 893        match self {
 894            Node::Leaf { items, .. } => items,
 895            Node::Internal { .. } => panic!("Internal nodes have no items"),
 896        }
 897    }
 898
 899    const fn is_underflowing(&self) -> bool {
 900        match self {
 901            Node::Internal { child_trees, .. } => child_trees.len() < TREE_BASE,
 902            Node::Leaf { items, .. } => items.len() < TREE_BASE,
 903        }
 904    }
 905}
 906
 907#[derive(Debug)]
 908pub enum Edit<T: KeyedItem> {
 909    Insert(T),
 910    Remove(T::Key),
 911}
 912
 913impl<T: KeyedItem> Edit<T> {
 914    fn key(&self) -> T::Key {
 915        match self {
 916            Edit::Insert(item) => item.key(),
 917            Edit::Remove(key) => key.clone(),
 918        }
 919    }
 920}
 921
 922fn sum<'a, T, I>(iter: I, cx: T::Context<'_>) -> T
 923where
 924    T: 'a + Summary,
 925    I: Iterator<Item = &'a T>,
 926{
 927    let mut sum = T::zero(cx);
 928    for value in iter {
 929        sum.add_summary(value, cx);
 930    }
 931    sum
 932}
 933
 934#[cfg(test)]
 935mod tests {
 936    use super::*;
 937    use rand::{distr::StandardUniform, prelude::*};
 938    use std::cmp;
 939
 940    #[ctor::ctor]
 941    fn init_logger() {
 942        zlog::init_test();
 943    }
 944
 945    #[test]
 946    fn test_extend_and_push_tree() {
 947        let mut tree1 = SumTree::default();
 948        tree1.extend(0..20, ());
 949
 950        let mut tree2 = SumTree::default();
 951        tree2.extend(50..100, ());
 952
 953        tree1.append(tree2, ());
 954        assert_eq!(tree1.items(()), (0..20).chain(50..100).collect::<Vec<u8>>());
 955    }
 956
 957    #[test]
 958    fn test_random() {
 959        let mut starting_seed = 0;
 960        if let Ok(value) = std::env::var("SEED") {
 961            starting_seed = value.parse().expect("invalid SEED variable");
 962        }
 963        let mut num_iterations = 100;
 964        if let Ok(value) = std::env::var("ITERATIONS") {
 965            num_iterations = value.parse().expect("invalid ITERATIONS variable");
 966        }
 967        let num_operations = std::env::var("OPERATIONS")
 968            .map_or(5, |o| o.parse().expect("invalid OPERATIONS variable"));
 969
 970        for seed in starting_seed..(starting_seed + num_iterations) {
 971            eprintln!("seed = {}", seed);
 972            let mut rng = StdRng::seed_from_u64(seed);
 973
 974            let rng = &mut rng;
 975            let mut tree = SumTree::<u8>::default();
 976            let count = rng.random_range(0..10);
 977            if rng.random() {
 978                tree.extend(rng.sample_iter(StandardUniform).take(count), ());
 979            } else {
 980                let items = rng
 981                    .sample_iter(StandardUniform)
 982                    .take(count)
 983                    .collect::<Vec<_>>();
 984                tree.par_extend(items, ());
 985            }
 986
 987            for _ in 0..num_operations {
 988                let splice_end = rng.random_range(0..tree.extent::<Count>(()).0 + 1);
 989                let splice_start = rng.random_range(0..splice_end + 1);
 990                let count = rng.random_range(0..10);
 991                let tree_end = tree.extent::<Count>(());
 992                let new_items = rng
 993                    .sample_iter(StandardUniform)
 994                    .take(count)
 995                    .collect::<Vec<u8>>();
 996
 997                let mut reference_items = tree.items(());
 998                reference_items.splice(splice_start..splice_end, new_items.clone());
 999
1000                tree = {
1001                    let mut cursor = tree.cursor::<Count>(());
1002                    let mut new_tree = cursor.slice(&Count(splice_start), Bias::Right);
1003                    if rng.random() {
1004                        new_tree.extend(new_items, ());
1005                    } else {
1006                        new_tree.par_extend(new_items, ());
1007                    }
1008                    cursor.seek(&Count(splice_end), Bias::Right);
1009                    new_tree.append(cursor.slice(&tree_end, Bias::Right), ());
1010                    new_tree
1011                };
1012
1013                assert_eq!(tree.items(()), reference_items);
1014                assert_eq!(
1015                    tree.iter().collect::<Vec<_>>(),
1016                    tree.cursor::<()>(()).collect::<Vec<_>>()
1017                );
1018
1019                log::info!("tree items: {:?}", tree.items(()));
1020
1021                let mut filter_cursor =
1022                    tree.filter::<_, Count>((), |summary| summary.contains_even);
1023                let expected_filtered_items = tree
1024                    .items(())
1025                    .into_iter()
1026                    .enumerate()
1027                    .filter(|(_, item)| (item & 1) == 0)
1028                    .collect::<Vec<_>>();
1029
1030                let mut item_ix = if rng.random() {
1031                    filter_cursor.next();
1032                    0
1033                } else {
1034                    filter_cursor.prev();
1035                    expected_filtered_items.len().saturating_sub(1)
1036                };
1037                while item_ix < expected_filtered_items.len() {
1038                    log::info!("filter_cursor, item_ix: {}", item_ix);
1039                    let actual_item = filter_cursor.item().unwrap();
1040                    let (reference_index, reference_item) = expected_filtered_items[item_ix];
1041                    assert_eq!(actual_item, &reference_item);
1042                    assert_eq!(filter_cursor.start().0, reference_index);
1043                    log::info!("next");
1044                    filter_cursor.next();
1045                    item_ix += 1;
1046
1047                    while item_ix > 0 && rng.random_bool(0.2) {
1048                        log::info!("prev");
1049                        filter_cursor.prev();
1050                        item_ix -= 1;
1051
1052                        if item_ix == 0 && rng.random_bool(0.2) {
1053                            filter_cursor.prev();
1054                            assert_eq!(filter_cursor.item(), None);
1055                            assert_eq!(filter_cursor.start().0, 0);
1056                            filter_cursor.next();
1057                        }
1058                    }
1059                }
1060                assert_eq!(filter_cursor.item(), None);
1061
1062                let mut before_start = false;
1063                let mut cursor = tree.cursor::<Count>(());
1064                let start_pos = rng.random_range(0..=reference_items.len());
1065                cursor.seek(&Count(start_pos), Bias::Right);
1066                let mut pos = rng.random_range(start_pos..=reference_items.len());
1067                cursor.seek_forward(&Count(pos), Bias::Right);
1068
1069                for i in 0..10 {
1070                    assert_eq!(cursor.start().0, pos);
1071
1072                    if pos > 0 {
1073                        assert_eq!(cursor.prev_item().unwrap(), &reference_items[pos - 1]);
1074                    } else {
1075                        assert_eq!(cursor.prev_item(), None);
1076                    }
1077
1078                    if pos < reference_items.len() && !before_start {
1079                        assert_eq!(cursor.item().unwrap(), &reference_items[pos]);
1080                    } else {
1081                        assert_eq!(cursor.item(), None);
1082                    }
1083
1084                    if before_start {
1085                        assert_eq!(cursor.next_item(), reference_items.first());
1086                    } else if pos + 1 < reference_items.len() {
1087                        assert_eq!(cursor.next_item().unwrap(), &reference_items[pos + 1]);
1088                    } else {
1089                        assert_eq!(cursor.next_item(), None);
1090                    }
1091
1092                    if i < 5 {
1093                        cursor.next();
1094                        if pos < reference_items.len() {
1095                            pos += 1;
1096                            before_start = false;
1097                        }
1098                    } else {
1099                        cursor.prev();
1100                        if pos == 0 {
1101                            before_start = true;
1102                        }
1103                        pos = pos.saturating_sub(1);
1104                    }
1105                }
1106            }
1107
1108            for _ in 0..10 {
1109                let end = rng.random_range(0..tree.extent::<Count>(()).0 + 1);
1110                let start = rng.random_range(0..end + 1);
1111                let start_bias = if rng.random() {
1112                    Bias::Left
1113                } else {
1114                    Bias::Right
1115                };
1116                let end_bias = if rng.random() {
1117                    Bias::Left
1118                } else {
1119                    Bias::Right
1120                };
1121
1122                let mut cursor = tree.cursor::<Count>(());
1123                cursor.seek(&Count(start), start_bias);
1124                let slice = cursor.slice(&Count(end), end_bias);
1125
1126                cursor.seek(&Count(start), start_bias);
1127                let summary = cursor.summary::<_, Sum>(&Count(end), end_bias);
1128
1129                assert_eq!(summary.0, slice.summary().sum);
1130            }
1131        }
1132    }
1133
1134    #[test]
1135    fn test_cursor() {
1136        // Empty tree
1137        let tree = SumTree::<u8>::default();
1138        let mut cursor = tree.cursor::<IntegersSummary>(());
1139        assert_eq!(
1140            cursor.slice(&Count(0), Bias::Right).items(()),
1141            Vec::<u8>::new()
1142        );
1143        assert_eq!(cursor.item(), None);
1144        assert_eq!(cursor.prev_item(), None);
1145        assert_eq!(cursor.next_item(), None);
1146        assert_eq!(cursor.start().sum, 0);
1147        cursor.prev();
1148        assert_eq!(cursor.item(), None);
1149        assert_eq!(cursor.prev_item(), None);
1150        assert_eq!(cursor.next_item(), None);
1151        assert_eq!(cursor.start().sum, 0);
1152        cursor.next();
1153        assert_eq!(cursor.item(), None);
1154        assert_eq!(cursor.prev_item(), None);
1155        assert_eq!(cursor.next_item(), None);
1156        assert_eq!(cursor.start().sum, 0);
1157
1158        // Single-element tree
1159        let mut tree = SumTree::<u8>::default();
1160        tree.extend(vec![1], ());
1161        let mut cursor = tree.cursor::<IntegersSummary>(());
1162        assert_eq!(
1163            cursor.slice(&Count(0), Bias::Right).items(()),
1164            Vec::<u8>::new()
1165        );
1166        assert_eq!(cursor.item(), Some(&1));
1167        assert_eq!(cursor.prev_item(), None);
1168        assert_eq!(cursor.next_item(), None);
1169        assert_eq!(cursor.start().sum, 0);
1170
1171        cursor.next();
1172        assert_eq!(cursor.item(), None);
1173        assert_eq!(cursor.prev_item(), Some(&1));
1174        assert_eq!(cursor.next_item(), None);
1175        assert_eq!(cursor.start().sum, 1);
1176
1177        cursor.prev();
1178        assert_eq!(cursor.item(), Some(&1));
1179        assert_eq!(cursor.prev_item(), None);
1180        assert_eq!(cursor.next_item(), None);
1181        assert_eq!(cursor.start().sum, 0);
1182
1183        let mut cursor = tree.cursor::<IntegersSummary>(());
1184        assert_eq!(cursor.slice(&Count(1), Bias::Right).items(()), [1]);
1185        assert_eq!(cursor.item(), None);
1186        assert_eq!(cursor.prev_item(), Some(&1));
1187        assert_eq!(cursor.next_item(), None);
1188        assert_eq!(cursor.start().sum, 1);
1189
1190        cursor.seek(&Count(0), Bias::Right);
1191        assert_eq!(
1192            cursor
1193                .slice(&tree.extent::<Count>(()), Bias::Right)
1194                .items(()),
1195            [1]
1196        );
1197        assert_eq!(cursor.item(), None);
1198        assert_eq!(cursor.prev_item(), Some(&1));
1199        assert_eq!(cursor.next_item(), None);
1200        assert_eq!(cursor.start().sum, 1);
1201
1202        // Multiple-element tree
1203        let mut tree = SumTree::default();
1204        tree.extend(vec![1, 2, 3, 4, 5, 6], ());
1205        let mut cursor = tree.cursor::<IntegersSummary>(());
1206
1207        assert_eq!(cursor.slice(&Count(2), Bias::Right).items(()), [1, 2]);
1208        assert_eq!(cursor.item(), Some(&3));
1209        assert_eq!(cursor.prev_item(), Some(&2));
1210        assert_eq!(cursor.next_item(), Some(&4));
1211        assert_eq!(cursor.start().sum, 3);
1212
1213        cursor.next();
1214        assert_eq!(cursor.item(), Some(&4));
1215        assert_eq!(cursor.prev_item(), Some(&3));
1216        assert_eq!(cursor.next_item(), Some(&5));
1217        assert_eq!(cursor.start().sum, 6);
1218
1219        cursor.next();
1220        assert_eq!(cursor.item(), Some(&5));
1221        assert_eq!(cursor.prev_item(), Some(&4));
1222        assert_eq!(cursor.next_item(), Some(&6));
1223        assert_eq!(cursor.start().sum, 10);
1224
1225        cursor.next();
1226        assert_eq!(cursor.item(), Some(&6));
1227        assert_eq!(cursor.prev_item(), Some(&5));
1228        assert_eq!(cursor.next_item(), None);
1229        assert_eq!(cursor.start().sum, 15);
1230
1231        cursor.next();
1232        cursor.next();
1233        assert_eq!(cursor.item(), None);
1234        assert_eq!(cursor.prev_item(), Some(&6));
1235        assert_eq!(cursor.next_item(), None);
1236        assert_eq!(cursor.start().sum, 21);
1237
1238        cursor.prev();
1239        assert_eq!(cursor.item(), Some(&6));
1240        assert_eq!(cursor.prev_item(), Some(&5));
1241        assert_eq!(cursor.next_item(), None);
1242        assert_eq!(cursor.start().sum, 15);
1243
1244        cursor.prev();
1245        assert_eq!(cursor.item(), Some(&5));
1246        assert_eq!(cursor.prev_item(), Some(&4));
1247        assert_eq!(cursor.next_item(), Some(&6));
1248        assert_eq!(cursor.start().sum, 10);
1249
1250        cursor.prev();
1251        assert_eq!(cursor.item(), Some(&4));
1252        assert_eq!(cursor.prev_item(), Some(&3));
1253        assert_eq!(cursor.next_item(), Some(&5));
1254        assert_eq!(cursor.start().sum, 6);
1255
1256        cursor.prev();
1257        assert_eq!(cursor.item(), Some(&3));
1258        assert_eq!(cursor.prev_item(), Some(&2));
1259        assert_eq!(cursor.next_item(), Some(&4));
1260        assert_eq!(cursor.start().sum, 3);
1261
1262        cursor.prev();
1263        assert_eq!(cursor.item(), Some(&2));
1264        assert_eq!(cursor.prev_item(), Some(&1));
1265        assert_eq!(cursor.next_item(), Some(&3));
1266        assert_eq!(cursor.start().sum, 1);
1267
1268        cursor.prev();
1269        assert_eq!(cursor.item(), Some(&1));
1270        assert_eq!(cursor.prev_item(), None);
1271        assert_eq!(cursor.next_item(), Some(&2));
1272        assert_eq!(cursor.start().sum, 0);
1273
1274        cursor.prev();
1275        assert_eq!(cursor.item(), None);
1276        assert_eq!(cursor.prev_item(), None);
1277        assert_eq!(cursor.next_item(), Some(&1));
1278        assert_eq!(cursor.start().sum, 0);
1279
1280        cursor.next();
1281        assert_eq!(cursor.item(), Some(&1));
1282        assert_eq!(cursor.prev_item(), None);
1283        assert_eq!(cursor.next_item(), Some(&2));
1284        assert_eq!(cursor.start().sum, 0);
1285
1286        let mut cursor = tree.cursor::<IntegersSummary>(());
1287        assert_eq!(
1288            cursor
1289                .slice(&tree.extent::<Count>(()), Bias::Right)
1290                .items(()),
1291            tree.items(())
1292        );
1293        assert_eq!(cursor.item(), None);
1294        assert_eq!(cursor.prev_item(), Some(&6));
1295        assert_eq!(cursor.next_item(), None);
1296        assert_eq!(cursor.start().sum, 21);
1297
1298        cursor.seek(&Count(3), Bias::Right);
1299        assert_eq!(
1300            cursor
1301                .slice(&tree.extent::<Count>(()), Bias::Right)
1302                .items(()),
1303            [4, 5, 6]
1304        );
1305        assert_eq!(cursor.item(), None);
1306        assert_eq!(cursor.prev_item(), Some(&6));
1307        assert_eq!(cursor.next_item(), None);
1308        assert_eq!(cursor.start().sum, 21);
1309
1310        // Seeking can bias left or right
1311        cursor.seek(&Count(1), Bias::Left);
1312        assert_eq!(cursor.item(), Some(&1));
1313        cursor.seek(&Count(1), Bias::Right);
1314        assert_eq!(cursor.item(), Some(&2));
1315
1316        // Slicing without resetting starts from where the cursor is parked at.
1317        cursor.seek(&Count(1), Bias::Right);
1318        assert_eq!(cursor.slice(&Count(3), Bias::Right).items(()), vec![2, 3]);
1319        assert_eq!(cursor.slice(&Count(6), Bias::Left).items(()), vec![4, 5]);
1320        assert_eq!(cursor.slice(&Count(6), Bias::Right).items(()), vec![6]);
1321    }
1322
1323    #[test]
1324    fn test_edit() {
1325        let mut tree = SumTree::<u8>::default();
1326
1327        let removed = tree.edit(vec![Edit::Insert(1), Edit::Insert(2), Edit::Insert(0)], ());
1328        assert_eq!(tree.items(()), vec![0, 1, 2]);
1329        assert_eq!(removed, Vec::<u8>::new());
1330        assert_eq!(tree.get(&0, ()), Some(&0));
1331        assert_eq!(tree.get(&1, ()), Some(&1));
1332        assert_eq!(tree.get(&2, ()), Some(&2));
1333        assert_eq!(tree.get(&4, ()), None);
1334
1335        let removed = tree.edit(vec![Edit::Insert(2), Edit::Insert(4), Edit::Remove(0)], ());
1336        assert_eq!(tree.items(()), vec![1, 2, 4]);
1337        assert_eq!(removed, vec![0, 2]);
1338        assert_eq!(tree.get(&0, ()), None);
1339        assert_eq!(tree.get(&1, ()), Some(&1));
1340        assert_eq!(tree.get(&2, ()), Some(&2));
1341        assert_eq!(tree.get(&4, ()), Some(&4));
1342    }
1343
1344    #[test]
1345    fn test_from_iter() {
1346        assert_eq!(
1347            SumTree::from_iter(0..100, ()).items(()),
1348            (0..100).collect::<Vec<_>>()
1349        );
1350
1351        // Ensure `from_iter` works correctly when the given iterator restarts
1352        // after calling `next` if `None` was already returned.
1353        let mut ix = 0;
1354        let iterator = std::iter::from_fn(|| {
1355            ix = (ix + 1) % 2;
1356            if ix == 1 { Some(1) } else { None }
1357        });
1358        assert_eq!(SumTree::from_iter(iterator, ()).items(()), vec![1]);
1359    }
1360
1361    #[derive(Clone, Default, Debug)]
1362    pub struct IntegersSummary {
1363        count: usize,
1364        sum: usize,
1365        contains_even: bool,
1366        max: u8,
1367    }
1368
1369    #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1370    struct Count(usize);
1371
1372    #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1373    struct Sum(usize);
1374
1375    impl Item for u8 {
1376        type Summary = IntegersSummary;
1377
1378        fn summary(&self, _cx: ()) -> Self::Summary {
1379            IntegersSummary {
1380                count: 1,
1381                sum: *self as usize,
1382                contains_even: (*self & 1) == 0,
1383                max: *self,
1384            }
1385        }
1386    }
1387
1388    impl KeyedItem for u8 {
1389        type Key = u8;
1390
1391        fn key(&self) -> Self::Key {
1392            *self
1393        }
1394    }
1395
1396    impl ContextLessSummary for IntegersSummary {
1397        fn zero() -> Self {
1398            Default::default()
1399        }
1400
1401        fn add_summary(&mut self, other: &Self) {
1402            self.count += other.count;
1403            self.sum += other.sum;
1404            self.contains_even |= other.contains_even;
1405            self.max = cmp::max(self.max, other.max);
1406        }
1407    }
1408
1409    impl Dimension<'_, IntegersSummary> for u8 {
1410        fn zero(_cx: ()) -> Self {
1411            Default::default()
1412        }
1413
1414        fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1415            *self = summary.max;
1416        }
1417    }
1418
1419    impl Dimension<'_, IntegersSummary> for Count {
1420        fn zero(_cx: ()) -> Self {
1421            Default::default()
1422        }
1423
1424        fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1425            self.0 += summary.count;
1426        }
1427    }
1428
1429    impl SeekTarget<'_, IntegersSummary, IntegersSummary> for Count {
1430        fn cmp(&self, cursor_location: &IntegersSummary, _: ()) -> Ordering {
1431            self.0.cmp(&cursor_location.count)
1432        }
1433    }
1434
1435    impl Dimension<'_, IntegersSummary> for Sum {
1436        fn zero(_cx: ()) -> Self {
1437            Default::default()
1438        }
1439
1440        fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1441            self.0 += summary.sum;
1442        }
1443    }
1444}