1 files changed, 137 insertions, 66 deletions

diff --git a/absl/container/internal/btree.h b/absl/container/internal/btree.h
index 047055a2..996afa61 100644
--- a/absl/container/internal/btree.h
+++ b/absl/container/internal/btree.h
@@ -255,6 +255,10 @@ struct common_params {
   static void move(Alloc *alloc, slot_type *src, slot_type *dest) {
     slot_policy::move(alloc, src, dest);
   }
+  static void move(Alloc *alloc, slot_type *first, slot_type *last,
+                   slot_type *result) {
+    slot_policy::move(alloc, first, last, result);
+  }
 };
 
 // A parameters structure for holding the type parameters for a btree_map.
@@ -332,6 +336,13 @@ struct set_slot_policy {
   static void move(Alloc * /*alloc*/, slot_type *src, slot_type *dest) {
     *dest = std::move(*src);
   }
+
+  template <typename Alloc>
+  static void move(Alloc *alloc, slot_type *first, slot_type *last,
+                   slot_type *result) {
+    for (slot_type *src = first, *dest = result; src != last; ++src, ++dest)
+      move(alloc, src, dest);
+  }
 };
 
 // A parameters structure for holding the type parameters for a btree_set.
@@ -748,10 +759,14 @@ class btree_node {
   template <typename... Args>
   void emplace_value(size_type i, allocator_type *alloc, Args &&... args);
 
-  // Removes the values at positions [i, i + to_erase), shifting all existing
-  // values and children after that range to the left by to_erase. Clears all
-  // children between [i, i + to_erase).
-  void remove_values(field_type i, field_type to_erase, allocator_type *alloc);
+  // Removes the value at position i, shifting all existing values and children
+  // at positions > i to the left by 1.
+  void remove_value(int i, allocator_type *alloc);
+
+  // Removes the values at positions [i, i + to_erase), shifting all values
+  // after that range to the left by to_erase. Does not change children at all.
+  void remove_values_ignore_children(int i, int to_erase,
+                                     allocator_type *alloc);
 
   // Rebalances a node with its right sibling.
   void rebalance_right_to_left(int to_move, btree_node *right,
@@ -763,7 +778,7 @@ class btree_node {
   void split(int insert_position, btree_node *dest, allocator_type *alloc);
 
   // Merges a node with its right sibling, moving all of the values and the
-  // delimiting key in the parent node onto itself, and deleting the src node.
+  // delimiting key in the parent node onto itself.
   void merge(btree_node *src, allocator_type *alloc);
 
   // Node allocation/deletion routines.
@@ -784,27 +799,9 @@ class btree_node {
     absl::container_internal::SanitizerPoisonMemoryRegion(
         &mutable_child(start()), (kNodeValues + 1) * sizeof(btree_node *));
   }
-
-  // Deletes a node and all of its children.
-  // TODO(ezb): don't use recursion here to avoid potential stack overflows.
-  static void clear_and_delete(btree_node *node, allocator_type *alloc) {
-    const field_type start = node->start();
-    const field_type finish = node->finish();
-    for (field_type i = start; i < finish; ++i) {
-      node->value_destroy(i, alloc);
-    }
-    if (node->leaf()) {
-      absl::container_internal::Deallocate<Alignment()>(
-          alloc, node, LeafSize(node->max_count()));
-    } else {
-      // If the node is empty, then there are no children so don't try clearing.
-      if (start < finish) {
-        for (field_type i = start; i <= finish; ++i) {
-          clear_and_delete(node->child(i), alloc);
-        }
-      }
-      absl::container_internal::Deallocate<Alignment()>(alloc, node,
-                                                        InternalSize());
+  void destroy(allocator_type *alloc) {
+    for (int i = start(); i < finish(); ++i) {
+      value_destroy(i, alloc);
     }
   }
 
@@ -1426,8 +1423,25 @@ class btree {
   }
 
   // Deletion helper routines.
+  void erase_same_node(iterator begin, iterator end);
+  iterator erase_from_leaf_node(iterator begin, size_type to_erase);
   iterator rebalance_after_delete(iterator iter);
 
+  // Deallocates a node of a certain size in bytes using the allocator.
+  void deallocate(const size_type size, node_type *node) {
+    absl::container_internal::Deallocate<node_type::Alignment()>(
+        mutable_allocator(), node, size);
+  }
+
+  void delete_internal_node(node_type *node) {
+    node->destroy(mutable_allocator());
+    deallocate(node_type::InternalSize(), node);
+  }
+  void delete_leaf_node(node_type *node) {
+    node->destroy(mutable_allocator());
+    deallocate(node_type::LeafSize(node->max_count()), node);
+  }
+
   // Rebalances or splits the node iter points to.
   void rebalance_or_split(iterator *iter);
 
@@ -1496,6 +1510,9 @@ class btree {
   template <typename K>
   iterator internal_find(const K &key) const;
 
+  // Deletes a node and all of its children.
+  void internal_clear(node_type *node);
+
   // Verifies the tree structure of node.
   int internal_verify(const node_type *node, const key_type *lo,
                       const key_type *hi) const;
@@ -1563,29 +1580,26 @@ inline void btree_node<P>::emplace_value(const size_type i,
 }
 
 template <typename P>
-inline void btree_node<P>::remove_values(const field_type i,
-                                         const field_type to_erase,
-                                         allocator_type *alloc) {
-  // Transfer values after the removed range into their new places.
-  const field_type orig_finish = finish();
-  const field_type src_i = i + to_erase;
-  for (field_type j = i; j < src_i; ++j) {
-    value_destroy(j, alloc);
+inline void btree_node<P>::remove_value(const int i, allocator_type *alloc) {
+  if (!leaf() && finish() > i + 1) {
+    assert(child(i + 1)->count() == 0);
+    for (size_type j = i + 1; j < finish(); ++j) {
+      set_child(j, child(j + 1));
+    }
+    clear_child(finish());
   }
-  transfer_n(orig_finish - src_i, i, src_i, this, alloc);
 
-  if (!leaf()) {
-    // Delete all children between begin and end.
-    for (field_type j = 0; j < to_erase; ++j) {
-      clear_and_delete(child(i + j + 1), alloc);
-    }
-    // Rotate children after end into new positions.
-    for (field_type j = i + to_erase + 1; j <= orig_finish; ++j) {
-      set_child(j - to_erase, child(j));
-      clear_child(j);
-    }
+  remove_values_ignore_children(i, /*to_erase=*/1, alloc);
+}
+
+template <typename P>
+inline void btree_node<P>::remove_values_ignore_children(
+    const int i, const int to_erase, allocator_type *alloc) {
+  params_type::move(alloc, slot(i + to_erase), finish_slot(), slot(i));
+  for (int j = finish() - to_erase; j < finish(); ++j) {
+    value_destroy(j, alloc);
   }
-  set_finish(orig_finish - to_erase);
+  set_finish(finish() - to_erase);
 }
 
 template <typename P>
@@ -1737,8 +1751,8 @@ void btree_node<P>::merge(btree_node *src, allocator_type *alloc) {
   set_finish(start() + 1 + count() + src->count());
   src->set_finish(src->start());
 
-  // Remove the value on the parent node and delete the src node.
-  parent()->remove_values(position(), /*to_erase=*/1, alloc);
+  // Remove the value on the parent node.
+  parent()->remove_value(position(), alloc);
 }
 
 ////
@@ -2020,7 +2034,7 @@ auto btree<P>::erase(iterator iter) -> iterator {
   bool internal_delete = false;
   if (!iter.node->leaf()) {
     // Deletion of a value on an internal node. First, move the largest value
-    // from our left child here, then delete that position (in remove_values()
+    // from our left child here, then delete that position (in remove_value()
     // below). We can get to the largest value from our left child by
     // decrementing iter.
     iterator internal_iter(iter);
@@ -2032,7 +2046,7 @@ auto btree<P>::erase(iterator iter) -> iterator {
   }
 
   // Delete the key from the leaf.
-  iter.node->remove_values(iter.position, /*to_erase=*/1, mutable_allocator());
+  iter.node->remove_value(iter.position, mutable_allocator());
   --size_;
 
   // We want to return the next value after the one we just erased. If we
@@ -2107,9 +2121,7 @@ auto btree<P>::erase_range(iterator begin, iterator end)
   }
 
   if (begin.node == end.node) {
-    assert(end.position > begin.position);
-    begin.node->remove_values(begin.position, end.position - begin.position,
-                              mutable_allocator());
+    erase_same_node(begin, end);
     size_ -= count;
     return {count, rebalance_after_delete(begin)};
   }
@@ -2119,11 +2131,8 @@ auto btree<P>::erase_range(iterator begin, iterator end)
     if (begin.node->leaf()) {
       const size_type remaining_to_erase = size_ - target_size;
       const size_type remaining_in_node = begin.node->finish() - begin.position;
-      const size_type to_erase =
-          (std::min)(remaining_to_erase, remaining_in_node);
-      begin.node->remove_values(begin.position, to_erase, mutable_allocator());
-      size_ -= to_erase;
-      begin = rebalance_after_delete(begin);
+      begin = erase_from_leaf_node(
+          begin, (std::min)(remaining_to_erase, remaining_in_node));
     } else {
       begin = erase(begin);
     }
@@ -2132,6 +2141,51 @@ auto btree<P>::erase_range(iterator begin, iterator end)
 }
 
 template <typename P>
+void btree<P>::erase_same_node(iterator begin, iterator end) {
+  assert(begin.node == end.node);
+  assert(end.position > begin.position);
+
+  node_type *node = begin.node;
+  size_type to_erase = end.position - begin.position;
+  if (!node->leaf()) {
+    // Delete all children between begin and end.
+    for (size_type i = 0; i < to_erase; ++i) {
+      internal_clear(node->child(begin.position + i + 1));
+    }
+    // Rotate children after end into new positions.
+    for (size_type i = begin.position + to_erase + 1; i <= node->finish();
+         ++i) {
+      node->set_child(i - to_erase, node->child(i));
+      node->clear_child(i);
+    }
+  }
+  node->remove_values_ignore_children(begin.position, to_erase,
+                                      mutable_allocator());
+
+  // Do not need to update rightmost_, because
+  // * either end == this->end(), and therefore node == rightmost_, and still
+  //   exists
+  // * or end != this->end(), and therefore rightmost_ hasn't been erased, since
+  //   it wasn't covered in [begin, end)
+}
+
+template <typename P>
+auto btree<P>::erase_from_leaf_node(iterator begin, size_type to_erase)
+    -> iterator {
+  node_type *node = begin.node;
+  assert(node->leaf());
+  assert(node->finish() > begin.position);
+  assert(begin.position + to_erase <= node->finish());
+
+  node->remove_values_ignore_children(begin.position, to_erase,
+                                      mutable_allocator());
+
+  size_ -= to_erase;
+
+  return rebalance_after_delete(begin);
+}
+
+template <typename P>
 template <typename K>
 auto btree<P>::erase_unique(const K &key) -> size_type {
   const iterator iter = internal_find(key);
@@ -2159,7 +2213,7 @@ auto btree<P>::erase_multi(const K &key) -> size_type {
 template <typename P>
 void btree<P>::clear() {
   if (!empty()) {
-    node_type::clear_and_delete(root(), mutable_allocator());
+    internal_clear(root());
   }
   mutable_root() = EmptyNode();
   rightmost_ = EmptyNode();
@@ -2300,7 +2354,12 @@ void btree<P>::rebalance_or_split(iterator *iter) {
 template <typename P>
 void btree<P>::merge_nodes(node_type *left, node_type *right) {
   left->merge(right, mutable_allocator());
-  if (rightmost_ == right) rightmost_ = left;
+  if (right->leaf()) {
+    if (rightmost_ == right) rightmost_ = left;
+    delete_leaf_node(right);
+  } else {
+    delete_internal_node(right);
+  }
 }
 
 template <typename P>
@@ -2357,20 +2416,20 @@ bool btree<P>::try_merge_or_rebalance(iterator *iter) {
 
 template <typename P>
 void btree<P>::try_shrink() {
-  node_type *orig_root = root();
-  if (orig_root->count() > 0) {
+  if (root()->count() > 0) {
     return;
   }
   // Deleted the last item on the root node, shrink the height of the tree.
-  if (orig_root->leaf()) {
+  if (root()->leaf()) {
     assert(size() == 0);
+    delete_leaf_node(root());
     mutable_root() = rightmost_ = EmptyNode();
   } else {
-    node_type *child = orig_root->start_child();
+    node_type *child = root()->start_child();
     child->make_root();
+    delete_internal_node(root());
     mutable_root() = child;
   }
-  node_type::clear_and_delete(orig_root, mutable_allocator());
 }
 
 template <typename P>
@@ -2415,7 +2474,7 @@ inline auto btree<P>::internal_emplace(iterator iter, Args &&... args)
                            old_root->start(), old_root, alloc);
       new_root->set_finish(old_root->finish());
       old_root->set_finish(old_root->start());
-      node_type::clear_and_delete(old_root, alloc);
+      delete_leaf_node(old_root);
       mutable_root() = rightmost_ = new_root;
     } else {
       rebalance_or_split(&iter);
@@ -2519,6 +2578,18 @@ auto btree<P>::internal_find(const K &key) const -> iterator {
 }
 
 template <typename P>
+void btree<P>::internal_clear(node_type *node) {
+  if (!node->leaf()) {
+    for (int i = node->start(); i <= node->finish(); ++i) {
+      internal_clear(node->child(i));
+    }
+    delete_internal_node(node);
+  } else {
+    delete_leaf_node(node);
+  }
+}
+
+template <typename P>
 int btree<P>::internal_verify(const node_type *node, const key_type *lo,
                               const key_type *hi) const {
   assert(node->count() > 0);