[Bf-blender-cvs] [d8d6f2a521a] functions: smallmap: lookup or insert using callback

[Jacques Lucke]

Commit: d8d6f2a521a3da974658f3aa9ccac5bce5fa97b5
Author: Jacques Lucke
Date:   Mon May 20 10:55:15 2019 +0200
Branches: functions
https://developer.blender.org/rBd8d6f2a521a3da974658f3aa9ccac5bce5fa97b5

smallmap: lookup or insert using callback

===================================================================

M	source/blender/blenlib/BLI_small_map.hpp
M	tests/gtests/blenlib/BLI_small_map_test.cc

===================================================================

diff --git a/source/blender/blenlib/BLI_small_map.hpp b/source/blender/blenlib/BLI_small_map.hpp
index 6a10250f406..cfe275376f3 100644
--- a/source/blender/blenlib/BLI_small_map.hpp
+++ b/source/blender/blenlib/BLI_small_map.hpp
@@ -121,6 +121,17 @@ template<typename K, typename V, uint N = 4> class SmallMap {
     return this->lookup_ptr(key);
   }
 
+  template<typename... Args>
+  V &lookup_ref_or_insert_func(const K &key, V (*create_value)(Args... args), Args &&... args)
+  {
+    V *value = this->lookup_ptr(key);
+    if (value != NULL) {
+      return *value;
+    }
+    this->add_new(key, create_value(std::forward<Args>(args)...));
+    return this->lookup_ref(key);
+  }
+
   uint size() const
   {
     return m_entries.size();
diff --git a/tests/gtests/blenlib/BLI_small_map_test.cc b/tests/gtests/blenlib/BLI_small_map_test.cc
index fdc3c826b75..f11ac8260ca 100644
--- a/tests/gtests/blenlib/BLI_small_map_test.cc
+++ b/tests/gtests/blenlib/BLI_small_map_test.cc
@@ -164,3 +164,64 @@ TEST(small_map, ItemIterator)
   EXPECT_TRUE(values.contains(9.0f));
   EXPECT_TRUE(values.contains(0.0f));
 }
+
+float return_42()
+{
+  return 42.0f;
+}
+
+TEST(small_map, LookupOrInsertFunc_NoArgs)
+{
+  IntFloatMap map;
+  EXPECT_EQ(map.lookup_ref_or_insert_func(0, return_42), 42.0f);
+  EXPECT_EQ(map.lookup(0), 42);
+}
+
+float return_identity(float a)
+{
+  return a;
+}
+
+TEST(small_map, LookupOrInsertFunc_SingleArg)
+{
+  IntFloatMap map;
+  EXPECT_EQ(map.lookup_ref_or_insert_func(1, return_identity, 5.0f), 5.0f);
+  EXPECT_EQ(map.lookup(1), 5.0f);
+}
+
+float add_func(float a, float b)
+{
+  return a + b;
+}
+
+TEST(small_map, LookupOrInsertFunc_TwoArgs)
+{
+  IntFloatMap map;
+  EXPECT_EQ(map.lookup_ref_or_insert_func(2, add_func, 4.0f, 6.0f), 10.0f);
+  EXPECT_EQ(map.lookup(2), 10.0f);
+}
+
+TEST(small_map, LookupOrInsertFunc_NoReinsert)
+{
+  IntFloatMap map;
+  EXPECT_EQ(map.lookup_ref_or_insert_func(2, return_identity, 4.0f), 4.0f);
+  EXPECT_EQ(map.lookup_ref_or_insert_func(2, return_identity, 6.0f), 4.0f);
+  EXPECT_EQ(map.lookup_ref_or_insert_func(2, return_identity, 8.0f), 4.0f);
+  EXPECT_EQ(map.size(), 1);
+}
+
+float inc_value_and_return_42(int *ptr)
+{
+  *ptr += 1;
+  return 42.0f;
+}
+
+TEST(small_map, LookupOrInsertFunc_FuncCalledOnce)
+{
+  int counter = 0;
+  IntFloatMap map;
+  EXPECT_EQ(map.lookup_ref_or_insert_func(0, inc_value_and_return_42, &counter), 42.0f);
+  EXPECT_EQ(counter, 1);
+  EXPECT_EQ(map.lookup_ref_or_insert_func(0, inc_value_and_return_42, &counter), 42.0f);
+  EXPECT_EQ(counter, 1);
+}