Add StringMap.LookupOrAdd and StringMap.LookupOrAddDefault

source/blender/blenlib/BLI_string_map.hh

@@ -341,6 +341,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.
    */

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");
+}