Add a new API for `extract_and_get_next()` in b-tree that returns both the extracted node and an iterator to the next element in the container.

Motivation: it can be useful, when calling `extract` to maintain an iterator next to the location of the extracted element. `std::set`, et al. allow for this because they have iterator stability. `absl::{flat,node}_hash_{set,map}` allow for this because they are guaranteed not to rehash when elements are removed so they also have iterator stability across calls to `extract()`. But b-tree doesn't support this use case without this API because removing elements can cause rebalancing, which invalidates iterators. We can get the next iterator without this API using `lower_bound(node_handle.value())`, but that requires an extra lookup.
PiperOrigin-RevId: 488721247
Change-Id: Id66f17311bf53678f536e4e4f070775f5ce0c542

1 files changed, 14 insertions, 0 deletions

diff --git a/absl/container/internal/btree_container.h b/absl/container/internal/btree_container.h
index 3e259861..2bff11db 100644
--- a/absl/container/internal/btree_container.h
+++ b/absl/container/internal/btree_container.h
@@ -65,6 +65,11 @@ class btree_container {
   using const_reverse_iterator = typename Tree::const_reverse_iterator;
   using node_type = typename Tree::node_handle_type;
 
+  struct extract_and_get_next_return_type {
+    node_type node;
+    iterator next;
+  };
+
   // Constructors/assignments.
   btree_container() : tree_(key_compare(), allocator_type()) {}
   explicit btree_container(const key_compare &comp,
@@ -165,6 +170,15 @@ class btree_container {
   }
 
   // Extract routines.
+  extract_and_get_next_return_type extract_and_get_next(
+      const_iterator position) {
+    // Use Construct instead of Transfer because the rebalancing code will
+    // destroy the slot later.
+    // Note: we rely on erase() taking place after Construct().
+    return {CommonAccess::Construct<node_type>(get_allocator(),
+                                               iterator(position).slot()),
+            erase(position)};
+  }
   node_type extract(iterator position) {
     // Use Construct instead of Transfer because the rebalancing code will
     // destroy the slot later.