[Bf-blender-cvs] [bb1f02510b6] master: BLI: support mutable lookup from multi value map

[Jacques Lucke]

Commit: bb1f02510b69a70ab9bde68be1a4d51c47d72d57
Author: Jacques Lucke
Date:   Thu Mar 4 17:37:46 2021 +0100
Branches: master
https://developer.blender.org/rBbb1f02510b69a70ab9bde68be1a4d51c47d72d57

BLI: support mutable lookup from multi value map

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

M	source/blender/blenlib/BLI_multi_value_map.hh
M	source/blender/blenlib/tests/BLI_multi_value_map_test.cc

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

diff --git a/source/blender/blenlib/BLI_multi_value_map.hh b/source/blender/blenlib/BLI_multi_value_map.hh
index 018f080e633..4113085a1e3 100644
--- a/source/blender/blenlib/BLI_multi_value_map.hh
+++ b/source/blender/blenlib/BLI_multi_value_map.hh
@@ -103,6 +103,22 @@ template<typename Key, typename Value> class MultiValueMap {
     return {};
   }
 
+  /**
+   * Get a mutable span to all the values that are stored for the given key.
+   */
+  MutableSpan<Value> lookup(const Key &key)
+  {
+    return this->lookup_as(key);
+  }
+  template<typename ForwardKey> MutableSpan<Value> lookup_as(const ForwardKey &key)
+  {
+    Vector<Value> *vector = map_.lookup_ptr_as(key);
+    if (vector != nullptr) {
+      return vector->as_mutable_span();
+    }
+    return {};
+  }
+
   /**
    * Note: This signature will change when the implementation changes.
    */
diff --git a/source/blender/blenlib/tests/BLI_multi_value_map_test.cc b/source/blender/blenlib/tests/BLI_multi_value_map_test.cc
index 7501fbe0d87..a910f2935d4 100644
--- a/source/blender/blenlib/tests/BLI_multi_value_map_test.cc
+++ b/source/blender/blenlib/tests/BLI_multi_value_map_test.cc
@@ -29,6 +29,27 @@ TEST(multi_value_map, LookupExistant)
   EXPECT_EQ(map.lookup(3)[0], 6);
 }
 
+TEST(multi_value_map, LookupMutable)
+{
+  MultiValueMap<int, int> map;
+  map.add(1, 2);
+  map.add(4, 5);
+  map.add(4, 6);
+  map.add(6, 7);
+
+  MutableSpan<int> span = map.lookup(4);
+  EXPECT_EQ(span.size(), 2);
+  span[0] = 10;
+  span[1] = 20;
+
+  map.add(4, 5);
+  MutableSpan<int> new_span = map.lookup(4);
+  EXPECT_EQ(new_span.size(), 3);
+  EXPECT_EQ(new_span[0], 10);
+  EXPECT_EQ(new_span[1], 20);
+  EXPECT_EQ(new_span[2], 5);
+}
+
 TEST(multi_value_map, AddMultiple)
 {
   MultiValueMap<int, int> map;