diff options
-rw-r--r-- | source/blender/blenlib/BLI_map.hh | 43 | ||||
-rw-r--r-- | tests/gtests/blenlib/BLI_map_test.cc | 20 |
2 files changed, 61 insertions, 2 deletions
diff --git a/source/blender/blenlib/BLI_map.hh b/source/blender/blenlib/BLI_map.hh index 13f1a967d65..9737367ebca 100644 --- a/source/blender/blenlib/BLI_map.hh +++ b/source/blender/blenlib/BLI_map.hh @@ -406,6 +406,28 @@ class Map { } /** + * Get the value that corresponds to the given key and remove it from the map. If the key is not + * in the map, return the given default value instead. + */ + Value pop_default(const Key &key, const Value &default_value) + { + return this->pop_default_as(key, default_value); + } + Value pop_default(const Key &key, Value &&default_value) + { + return this->pop_default_as(key, std::move(default_value)); + } + + /** + * Same as `pop_default`, but accepts other key types that are supported by the hash function. + */ + template<typename ForwardKey, typename ForwardValue> + Value pop_default_as(const ForwardKey &key, ForwardValue &&default_value) + { + return this->pop_default__impl(key, std::forward<ForwardValue>(default_value), m_hash(key)); + } + + /** * This method can be used to implement more complex custom behavior without having to do * multiple lookups * @@ -1039,7 +1061,7 @@ class Map { MAP_SLOT_PROBING_BEGIN (hash, slot) { if (slot.contains(key, m_is_equal, hash)) { - Value value = *slot.value(); + Value value = std::move(*slot.value()); slot.remove(); return value; } @@ -1051,7 +1073,7 @@ class Map { { MAP_SLOT_PROBING_BEGIN (hash, slot) { if (slot.contains(key, m_is_equal, hash)) { - Optional<Value> value = *slot.value(); + Optional<Value> value = std::move(*slot.value()); slot.remove(); m_removed_slots++; return value; @@ -1063,6 +1085,23 @@ class Map { MAP_SLOT_PROBING_END(); } + template<typename ForwardKey, typename ForwardValue> + Value pop_default__impl(const ForwardKey &key, ForwardValue &&default_value, uint32_t hash) + { + MAP_SLOT_PROBING_BEGIN (hash, slot) { + if (slot.contains(key, m_is_equal, hash)) { + Value value = std::move(*slot.value()); + slot.remove(); + m_removed_slots++; + return value; + } + if (slot.is_empty()) { + return std::forward<ForwardValue>(default_value); + } + } + MAP_SLOT_PROBING_END(); + } + template<typename ForwardKey, typename CreateValueF, typename ModifyValueF> auto add_or_modify__impl(ForwardKey &&key, const CreateValueF &create_value, diff --git a/tests/gtests/blenlib/BLI_map_test.cc b/tests/gtests/blenlib/BLI_map_test.cc index aa064e2f427..b14ba40ed7c 100644 --- a/tests/gtests/blenlib/BLI_map_test.cc +++ b/tests/gtests/blenlib/BLI_map_test.cc @@ -99,6 +99,25 @@ TEST(map, PopTry) EXPECT_EQ(map.size(), 0); } +TEST(map, PopDefault) +{ + Map<int, int> map; + map.add(1, 4); + map.add(2, 7); + map.add(3, 8); + EXPECT_EQ(map.size(), 3); + EXPECT_EQ(map.pop_default(4, 10), 10); + EXPECT_EQ(map.size(), 3); + EXPECT_EQ(map.pop_default(1, 10), 4); + EXPECT_EQ(map.size(), 2); + EXPECT_EQ(map.pop_default(2, 20), 7); + EXPECT_EQ(map.size(), 1); + EXPECT_EQ(map.pop_default(2, 20), 20); + EXPECT_EQ(map.size(), 1); + EXPECT_EQ(map.pop_default(3, 0), 8); + EXPECT_EQ(map.size(), 0); +} + TEST(map, PopItemMany) { Map<int, int> map; @@ -373,6 +392,7 @@ TEST(map, UniquePtrValue) map.add(6, std::unique_ptr<int>(new int())); map.add_overwrite(7, std::unique_ptr<int>(new int())); map.lookup_or_add(8, std::unique_ptr<int>(new int())); + map.pop_default(9, std::unique_ptr<int>(new int())); EXPECT_EQ(map.lookup(1).get(), value1_ptr); EXPECT_EQ(map.lookup_ptr(100), nullptr); |