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, 32 insertions, 2 deletions

diff --git a/absl/container/btree_set.h b/absl/container/btree_set.h
index d93bdbf6..51dc42b7 100644
--- a/absl/container/btree_set.h
+++ b/absl/container/btree_set.h
@@ -43,8 +43,8 @@
 // pointers, and references to elements. Such invalidations are typically only
 // an issue if insertion and deletion operations are interleaved with the use of
 // more than one iterator, pointer, or reference simultaneously. For this
-// reason, `insert()` and `erase()` return a valid iterator at the current
-// position (and `extract()` cannot be used in this way).
+// reason, `insert()`, `erase()`, and `extract_and_get_next()` return a valid
+// iterator at the current position.
 //
 // Another API difference is that btree iterators can be subtracted, and this
 // is faster than using std::distance.
@@ -293,6 +293,21 @@ class btree_set
   // It does NOT refer to the data layout of the underlying btree.
   using Base::extract;
 
+  // btree_set::extract_and_get_next()
+  //
+  // Extracts the indicated element, erasing it in the process, and returns it
+  // as a C++17-compatible node handle along with an iterator to the next
+  // element.
+  //
+  // extract_and_get_next_return_type extract_and_get_next(
+  //     const_iterator position):
+  //
+  //   Extracts the element at the indicated position, returns a struct
+  //   containing a member named `node`: a node handle owning that extracted
+  //   data and a member named `next`: an iterator pointing to the next element
+  //   in the btree.
+  using Base::extract_and_get_next;
+
   // btree_set::merge()
   //
   // Extracts elements from a given `source` btree_set into this
@@ -615,6 +630,21 @@ class btree_multiset
   // It does NOT refer to the data layout of the underlying btree.
   using Base::extract;
 
+  // btree_multiset::extract_and_get_next()
+  //
+  // Extracts the indicated element, erasing it in the process, and returns it
+  // as a C++17-compatible node handle along with an iterator to the next
+  // element.
+  //
+  // extract_and_get_next_return_type extract_and_get_next(
+  //     const_iterator position):
+  //
+  //   Extracts the element at the indicated position, returns a struct
+  //   containing a member named `node`: a node handle owning that extracted
+  //   data and a member named `next`: an iterator pointing to the next element
+  //   in the btree.
+  using Base::extract_and_get_next;
+
   // btree_multiset::merge()
   //
   // Extracts all elements from a given `source` btree_multiset into this