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, 34 insertions, 4 deletions

diff --git a/absl/container/btree_map.h b/absl/container/btree_map.h
index 819a925f..cd3ee2b4 100644
--- a/absl/container/btree_map.h
+++ b/absl/container/btree_map.h
@@ -42,10 +42,10 @@
 // Importantly, insertions and deletions may invalidate outstanding iterators,
 // 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). Another important
-// difference is that key-types must be copy-constructible.
+// more than one iterator, pointer, or reference simultaneously.  For this
+// reason, `insert()`, `erase()`, and `extract_and_get_next()` return a valid
+// iterator at the current position. Another important difference is that
+// key-types must be copy-constructible.
 //
 // Another API difference is that btree iterators can be subtracted, and this
 // is faster than using std::distance.
@@ -351,6 +351,21 @@ class btree_map
   // It does NOT refer to the data layout of the underlying btree.
   using Base::extract;
 
+  // btree_map::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_map::merge()
   //
   // Extracts elements from a given `source` btree_map into this
@@ -702,6 +717,21 @@ class btree_multimap
   // It does NOT refer to the data layout of the underlying btree.
   using Base::extract;
 
+  // btree_multimap::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_multimap::merge()
   //
   // Extracts all elements from a given `source` btree_multimap into this