diff options
Diffstat (limited to 'absl/container')
-rw-r--r-- | absl/container/btree_test.cc | 48 | ||||
-rw-r--r-- | absl/container/internal/btree.h | 39 | ||||
-rw-r--r-- | absl/container/internal/btree_container.h | 17 |
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; |