@@ -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>,