BLI: add MultiValueMap

This is a convenience wrapper for `Map<Key, Vector<Value>>`.
It does not provide any performance benefits (yet). I need this
kind of map in a couple of places and before I was duplicating
the lookup logic in many places.

source/blender/blenlib/BLI_multi_value_map.hh

@@ -0,0 +1,133 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#ifndef __BLI_MULTI_VALUE_MAP_HH__
+#define __BLI_MULTI_VALUE_MAP_HH__
+
+/** \file
+ * \ingroup bli
+ *
+ * A `blender::MultiValueMap<Key, Value>` is an unordered associative container that stores
+ * key-value pairs. It is different from `blender::Map` in that it can store multiple values for
+ * the same key. The list of values that corresponds to a specific key can contain duplicates.
+ *
+ * This data structure is different from a `std::multi_map`, because multi_map can store the same
+ * key more than once and MultiValueMap can't.
+ *
+ * Currently, this class exists mainly for convenience. There are no performance benefits over
+ * using Map<Key, Vector<Value>>. In the future, a better implementation for this data structure
+ * can be developed.
+ */
+
+#include "BLI_map.hh"
+#include "BLI_vector.hh"
+
+namespace blender {
+
+template<typename Key, typename Value> class MultiValueMap {
+ private:
+  using MapType = Map<Key, Vector<Value>>;
+  MapType map_;
+
+ public:
+  /**
+   * Add a new value for the given key. If the map contains the key already, the value will be
+   * appended to the list of corresponding values.
+   */
+  void add(const Key &key, const Value &value)
+  {
+    this->add_as(key, value);
+  }
+  void add(const Key &key, Value &&value)
+  {
+    this->add_as(key, std::move(value));
+  }
+  void add(Key &&key, const Value &value)
+  {
+    this->add_as(std::move(key), value);
+  }
+  void add(Key &&key, Value &&value)
+  {
+    this->add_as(std::move(key), std::move(value));
+  }
+  template<typename ForwardKey, typename ForwardValue>
+  void add_as(ForwardKey &&key, ForwardValue &&value)
+  {
+    Vector<Value> &vector = map_.lookup_or_add_default_as(std::forward<ForwardKey>(key));
+    vector.append(std::forward<ForwardValue>(value));
+  }
+
+  /**
+   * Add all given values to the key.
+   */
+  void add_multiple(const Key &key, Span<Value> values)
+  {
+    this->add_multiple_as(key, values);
+  }
+  void add_multiple(Key &&key, Span<Value> values)
+  {
+    this->add_multiple_as(std::move(key), values);
+  }
+  template<typename ForwardKey> void add_multiple_as(ForwardKey &&key, Span<Value> values)
+  {
+    Vector<Value> &vector = map_.lookup_or_add_default_as(std::forward<ForwardKey>(key));
+    vector.extend(values);
+  }
+
+  /**
+   * Get a span to all the values that are stored for the given key.
+   */
+  Span<Value> lookup(const Key &key) const
+  {
+    return this->lookup_as(key);
+  }
+  template<typename ForwardKey> Span<Value> lookup_as(const ForwardKey &key) const
+  {
+    const Vector<Value> *vector = map_.lookup_ptr_as(key);
+    if (vector != nullptr) {
+      return vector->as_span();
+    }
+    return {};
+  }
+
+  /**
+   * Note: This signature will change when the implementation changes.
+   */
+  typename MapType::ItemIterator items() const
+  {
+    return map_.items();
+  }
+
+  /**
+   * Note: This signature will change when the implementation changes.
+   */
+  typename MapType::KeyIterator keys() const
+  {
+    return map_.keys();
+  }
+
+  /**
+   * Note: This signature will change when the implementation changes.
+   */
+  typename MapType::ValueIterator values() const
+  {
+    return map_.values();
+  }
+};
+
+}  // namespace blender
+
+#endif /* __BLI_MULTI_VALUE_MAP_HH__ */

source/blender/blenlib/CMakeLists.txt

@@ -353,6 +353,7 @@ if(WITH_GTESTS)
     tests/BLI_map_test.cc
     tests/BLI_math_base_safe_test.cc
     tests/BLI_memory_utils_test.cc
+    tests/BLI_multi_value_map_test.cc
     tests/BLI_set_test.cc
     tests/BLI_span_test.cc
     tests/BLI_stack_cxx_test.cc

source/blender/blenlib/tests/BLI_multi_value_map_test.cc

@@ -0,0 +1,109 @@
+/* Apache License, Version 2.0 */
+
+#include "BLI_multi_value_map.hh"
+#include "BLI_vector.hh"
+#include "testing/testing.h"
+
+namespace blender::tests {
+
+TEST(multi_value_map, LookupNotExistant)
+{
+  MultiValueMap<int, int> map;
+  EXPECT_EQ(map.lookup(5).size(), 0);
+  map.add(2, 5);
+  EXPECT_EQ(map.lookup(5).size(), 0);
+}
+
+TEST(multi_value_map, LookupExistant)
+{
+  MultiValueMap<int, int> map;
+  map.add(2, 4);
+  map.add(2, 5);
+  map.add(3, 6);
+
+  EXPECT_EQ(map.lookup(2).size(), 2);
+  EXPECT_EQ(map.lookup(2)[0], 4);
+  EXPECT_EQ(map.lookup(2)[1], 5);
+
+  EXPECT_EQ(map.lookup(3).size(), 1);
+  EXPECT_EQ(map.lookup(3)[0], 6);
+}
+
+TEST(multi_value_map, AddMultiple)
+{
+  MultiValueMap<int, int> map;
+  map.add_multiple(2, {4, 5, 6});
+  map.add_multiple(2, {1, 2});
+  map.add_multiple(5, {7, 5, 3});
+
+  EXPECT_EQ(map.lookup(2).size(), 5);
+  EXPECT_EQ(map.lookup(2)[0], 4);
+  EXPECT_EQ(map.lookup(2)[1], 5);
+  EXPECT_EQ(map.lookup(2)[2], 6);
+  EXPECT_EQ(map.lookup(2)[3], 1);
+  EXPECT_EQ(map.lookup(2)[4], 2);
+
+  EXPECT_EQ(map.lookup(5).size(), 3);
+  EXPECT_EQ(map.lookup(5)[0], 7);
+  EXPECT_EQ(map.lookup(5)[1], 5);
+  EXPECT_EQ(map.lookup(5)[2], 3);
+}
+
+TEST(multi_value_map, Keys)
+{
+  MultiValueMap<int, int> map;
+  map.add(5, 7);
+  map.add(5, 7);
+  map.add_multiple(2, {6, 7, 8});
+
+  Vector<int> keys;
+  for (int key : map.keys()) {
+    keys.append(key);
+  }
+
+  EXPECT_EQ(keys.size(), 2);
+  EXPECT_TRUE(keys.contains(5));
+  EXPECT_TRUE(keys.contains(2));
+}
+
+TEST(multi_value_map, Values)
+{
+  MultiValueMap<int, int> map;
+  map.add(3, 5);
+  map.add_multiple(3, {1, 2});
+  map.add(6, 1);
+
+  Vector<Span<int>> values;
+  for (Span<int> value_span : map.values()) {
+    values.append(value_span);
+  }
+
+  EXPECT_EQ(values.size(), 2);
+}
+
+TEST(multi_value_map, Items)
+{
+  MultiValueMap<int, int> map;
+  map.add_multiple(4, {1, 2, 3});
+
+  for (auto &&item : map.items()) {
+    int key = item.key;
+    Span<int> values = item.value;
+    EXPECT_EQ(key, 4);
+    EXPECT_EQ(values.size(), 3);
+    EXPECT_EQ(values[0], 1);
+    EXPECT_EQ(values[1], 2);
+    EXPECT_EQ(values[2], 3);
+  }
+}
+
+TEST(multi_value_map, UniquePtr)
+{
+  /* Mostly testing if it compiles here. */
+  MultiValueMap<std::unique_ptr<int>, std::unique_ptr<int>> map;
+  map.add(std::make_unique<int>(4), std::make_unique<int>(6));
+  map.add(std::make_unique<int>(4), std::make_unique<int>(7));
+  EXPECT_EQ(map.lookup(std::make_unique<int>(10)).size(), 0);
+}
+
+}  // namespace blender::tests

source/blender/functions/intern/multi_function_network_optimization.cc

@@ -28,6 +28,7 @@
 #include "BLI_disjoint_set.hh"
 #include "BLI_ghash.h"
 #include "BLI_map.hh"
+#include "BLI_multi_value_map.hh"
 #include "BLI_rand.h"
 #include "BLI_stack.hh"
 
@@ -403,15 +404,15 @@ static Array<uint64_t> compute_node_hashes(MFNetwork &network)
   return node_hashes;
 }
 
-static Map<uint64_t, Vector<MFNode *, 1>> group_nodes_by_hash(MFNetwork &network,
-                                                              Span<uint64_t> node_hashes)
+static MultiValueMap<uint64_t, MFNode *> group_nodes_by_hash(MFNetwork &network,
+                                                             Span<uint64_t> node_hashes)
 {
-  Map<uint64_t, Vector<MFNode *, 1>> nodes_by_hash;
+  MultiValueMap<uint64_t, MFNode *> nodes_by_hash;
   for (int id : IndexRange(network.node_id_amount())) {
     MFNode *node = network.node_or_null_by_id(id);
     if (node != nullptr) {
       uint64_t node_hash = node_hashes[id];
-      nodes_by_hash.lookup_or_add_default(node_hash).append(node);
+      nodes_by_hash.add(node_hash, node);
     }
   }
   return nodes_by_hash;
@@ -456,7 +457,7 @@ static bool nodes_output_same_values(DisjointSet &cache, const MFNode &a, const
 }
 
 static void relink_duplicate_nodes(MFNetwork &network,
-                                   Map<uint64_t, Vector<MFNode *, 1>> &nodes_by_hash)
+                                   MultiValueMap<uint64_t, MFNode *> &nodes_by_hash)
 {
   DisjointSet same_node_cache{network.node_id_amount()};
 
@@ -494,7 +495,7 @@ static void relink_duplicate_nodes(MFNetwork &network,
 void common_subnetwork_elimination(MFNetwork &network)
 {
   Array<uint64_t> node_hashes = compute_node_hashes(network);
-  Map<uint64_t, Vector<MFNode *, 1>> nodes_by_hash = group_nodes_by_hash(network, node_hashes);
+  MultiValueMap<uint64_t, MFNode *> nodes_by_hash = group_nodes_by_hash(network, node_hashes);
   relink_duplicate_nodes(network, nodes_by_hash);
 }
 

source/blender/nodes/NOD_derived_node_tree.hh

@@ -181,7 +181,7 @@ class DerivedNodeTree : NonCopyable, NonMovable {
   Vector<DInputSocket *> input_sockets_;
   Vector<DOutputSocket *> output_sockets_;
 
-  Map<const bNodeType *, Vector<DNode *>> nodes_by_type_;
+  MultiValueMap<const bNodeType *, DNode *> nodes_by_type_;
 
  public:
   DerivedNodeTree(bNodeTree *btree, NodeTreeRefMap &node_tree_refs);
@@ -483,13 +483,7 @@ inline Span<const DNode *> DerivedNodeTree::nodes_by_type(StringRefNull idname)
 
 inline Span<const DNode *> DerivedNodeTree::nodes_by_type(const bNodeType *nodetype) const
 {
-  const Vector<DNode *> *nodes = nodes_by_type_.lookup_ptr(nodetype);
-  if (nodes == nullptr) {
-    return {};
-  }
-  else {
-    return *nodes;
-  }
+  return nodes_by_type_.lookup(nodetype);
 }
 
 inline Span<const DSocket *> DerivedNodeTree::sockets() const

source/blender/nodes/NOD_node_tree_ref.hh

@@ -47,6 +47,7 @@
 #include "BLI_array.hh"
 #include "BLI_linear_allocator.hh"
 #include "BLI_map.hh"
+#include "BLI_multi_value_map.hh"
 #include "BLI_string_ref.hh"
 #include "BLI_timeit.hh"
 #include "BLI_utility_mixins.hh"
@@ -162,7 +163,7 @@ class NodeTreeRef : NonCopyable, NonMovable {
   Vector<SocketRef *> sockets_by_id_;
   Vector<InputSocketRef *> input_sockets_;
   Vector<OutputSocketRef *> output_sockets_;
-  Map<const bNodeType *, Vector<NodeRef *>> nodes_by_type_;
+  MultiValueMap<const bNodeType *, NodeRef *> nodes_by_type_;
 
  public:
   NodeTreeRef(bNodeTree *btree);
@@ -411,13 +412,7 @@ inline Span<const NodeRef *> NodeTreeRef::nodes_by_type(StringRefNull idname) co
 
 inline Span<const NodeRef *> NodeTreeRef::nodes_by_type(const bNodeType *nodetype) const
 {
-  const Vector<NodeRef *> *nodes = nodes_by_type_.lookup_ptr(nodetype);
-  if (nodes == nullptr) {
-    return {};
-  }
-  else {
-    return *nodes;
-  }
+  return nodes_by_type_.lookup(nodetype);
 }
 
 inline Span<const SocketRef *> NodeTreeRef::sockets() const

source/blender/nodes/intern/derived_node_tree.cc

@@ -321,7 +321,7 @@ BLI_NOINLINE void DerivedNodeTree::store_in_this_and_init_ids(
     node->id_ = node_index;
 
     const bNodeType *nodetype = node->node_ref_->bnode()->typeinfo;
-    nodes_by_type_.lookup_or_add_default(nodetype).append(node);
+    nodes_by_type_.add(nodetype, node);
 
     for (DInputSocket *socket : node->inputs_) {
       socket->id_ = sockets_by_id_.append_and_get_index(socket);

source/blender/nodes/intern/node_tree_ref.cc

@@ -79,7 +79,7 @@ NodeTreeRef::NodeTreeRef(bNodeTree *btree) : btree_(btree)
 
   for (NodeRef *node : nodes_by_id_) {
     const bNodeType *nodetype = node->bnode_->typeinfo;
-    nodes_by_type_.lookup_or_add_default(nodetype).append(node);
+    nodes_by_type_.add(nodetype, node);
   }
 }
 

source/blender/simulation/intern/simulation_collect_influences.cc

@@ -122,13 +122,12 @@ static void find_and_deduplicate_particle_attribute_nodes(nodes::MFNetworkTreeMa
     return;
   }
 
-  Map<std::pair<std::string, fn::MFDataType>, Vector<fn::MFNode *>>
+  MultiValueMap<std::pair<std::string, fn::MFDataType>, fn::MFNode *>
       attribute_nodes_by_name_and_type;
   for (int i : attribute_names->index_range()) {
-    attribute_nodes_by_name_and_type
-        .lookup_or_add_default(
-            {(*attribute_names)[i], name_sockets[i]->node().output(0).data_type()})
-        .append(&name_sockets[i]->node());
+    attribute_nodes_by_name_and_type.add(
+        {(*attribute_names)[i], name_sockets[i]->node().output(0).data_type()},
+        &name_sockets[i]->node());
   }
 
   Map<const fn::MFOutputSocket *, std::string> attribute_inputs;
@@ -237,11 +236,11 @@ class ParticleFunctionForce : public ParticleForce {
   }
 };
 
-static Vector<const ParticleForce *> create_forces_for_particle_simulation(
-    const nodes::DNode &simulation_node,
-    nodes::MFNetworkTreeMap &network_map,
-    ResourceCollector &resources,
-    DummyDataSources &data_sources)
+static void create_forces_for_particle_simulation(const nodes::DNode &simulation_node,
+                                                  nodes::MFNetworkTreeMap &network_map,
+                                                  ResourceCollector &resources,
+                                                  DummyDataSources &data_sources,
+                                                  SimulationInfluences &r_influences)
 {
   Vector<const ParticleForce *> forces;
   for (const nodes::DOutputSocket *origin_socket :
@@ -264,7 +263,9 @@ static Vector<const ParticleForce *> create_forces_for_particle_simulation(
     const ParticleForce &force = resources.construct<ParticleFunctionForce>(AT, *particle_fn);
     forces.append(&force);
   }
-  return forces;
+
+  std::string particle_name = dnode_to_path(simulation_node);
+  r_influences.particle_forces.add_multiple(std::move(particle_name), forces);
 }
 
 static void collect_forces(nodes::MFNetworkTreeMap &network_map,
@@ -273,10 +274,8 @@ static void collect_forces(nodes::MFNetworkTreeMap &network_map,
                            SimulationInfluences &r_influences)
 {
   for (const nodes::DNode *dnode : get_particle_simulation_nodes(network_map.tree())) {
-    std::string name = dnode_to_path(*dnode);
-    Vector<const ParticleForce *> forces = create_forces_for_particle_simulation(
-        *dnode, network_map, resources, data_sources);
-    r_influences.particle_forces.add_new(std::move(name), std::move(forces));
+    create_forces_for_particle_simulation(
+        *dnode, network_map, resources, data_sources, r_influences);
   }
 }
 

source/blender/simulation/intern/simulation_solver.cc

@@ -137,7 +137,7 @@ BLI_NOINLINE static void simulate_particle_chunk(SimulationSolveContext &solve_c
 {
   int particle_amount = attributes.size();
   Array<float3> force_vectors{particle_amount, {0, 0, 0}};
-  Span<const ParticleForce *> forces = solve_context.influences().get_particle_forces(
+  Span<const ParticleForce *> forces = solve_context.influences().particle_forces.lookup_as(
       state.head.name);
 
   ParticleChunkContext particle_chunk_context{IndexMask(particle_amount), attributes};

source/blender/simulation/intern/simulation_solver_influences.hh

@@ -18,6 +18,7 @@
 #define __SIM_SIMULATION_SOLVER_INFLUENCES_HH__
 
 #include "BLI_float3.hh"
+#include "BLI_multi_value_map.hh"
 #include "BLI_span.hh"
 
 #include "DNA_simulation_types.h"
@@ -48,31 +49,21 @@ class ParticleForce {
 };
 
 struct SimulationInfluences {
-  /* TODO: Use a special type for a map that contains vectors. */
-  Map<std::string, Vector<const ParticleForce *>> particle_forces;
+  MultiValueMap<std::string, const ParticleForce *> particle_forces;
   Map<std::string, fn::AttributesInfoBuilder *> particle_attributes_builder;
   Vector<const ParticleEmitter *> particle_emitters;
-
-  Span<const ParticleForce *> get_particle_forces(StringRef particle_name) const
-  {
-    const Vector<const ParticleForce *> *forces = particle_forces.lookup_ptr(particle_name);
-    if (forces == nullptr) {
-      return {};
-    }
-    return *forces;
-  }
 };
 
 class SimulationStateMap {
  private:
   Map<StringRefNull, SimulationState *> states_by_name_;
-  Map<StringRefNull, Vector<SimulationState *>> states_by_type_;
+  MultiValueMap<StringRefNull, SimulationState *> states_by_type_;
 
  public:
   void add(SimulationState *state)
   {
     states_by_name_.add_new(state->name, state);
-    states_by_type_.lookup_or_add_default(state->type).append(state);
+    states_by_type_.add(state->type, state);
   }
 
   template<typename StateType> StateType *lookup(StringRef name) const
@@ -101,13 +92,7 @@ class SimulationStateMap {
 
   Span<SimulationState *> lookup_type(StringRef type) const
   {
-    const Vector<SimulationState *> *states = states_by_type_.lookup_ptr_as(type);
-    if (states == nullptr) {
-      return {};
-    }
-    else {
-      return states->as_span();
-    }
+    return states_by_type_.lookup_as(type);
   }
 };