diff options
-rw-r--r-- | source/blender/blenlib/BLI_string_map.hh | 21 | ||||
-rw-r--r-- | tests/gtests/blenlib/BLI_string_map_test.cc | 24 |
2 files changed, 45 insertions, 0 deletions
diff --git a/source/blender/blenlib/BLI_string_map.hh b/source/blender/blenlib/BLI_string_map.hh index ed23ea3aaa0..caa7e16d1f3 100644 --- a/source/blender/blenlib/BLI_string_map.hh +++ b/source/blender/blenlib/BLI_string_map.hh @@ -342,6 +342,27 @@ template<typename T, typename Allocator = GuardedAllocator> class StringMap { } /** + * Return the value that corresponds to the given key. + * If it does not exist yet, create and insert it first. + */ + template<typename CreateValueF> T &lookup_or_add(StringRef key, const CreateValueF &create_value) + { + return *this->add_or_modify( + key, + [&](T *value) { return new (value) T(create_value()); }, + [](T *value) { return value; }); + } + + /** + * Return the value that corresponds to the given key. + * If it does not exist yet, insert a new default constructed value and return that. + */ + T &lookup_or_add_default(StringRef key) + { + return this->lookup_or_add(key, []() { return T(); }); + } + + /** * Do a linear search over all items to find a key for a value. */ StringRefNull find_key_for_value(const T &value) const diff --git a/tests/gtests/blenlib/BLI_string_map_test.cc b/tests/gtests/blenlib/BLI_string_map_test.cc index 41cda920a89..6acad0ce581 100644 --- a/tests/gtests/blenlib/BLI_string_map_test.cc +++ b/tests/gtests/blenlib/BLI_string_map_test.cc @@ -249,3 +249,27 @@ TEST(string_map, AddOrModify) EXPECT_FALSE(map.add_or_modify("Hello", create_func, modify_func)); EXPECT_EQ(map.lookup("Hello"), 15); } + +TEST(string_map, LookupOrAdd) +{ + StringMap<int> map; + auto return_5 = []() { return 5; }; + auto return_8 = []() { return 8; }; + + int &a = map.lookup_or_add("A", return_5); + EXPECT_EQ(a, 5); + EXPECT_EQ(map.lookup_or_add("A", return_8), 5); + EXPECT_EQ(map.lookup_or_add("B", return_8), 8); +} + +TEST(string_map, LookupOrAddDefault) +{ + StringMap<std::string> map; + + std::string &a = map.lookup_or_add_default("A"); + EXPECT_EQ(a.size(), 0); + a += "Test"; + EXPECT_EQ(a.size(), 4); + std::string &b = map.lookup_or_add_default("A"); + EXPECT_EQ(b, "Test"); +} |