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
diff --git a/absl/container/btree_map.h b/absl/container/btree_map.h
index 819a925..cd3ee2b 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 @@
// 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 @@
// 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
diff --git a/absl/container/btree_set.h b/absl/container/btree_set.h
index d93bdbf..51dc42b 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 @@
// 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 @@
// 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
diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index 404ccde..28dda8a 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -2123,6 +2123,79 @@
}
}
+TEST(Btree, ExtractAndGetNextSet) {
+ absl::btree_set<int> src = {1, 2, 3, 4, 5};
+ auto it = src.find(3);
+ auto extracted_and_next = src.extract_and_get_next(it);
+ EXPECT_THAT(src, ElementsAre(1, 2, 4, 5));
+ EXPECT_EQ(extracted_and_next.node.value(), 3);
+ EXPECT_EQ(*extracted_and_next.next, 4);
+}
+
+TEST(Btree, ExtractAndGetNextMultiSet) {
+ absl::btree_multiset<int> src = {1, 2, 3, 4, 5};
+ auto it = src.find(3);
+ auto extracted_and_next = src.extract_and_get_next(it);
+ EXPECT_THAT(src, ElementsAre(1, 2, 4, 5));
+ EXPECT_EQ(extracted_and_next.node.value(), 3);
+ EXPECT_EQ(*extracted_and_next.next, 4);
+}
+
+TEST(Btree, ExtractAndGetNextMap) {
+ absl::btree_map<int, int> src = {{1, 2}, {3, 4}, {5, 6}};
+ auto it = src.find(3);
+ auto extracted_and_next = src.extract_and_get_next(it);
+ EXPECT_THAT(src, ElementsAre(Pair(1, 2), Pair(5, 6)));
+ EXPECT_EQ(extracted_and_next.node.key(), 3);
+ EXPECT_EQ(extracted_and_next.node.mapped(), 4);
+ EXPECT_THAT(*extracted_and_next.next, Pair(5, 6));
+}
+
+TEST(Btree, ExtractAndGetNextMultiMap) {
+ absl::btree_multimap<int, int> src = {{1, 2}, {3, 4}, {5, 6}};
+ auto it = src.find(3);
+ auto extracted_and_next = src.extract_and_get_next(it);
+ EXPECT_THAT(src, ElementsAre(Pair(1, 2), Pair(5, 6)));
+ EXPECT_EQ(extracted_and_next.node.key(), 3);
+ EXPECT_EQ(extracted_and_next.node.mapped(), 4);
+ EXPECT_THAT(*extracted_and_next.next, Pair(5, 6));
+}
+
+TEST(Btree, ExtractAndGetNextEndIter) {
+ absl::btree_set<int> src = {1, 2, 3, 4, 5};
+ auto it = src.find(5);
+ auto extracted_and_next = src.extract_and_get_next(it);
+ EXPECT_THAT(src, ElementsAre(1, 2, 3, 4));
+ EXPECT_EQ(extracted_and_next.node.value(), 5);
+ EXPECT_EQ(extracted_and_next.next, src.end());
+}
+
+TEST(Btree, ExtractDoesntCauseExtraMoves) {
+#ifdef _MSC_VER
+ GTEST_SKIP() << "This test fails on MSVC.";
+#endif
+
+ using Set = absl::btree_set<MovableOnlyInstance>;
+ std::array<std::function<void(Set &)>, 3> extracters = {
+ [](Set &s) { auto node = s.extract(s.begin()); },
+ [](Set &s) { auto ret = s.extract_and_get_next(s.begin()); },
+ [](Set &s) { auto node = s.extract(MovableOnlyInstance(0)); }};
+
+ InstanceTracker tracker;
+ for (int i = 0; i < 3; ++i) {
+ Set s;
+ s.insert(MovableOnlyInstance(0));
+ tracker.ResetCopiesMovesSwaps();
+
+ extracters[i](s);
+ // We expect to see exactly 1 move: from the original slot into the
+ // extracted node.
+ EXPECT_EQ(tracker.copies(), 0) << i;
+ EXPECT_EQ(tracker.moves(), 1) << i;
+ EXPECT_EQ(tracker.swaps(), 0) << i;
+ }
+}
+
// For multisets, insert with hint also affects correctness because we need to
// insert immediately before the hint if possible.
struct InsertMultiHintData {
diff --git a/absl/container/internal/btree_container.h b/absl/container/internal/btree_container.h
index 3e25986..2bff11d 100644
--- a/absl/container/internal/btree_container.h
+++ b/absl/container/internal/btree_container.h
@@ -65,6 +65,11 @@
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 @@
}
// 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.
