extract IdMultiMap data structure

3 files changed, 94 insertions, 100 deletions

diff --git a/source/blender/functions/FN_node_tree_multi_function_network.h b/source/blender/functions/FN_node_tree_multi_function_network.h
index 4ef35342df2..7532d475c29 100644
--- a/source/blender/functions/FN_node_tree_multi_function_network.h
+++ b/source/blender/functions/FN_node_tree_multi_function_network.h
@@ -11,8 +11,67 @@ namespace FN {
 
 using BLI::MultiMap;
 
-#define InlinedTreeMFSocketMap_UNMAPPED UINT_MAX
-#define InlinedTreeMFSocketMap_MULTIMAPPED (UINT_MAX - 1)
+#define IdMultiMap_UNMAPPED UINT_MAX
+#define IdMultiMap_MULTIMAPPED (UINT_MAX - 1)
+
+class IdMultiMap {
+ private:
+  Array<uint> m_single_mapping;
+  MultiMap<uint, uint> m_fallback_multimap;
+
+ public:
+  IdMultiMap(uint max_key_id) : m_single_mapping(max_key_id, IdMultiMap_UNMAPPED)
+  {
+  }
+
+  bool contains(uint key_id) const
+  {
+    return m_single_mapping[key_id] != IdMultiMap_UNMAPPED;
+  }
+
+  ArrayRef<uint> lookup(uint key_id) const
+  {
+    const uint &stored_value = m_single_mapping[key_id];
+    switch (stored_value) {
+      case IdMultiMap_UNMAPPED: {
+        return {};
+      }
+      case IdMultiMap_MULTIMAPPED: {
+        return m_fallback_multimap.lookup(key_id);
+      }
+      default:
+        return ArrayRef<uint>(&stored_value, 1);
+    }
+  }
+
+  uint lookup_single(uint key_id) const
+  {
+    uint stored_value = m_single_mapping[key_id];
+    BLI_assert(stored_value != IdMultiMap_UNMAPPED && stored_value != IdMultiMap_MULTIMAPPED);
+    return stored_value;
+  }
+
+  void add(uint key_id, uint value_id)
+  {
+    uint &stored_value = m_single_mapping[key_id];
+    switch (stored_value) {
+      case IdMultiMap_UNMAPPED: {
+        stored_value = value_id;
+        break;
+      }
+      case IdMultiMap_MULTIMAPPED: {
+        m_fallback_multimap.add(key_id, value_id);
+        break;
+      }
+      default: {
+        uint other_value_id = stored_value;
+        stored_value = IdMultiMap_MULTIMAPPED;
+        m_fallback_multimap.add_multiple_new(key_id, {other_value_id, value_id});
+        break;
+      }
+    }
+  }
+};
 
 class InlinedTreeMFSocketMap {
  private:
@@ -21,67 +80,51 @@ class InlinedTreeMFSocketMap {
    */
   const FunctionNodeTree *m_function_tree;
   const MFNetwork *m_network;
-  Array<uint> m_single_socket_by_fsocket;
-  MultiMap<uint, uint> m_multiple_inputs_by_fsocket;
+  IdMultiMap m_socket_by_fsocket;
   Array<uint> m_fsocket_by_socket;
 
  public:
   InlinedTreeMFSocketMap(const FunctionNodeTree &function_tree,
                          const MFNetwork &network,
-                         Array<uint> single_socket_by_fsocket,
-                         MultiMap<uint, uint> multiple_inputs_by_fsocket,
+                         IdMultiMap socket_by_fsocket,
                          Array<uint> fsocket_by_socket)
       : m_function_tree(&function_tree),
         m_network(&network),
-        m_single_socket_by_fsocket(std::move(single_socket_by_fsocket)),
-        m_multiple_inputs_by_fsocket(std::move(multiple_inputs_by_fsocket)),
+        m_socket_by_fsocket(std::move(socket_by_fsocket)),
         m_fsocket_by_socket(std::move(fsocket_by_socket))
   {
   }
 
   bool is_mapped(const FSocket &fsocket) const
   {
-    return m_single_socket_by_fsocket[fsocket.id()] < InlinedTreeMFSocketMap_MULTIMAPPED;
+    return m_socket_by_fsocket.contains(fsocket.id());
   }
 
   bool is_mapped(const MFSocket &socket) const
   {
-    return m_fsocket_by_socket[socket.id()] != InlinedTreeMFSocketMap_UNMAPPED;
+    return m_fsocket_by_socket[socket.id()] != IdMultiMap_UNMAPPED;
   }
 
   const MFInputSocket &lookup_singly_mapped_input_socket(const FInputSocket &fsocket) const
   {
-    BLI_assert(this->lookup_socket(fsocket).size() == 1);
-    uint mapped_id = m_single_socket_by_fsocket[fsocket.id()];
+    uint mapped_id = m_socket_by_fsocket.lookup_single(fsocket.id());
     return m_network->socket_by_id(mapped_id).as_input();
   }
 
   Vector<const MFInputSocket *> lookup_socket(const FInputSocket &fsocket) const
   {
-    uint id = fsocket.id();
-    uint mapped_value = m_single_socket_by_fsocket[id];
-    switch (mapped_value) {
-      case InlinedTreeMFSocketMap_UNMAPPED: {
-        return {};
-      }
-      case InlinedTreeMFSocketMap_MULTIMAPPED: {
-        Vector<const MFInputSocket *> sockets;
-        for (uint mapped_id : m_multiple_inputs_by_fsocket.lookup(id)) {
-          sockets.append(&m_network->socket_by_id(mapped_id).as_input());
-        }
-        return sockets;
-      }
-      default: {
-        uint mapped_id = mapped_value;
-        const MFInputSocket &socket = m_network->socket_by_id(mapped_id).as_input();
-        return {&socket};
-      }
+    ArrayRef<uint> mapped_ids = m_socket_by_fsocket.lookup(fsocket.id());
+    Vector<const MFInputSocket *> sockets;
+    sockets.reserve(mapped_ids.size());
+    for (uint mapped_id : mapped_ids) {
+      sockets.append(&m_network->socket_by_id(mapped_id).as_input());
     }
+    return sockets;
   }
 
   const MFOutputSocket &lookup_socket(const FOutputSocket &fsocket) const
   {
-    uint mapped_id = m_single_socket_by_fsocket[fsocket.id()];
+    uint mapped_id = m_socket_by_fsocket.lookup_single(fsocket.id());
     return m_network->socket_by_id(mapped_id).as_output();
   }
 
diff --git a/source/blender/functions/intern/node_tree_multi_function_network/builder.cc b/source/blender/functions/intern/node_tree_multi_function_network/builder.cc
index 5691977cd05..bee93531afb 100644
--- a/source/blender/functions/intern/node_tree_multi_function_network/builder.cc
+++ b/source/blender/functions/intern/node_tree_multi_function_network/builder.cc
@@ -13,7 +13,7 @@ FunctionTreeMFNetworkBuilder::FunctionTreeMFNetworkBuilder(
       m_preprocessed_function_tree_data(preprocessed_function_tree_data),
       m_function_tree_mappings(function_tree_mappings),
       m_resources(resources),
-      m_single_socket_by_fsocket(function_tree.socket_count(), InlinedTreeMFSocketMap_UNMAPPED),
+      m_socket_by_fsocket(function_tree.socket_count()),
       m_builder(BLI::make_unique<MFNetworkBuilder>())
 {
 }
@@ -205,31 +205,16 @@ std::unique_ptr<FunctionTreeMFNetwork> FunctionTreeMFNetworkBuilder::build()
 
   auto network = BLI::make_unique<MFNetwork>(std::move(m_builder));
 
-  Array<uint> fsocket_by_socket(network->socket_ids().size(), InlinedTreeMFSocketMap_UNMAPPED);
-  for (uint fsocket_id : m_single_socket_by_fsocket.index_range()) {
-    switch (m_single_socket_by_fsocket[fsocket_id]) {
-      case InlinedTreeMFSocketMap_UNMAPPED: {
-        break;
-      }
-      case InlinedTreeMFSocketMap_MULTIMAPPED: {
-        for (uint socket_id : m_multiple_inputs_by_fsocket.lookup(fsocket_id)) {
-          fsocket_by_socket[socket_id] = fsocket_id;
-        }
-        break;
-      }
-      default: {
-        uint socket_id = m_single_socket_by_fsocket[fsocket_id];
-        fsocket_by_socket[socket_id] = fsocket_id;
-        break;
-      }
+  Array<uint> fsocket_by_socket(network->socket_ids().size(), IdMultiMap_UNMAPPED);
+  for (uint fsocket_id : IndexRange(m_function_tree.socket_count())) {
+    ArrayRef<uint> mapped_ids = m_socket_by_fsocket.lookup(fsocket_id);
+    for (uint mapped_id : mapped_ids) {
+      fsocket_by_socket[mapped_id] = fsocket_id;
     }
   }
 
-  InlinedTreeMFSocketMap socket_map(m_function_tree,
-                                    *network,
-                                    std::move(m_single_socket_by_fsocket),
-                                    std::move(m_multiple_inputs_by_fsocket),
-                                    std::move(fsocket_by_socket));
+  InlinedTreeMFSocketMap socket_map(
+      m_function_tree, *network, std::move(m_socket_by_fsocket), std::move(fsocket_by_socket));
 
   return BLI::make_unique<FunctionTreeMFNetwork>(
       m_function_tree, std::move(network), std::move(socket_map));
diff --git a/source/blender/functions/intern/node_tree_multi_function_network/builder.h b/source/blender/functions/intern/node_tree_multi_function_network/builder.h
index 7581eef06a7..07819ecba9b 100644
--- a/source/blender/functions/intern/node_tree_multi_function_network/builder.h
+++ b/source/blender/functions/intern/node_tree_multi_function_network/builder.h
@@ -64,12 +64,7 @@ class FunctionTreeMFNetworkBuilder : BLI::NonCopyable, BLI::NonMovable {
   const VTreeMultiFunctionMappings &m_function_tree_mappings;
   ResourceCollector &m_resources;
 
-  /* By default store mapping between fsockets and builder sockets in an array.
-   * Input fsockets can be mapped to multiple new sockets. So fallback to a multimap in this case.
-   */
-  Array<uint> m_single_socket_by_fsocket;
-  MultiMap<uint, uint> m_multiple_inputs_by_fsocket;
-  static constexpr intptr_t MULTI_MAP_INDICATOR = 1;
+  IdMultiMap m_socket_by_fsocket;
 
   Map<const FGroupInput *, MFBuilderOutputSocket *> m_group_inputs_mapping;
 
@@ -152,31 +147,12 @@ class FunctionTreeMFNetworkBuilder : BLI::NonCopyable, BLI::NonMovable {
 
   void map_sockets(const FInputSocket &fsocket, MFBuilderInputSocket &socket)
   {
-    switch (m_single_socket_by_fsocket[fsocket.id()]) {
-      case InlinedTreeMFSocketMap_UNMAPPED: {
-        m_single_socket_by_fsocket[fsocket.id()] = socket.id();
-        break;
-      }
-      case InlinedTreeMFSocketMap_MULTIMAPPED: {
-        BLI_assert(!m_multiple_inputs_by_fsocket.lookup(fsocket.id()).contains(socket.id()));
-        m_multiple_inputs_by_fsocket.add(fsocket.id(), socket.id());
-        break;
-      }
-      default: {
-        uint already_inserted_id = m_single_socket_by_fsocket[fsocket.id()];
-        BLI_assert(already_inserted_id != socket.id());
-        m_multiple_inputs_by_fsocket.add_multiple_new(fsocket.id(),
-                                                      {already_inserted_id, socket.id()});
-        m_single_socket_by_fsocket[fsocket.id()] = InlinedTreeMFSocketMap_MULTIMAPPED;
-        break;
-      }
-    }
+    m_socket_by_fsocket.add(fsocket.id(), socket.id());
   }
 
   void map_sockets(const FOutputSocket &fsocket, MFBuilderOutputSocket &socket)
   {
-    BLI_assert(m_single_socket_by_fsocket[fsocket.id()] == InlinedTreeMFSocketMap_UNMAPPED);
-    m_single_socket_by_fsocket[fsocket.id()] = socket.id();
+    m_socket_by_fsocket.add(fsocket.id(), socket.id());
   }
 
   void map_sockets(ArrayRef<const FInputSocket *> fsockets,
@@ -209,7 +185,7 @@ class FunctionTreeMFNetworkBuilder : BLI::NonCopyable, BLI::NonMovable {
 
   bool fsocket_is_mapped(const FSocket &fsocket) const
   {
-    return m_single_socket_by_fsocket[fsocket.id()] != InlinedTreeMFSocketMap_UNMAPPED;
+    return m_socket_by_fsocket.contains(fsocket.id());
   }
 
   void assert_fnode_is_mapped_correctly(const FNode &fnode) const;
@@ -220,9 +196,7 @@ class FunctionTreeMFNetworkBuilder : BLI::NonCopyable, BLI::NonMovable {
 
   MFBuilderSocket &lookup_single_socket(const FSocket &fsocket) const
   {
-    uint mapped_id = m_single_socket_by_fsocket[fsocket.id()];
-    BLI_assert(
-        !ELEM(mapped_id, InlinedTreeMFSocketMap_MULTIMAPPED, InlinedTreeMFSocketMap_UNMAPPED));
+    uint mapped_id = m_socket_by_fsocket.lookup_single(fsocket.id());
     return *m_builder->sockets_by_id()[mapped_id];
   }
 
@@ -233,22 +207,14 @@ class FunctionTreeMFNetworkBuilder : BLI::NonCopyable, BLI::NonMovable {
 
   Vector<MFBuilderInputSocket *> lookup_socket(const FInputSocket &fsocket) const
   {
+    ArrayRef<uint> mapped_ids = m_socket_by_fsocket.lookup(fsocket.id());
+    ArrayRef<MFBuilderSocket *> sockets_by_id = m_builder->sockets_by_id();
+
     Vector<MFBuilderInputSocket *> sockets;
-    switch (m_single_socket_by_fsocket[fsocket.id()]) {
-      case InlinedTreeMFSocketMap_UNMAPPED: {
-        break;
-      }
-      case InlinedTreeMFSocketMap_MULTIMAPPED: {
-        for (uint mapped_id : m_multiple_inputs_by_fsocket.lookup(fsocket.id())) {
-          sockets.append(&m_builder->sockets_by_id()[mapped_id]->as_input());
-        }
-        break;
-      }
-      default: {
-        uint mapped_id = m_single_socket_by_fsocket[fsocket.id()];
-        sockets.append(&m_builder->sockets_by_id()[mapped_id]->as_input());
-        break;
-      }
+    sockets.reserve(mapped_ids.size());
+    for (uint mapped_id : mapped_ids) {
+      MFBuilderSocket &socket = *sockets_by_id[mapped_id];
+      sockets.append(&socket.as_input());
     }
     return sockets;
   }