summaryrefslogtreecommitdiff
path: root/absl/container
diff options
context:
space:
mode:
Diffstat (limited to 'absl/container')
-rw-r--r--absl/container/btree_test.cc48
-rw-r--r--absl/container/internal/btree.h39
-rw-r--r--absl/container/internal/btree_container.h17
3 files changed, 83 insertions, 21 deletions
diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index a933386a..367d75be 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -2038,6 +2038,30 @@ TEST(Btree, ExtractAndInsertNodeHandleMultiMap) {
EXPECT_EQ(res, ++other.begin());
}
+TEST(Btree, ExtractMultiMapEquivalentKeys) {
+ // Note: using string keys means a three-way comparator.
+ absl::btree_multimap<std::string, int> map;
+ for (int i = 0; i < 100; ++i) {
+ for (int j = 0; j < 100; ++j) {
+ map.insert({absl::StrCat(i), j});
+ }
+ }
+
+ for (int i = 0; i < 100; ++i) {
+ const std::string key = absl::StrCat(i);
+ auto node_handle = map.extract(key);
+ EXPECT_EQ(node_handle.key(), key);
+ EXPECT_EQ(node_handle.mapped(), 0) << i;
+ }
+
+ for (int i = 0; i < 100; ++i) {
+ const std::string key = absl::StrCat(i);
+ auto node_handle = map.extract(key);
+ EXPECT_EQ(node_handle.key(), key);
+ EXPECT_EQ(node_handle.mapped(), 1) << i;
+ }
+}
+
// For multisets, insert with hint also affects correctness because we need to
// insert immediately before the hint if possible.
struct InsertMultiHintData {
@@ -2726,7 +2750,6 @@ TYPED_TEST_SUITE(BtreeMultiKeyTest, MultiKeyComps);
TYPED_TEST(BtreeMultiKeyTest, EqualRange) {
absl::btree_set<MultiKey, TypeParam> set;
-
for (int i = 0; i < 100; ++i) {
for (int j = 0; j < 100; ++j) {
set.insert({i, j});
@@ -2736,11 +2759,32 @@ TYPED_TEST(BtreeMultiKeyTest, EqualRange) {
for (int i = 0; i < 100; ++i) {
auto equal_range = set.equal_range(i);
EXPECT_EQ(equal_range.first->i1, i);
- EXPECT_EQ(equal_range.first->i2, 0);
+ EXPECT_EQ(equal_range.first->i2, 0) << i;
EXPECT_EQ(std::distance(equal_range.first, equal_range.second), 100) << i;
}
}
+TYPED_TEST(BtreeMultiKeyTest, Extract) {
+ absl::btree_set<MultiKey, TypeParam> set;
+ for (int i = 0; i < 100; ++i) {
+ for (int j = 0; j < 100; ++j) {
+ set.insert({i, j});
+ }
+ }
+
+ for (int i = 0; i < 100; ++i) {
+ auto node_handle = set.extract(i);
+ EXPECT_EQ(node_handle.value().i1, i);
+ EXPECT_EQ(node_handle.value().i2, 0) << i;
+ }
+
+ for (int i = 0; i < 100; ++i) {
+ auto node_handle = set.extract(i);
+ EXPECT_EQ(node_handle.value().i1, i);
+ EXPECT_EQ(node_handle.value().i2, 1) << i;
+ }
+}
+
TYPED_TEST(BtreeMultiKeyTest, Erase) {
absl::btree_set<MultiKey, TypeParam> set = {
{1, 1}, {2, 1}, {2, 2}, {3, 1}};
diff --git a/absl/container/internal/btree.h b/absl/container/internal/btree.h
index d863cb30..6f5f01b8 100644
--- a/absl/container/internal/btree.h
+++ b/absl/container/internal/btree.h
@@ -1211,7 +1211,7 @@ class btree {
return const_reverse_iterator(begin());
}
- // Finds the first element whose key is not less than key.
+ // Finds the first element whose key is not less than `key`.
template <typename K>
iterator lower_bound(const K &key) {
return internal_end(internal_lower_bound(key).value);
@@ -1221,7 +1221,12 @@ class btree {
return internal_end(internal_lower_bound(key).value);
}
- // Finds the first element whose key is greater than key.
+ // Finds the first element whose key is not less than `key` and also returns
+ // whether that element is equal to `key`.
+ template <typename K>
+ std::pair<iterator, bool> lower_bound_equal(const K &key) const;
+
+ // Finds the first element whose key is greater than `key`.
template <typename K>
iterator upper_bound(const K &key) {
return internal_end(internal_upper_bound(key));
@@ -1303,8 +1308,8 @@ class btree {
// to the element after the last erased element.
std::pair<size_type, iterator> erase_range(iterator begin, iterator end);
- // Finds the iterator corresponding to a key or returns end() if the key is
- // not present.
+ // Finds an element with key equivalent to `key` or returns `end()` if `key`
+ // is not present.
template <typename K>
iterator find(const K &key) {
return internal_end(internal_find(key));
@@ -1905,12 +1910,23 @@ constexpr bool btree<P>::static_assert_validation() {
template <typename P>
template <typename K>
-auto btree<P>::equal_range(const K &key) -> std::pair<iterator, iterator> {
+auto btree<P>::lower_bound_equal(const K &key) const
+ -> std::pair<iterator, bool> {
const SearchResult<iterator, is_key_compare_to::value> res =
internal_lower_bound(key);
- const iterator lower = internal_end(res.value);
- if (res.HasMatch() ? !res.IsEq()
- : lower == end() || compare_keys(key, lower.key())) {
+ const iterator lower = iterator(internal_end(res.value));
+ const bool equal = res.HasMatch()
+ ? res.IsEq()
+ : lower != end() && !compare_keys(key, lower.key());
+ return {lower, equal};
+}
+
+template <typename P>
+template <typename K>
+auto btree<P>::equal_range(const K &key) -> std::pair<iterator, iterator> {
+ const std::pair<iterator, bool> lower_and_equal = lower_bound_equal(key);
+ const iterator lower = lower_and_equal.first;
+ if (!lower_and_equal.second) {
return {lower, lower};
}
@@ -2510,14 +2526,17 @@ template <typename P>
template <typename K>
auto btree<P>::internal_lower_bound(const K &key) const
-> SearchResult<iterator, is_key_compare_to::value> {
+ if (!params_type::template can_have_multiple_equivalent_keys<K>()) {
+ SearchResult<iterator, is_key_compare_to::value> ret = internal_locate(key);
+ ret.value = internal_last(ret.value);
+ return ret;
+ }
iterator iter(const_cast<node_type *>(root()));
SearchResult<int, is_key_compare_to::value> res;
bool seen_eq = false;
for (;;) {
res = iter.node->lower_bound(key, key_comp());
iter.position = res.value;
- // TODO(ezb): we should be able to terminate early on IsEq() if there can't
- // be multiple equivalent keys in container for this lookup type.
if (iter.node->leaf()) {
break;
}
diff --git a/absl/container/internal/btree_container.h b/absl/container/internal/btree_container.h
index 887eda41..03be708e 100644
--- a/absl/container/internal/btree_container.h
+++ b/absl/container/internal/btree_container.h
@@ -338,13 +338,12 @@ class btree_set_container : public btree_container<Tree> {
}
// Node extraction routines.
- // TODO(ezb): when the comparator is heterogeneous and has different
- // equivalence classes for different lookup types, we should extract the first
- // equivalent value if there are multiple.
template <typename K = key_type>
node_type extract(const key_arg<K> &key) {
- auto it = this->find(key);
- return it == this->end() ? node_type() : extract(it);
+ const std::pair<iterator, bool> lower_and_equal =
+ this->tree_.lower_bound_equal(key);
+ return lower_and_equal.second ? extract(lower_and_equal.first)
+ : node_type();
}
using super_type::extract;
@@ -621,12 +620,12 @@ class btree_multiset_container : public btree_container<Tree> {
}
// Node extraction routines.
- // TODO(ezb): we are supposed to extract the first equivalent key if there are
- // multiple, but this isn't guaranteed to extract the first one.
template <typename K = key_type>
node_type extract(const key_arg<K> &key) {
- auto it = this->find(key);
- return it == this->end() ? node_type() : extract(it);
+ const std::pair<iterator, bool> lower_and_equal =
+ this->tree_.lower_bound_equal(key);
+ return lower_and_equal.second ? extract(lower_and_equal.first)
+ : node_type();
}
using super_type::extract;