sum_tree: Make SumTree::append run in logarithmic time

crates/rope/benches/rope_benchmark.rs

@@ -238,6 +238,35 @@ fn rope_benchmarks(c: &mut Criterion) {
         });
     }
     group.finish();
+
+    let mut group = c.benchmark_group("append many");
+    group.throughput(Throughput::Bytes(128 * 100_000));
+
+    group.bench_function("small to large", |b| {
+        b.iter(|| {
+            let mut rope = Rope::new();
+            let small = Rope::from("A".repeat(128));
+            for _ in 0..100_000 {
+                rope.append(small.clone());
+            }
+            assert_eq!(rope.len(), 128 * 100_000);
+        });
+    });
+
+    group.bench_function("large to small", |b| {
+        b.iter(|| {
+            let mut rope = Rope::new();
+            let small = Rope::from("A".repeat(128));
+            for _ in 0..100_000 {
+                let large = rope;
+                rope = small.clone();
+                rope.append(large);
+            }
+            assert_eq!(rope.len(), 128 * 100_000);
+        });
+    });
+
+    group.finish();
 }
 
 criterion_group!(benches, rope_benchmarks);

crates/sum_tree/src/sum_tree.rs

@@ -620,13 +620,15 @@ impl<T: Item> SumTree<T> {
         );
     }
 
-    pub fn append(&mut self, other: Self, cx: <T::Summary as Summary>::Context<'_>) {
+    pub fn append(&mut self, mut other: Self, cx: <T::Summary as Summary>::Context<'_>) {
         if self.is_empty() {
             *self = other;
         } else if !other.0.is_leaf() || !other.0.items().is_empty() {
             if self.0.height() < other.0.height() {
-                for tree in other.0.child_trees() {
-                    self.append(tree.clone(), cx);
+                if let Some(tree) = Self::append_large(self.clone(), &mut other, cx) {
+                    *self = Self::from_child_trees(tree, other, cx);
+                } else {
+                    *self = other;
                 }
             } else if let Some(split_tree) = self.push_tree_recursive(other, cx) {
                 *self = Self::from_child_trees(self.clone(), split_tree, cx);
@@ -754,6 +756,186 @@ impl<T: Item> SumTree<T> {
         }
     }
 
+    // appends the `large` tree to a `small` tree, assumes small.height() <= large.height()
+    fn append_large(
+        small: Self,
+        large: &mut Self,
+        cx: <T::Summary as Summary>::Context<'_>,
+    ) -> Option<Self> {
+        if small.0.height() == large.0.height() {
+            if !small.0.is_underflowing() {
+                Some(small)
+            } else {
+                Self::merge_into_right(small, large, cx)
+            }
+        } else {
+            debug_assert!(small.0.height() < large.0.height());
+            let Node::Internal {
+                height,
+                summary,
+                child_summaries,
+                child_trees,
+            } = Arc::make_mut(&mut large.0)
+            else {
+                unreachable!();
+            };
+            let mut full_summary = small.summary().clone();
+            Summary::add_summary(&mut full_summary, summary, cx);
+            *summary = full_summary;
+
+            let first = child_trees.first_mut().unwrap();
+            let res = Self::append_large(small, first, cx);
+            *child_summaries.first_mut().unwrap() = first.summary().clone();
+            if let Some(tree) = res {
+                if child_trees.len() < 2 * TREE_BASE {
+                    child_summaries.insert(0, tree.summary().clone());
+                    child_trees.insert(0, tree);
+                    None
+                } else {
+                    let new_child_summaries = {
+                        let mut res = ArrayVec::from_iter([tree.summary().clone()]);
+                        res.extend(child_summaries.drain(..TREE_BASE));
+                        res
+                    };
+                    let tree = SumTree(Arc::new(Node::Internal {
+                        height: *height,
+                        summary: sum(new_child_summaries.iter(), cx),
+                        child_summaries: new_child_summaries,
+                        child_trees: {
+                            let mut res = ArrayVec::from_iter([tree]);
+                            res.extend(child_trees.drain(..TREE_BASE));
+                            res
+                        },
+                    }));
+
+                    *summary = sum(child_summaries.iter(), cx);
+                    Some(tree)
+                }
+            } else {
+                None
+            }
+        }
+    }
+
+    // Merge two nodes into `large`.
+    //
+    // `large` will contain the contents of `small` followed by its own data.
+    // If the combined data exceed the node capacity, returns a new node that
+    // holds the first half of the merged items and `large` is left with the
+    // second half
+    //
+    // The nodes must be on the same height
+    // It only makes sense to call this when `small` is underflowing
+    fn merge_into_right(
+        small: Self,
+        large: &mut Self,
+        cx: <<T as Item>::Summary as Summary>::Context<'_>,
+    ) -> Option<SumTree<T>> {
+        debug_assert_eq!(small.0.height(), large.0.height());
+        match (small.0.as_ref(), Arc::make_mut(&mut large.0)) {
+            (
+                Node::Internal {
+                    summary: small_summary,
+                    child_summaries: small_child_summaries,
+                    child_trees: small_child_trees,
+                    ..
+                },
+                Node::Internal {
+                    summary,
+                    child_summaries,
+                    child_trees,
+                    height,
+                },
+            ) => {
+                let total_child_count = child_trees.len() + small_child_trees.len();
+                if total_child_count <= 2 * TREE_BASE {
+                    let mut all_trees = small_child_trees.clone();
+                    all_trees.extend(child_trees.drain(..));
+                    *child_trees = all_trees;
+
+                    let mut all_summaries = small_child_summaries.clone();
+                    all_summaries.extend(child_summaries.drain(..));
+                    *child_summaries = all_summaries;
+
+                    let mut full_summary = small_summary.clone();
+                    Summary::add_summary(&mut full_summary, summary, cx);
+                    *summary = full_summary;
+                    None
+                } else {
+                    let midpoint = total_child_count.div_ceil(2);
+                    let mut all_trees = small_child_trees.iter().chain(child_trees.iter()).cloned();
+                    let left_trees = all_trees.by_ref().take(midpoint).collect();
+                    *child_trees = all_trees.collect();
+
+                    let mut all_summaries = small_child_summaries
+                        .iter()
+                        .chain(child_summaries.iter())
+                        .cloned();
+                    let left_summaries: ArrayVec<_, { 2 * TREE_BASE }> =
+                        all_summaries.by_ref().take(midpoint).collect();
+                    *child_summaries = all_summaries.collect();
+
+                    *summary = sum(child_summaries.iter(), cx);
+                    Some(SumTree(Arc::new(Node::Internal {
+                        height: *height,
+                        summary: sum(left_summaries.iter(), cx),
+                        child_summaries: left_summaries,
+                        child_trees: left_trees,
+                    })))
+                }
+            }
+            (
+                Node::Leaf {
+                    summary: small_summary,
+                    items: small_items,
+                    item_summaries: small_item_summaries,
+                },
+                Node::Leaf {
+                    summary,
+                    items,
+                    item_summaries,
+                },
+            ) => {
+                let total_child_count = small_items.len() + items.len();
+                if total_child_count <= 2 * TREE_BASE {
+                    let mut all_items = small_items.clone();
+                    all_items.extend(items.drain(..));
+                    *items = all_items;
+
+                    let mut all_summaries = small_item_summaries.clone();
+                    all_summaries.extend(item_summaries.drain(..));
+                    *item_summaries = all_summaries;
+
+                    let mut full_summary = small_summary.clone();
+                    Summary::add_summary(&mut full_summary, summary, cx);
+                    *summary = full_summary;
+                    None
+                } else {
+                    let midpoint = total_child_count.div_ceil(2);
+                    let mut all_items = small_items.iter().chain(items.iter()).cloned();
+                    let left_items = all_items.by_ref().take(midpoint).collect();
+                    *items = all_items.collect();
+
+                    let mut all_summaries = small_item_summaries
+                        .iter()
+                        .chain(item_summaries.iter())
+                        .cloned();
+                    let left_summaries: ArrayVec<_, { 2 * TREE_BASE }> =
+                        all_summaries.by_ref().take(midpoint).collect();
+                    *item_summaries = all_summaries.collect();
+
+                    *summary = sum(item_summaries.iter(), cx);
+                    Some(SumTree(Arc::new(Node::Leaf {
+                        items: left_items,
+                        summary: sum(left_summaries.iter(), cx),
+                        item_summaries: left_summaries,
+                    })))
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
+
     fn from_child_trees(
         left: SumTree<T>,
         right: SumTree<T>,