only map dummy sockets when function network has been created

4 files changed, 78 insertions, 25 deletions

diff --git a/source/blender/blenlib/BLI_index_map.h b/source/blender/blenlib/BLI_index_map.h
new file mode 100644
index 00000000000..c8668880c42
--- /dev/null
+++ b/source/blender/blenlib/BLI_index_map.h
@@ -0,0 +1,48 @@
+#ifndef __BLI_INDEX_MAP_H__
+#define __BLI_INDEX_MAP_H__
+
+#include "BLI_array_cxx.h"
+
+namespace BLI {
+
+template<typename ValueT, uint N = 4, typename Allocator = GuardedAllocator> class IndexMap {
+ private:
+  Array<ValueT, N, Allocator> m_array;
+  ValueT m_sentinel;
+
+ public:
+  IndexMap(uint size, ValueT sentinel) : m_array(size, sentinel), m_sentinel(sentinel)
+  {
+  }
+
+  uint size() const
+  {
+    return m_array.size();
+  }
+
+  void add(uint key, const ValueT &value)
+  {
+    m_array[key] = value;
+  }
+
+  void add_new(uint key, const ValueT &value)
+  {
+    BLI_assert(m_array[key] == m_sentinel);
+    m_array[key] = value;
+  }
+
+  bool contains(uint key) const
+  {
+    return m_array[key] != m_sentinel;
+  }
+
+  const ValueT &lookup(uint key) const
+  {
+    BLI_assert(this->contains(key));
+    return m_array[key];
+  }
+};
+
+}  // namespace BLI
+
+#endif /* __BLI_INDEX_MAP_H__ */
diff --git a/source/blender/blenlib/CMakeLists.txt b/source/blender/blenlib/CMakeLists.txt
index 81d0f6f3a60..589ecaaecfd 100644
--- a/source/blender/blenlib/CMakeLists.txt
+++ b/source/blender/blenlib/CMakeLists.txt
@@ -276,6 +276,7 @@ set(SRC
   BLI_index_mask.h
   BLI_parallel.h
   BLI_string_multi_map.h
+  BLI_index_map.h
 )
 
 set(LIB
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 14c5035a501..206d285384f 100644
--- a/source/blender/functions/FN_node_tree_multi_function_network.h
+++ b/source/blender/functions/FN_node_tree_multi_function_network.h
@@ -4,11 +4,13 @@
 #include "FN_node_tree.h"
 
 #include "BLI_multi_map.h"
+#include "BLI_index_map.h"
 
 #include "FN_multi_function_network.h"
 
 namespace FN {
 
+using BLI::IndexMap;
 using BLI::MultiMap;
 
 #define IdMultiMap_UNMAPPED UINT_MAX
@@ -75,60 +77,54 @@ class IdMultiMap {
 
 class InlinedTreeMFSocketMap {
  private:
-  /* An input fsocket can be mapped to multiple sockets.
-   * An output fsocket can be mapped to at most one socket.
-   */
   const FunctionNodeTree *m_function_tree;
   const MFNetwork *m_network;
-  IdMultiMap m_socket_by_fsocket;
-  Array<uint> m_fsocket_by_socket;
+
+  IndexMap<const MFSocket *> m_dummy_socket_by_fsocket_id;
+  IndexMap<const FSocket *> m_fsocket_by_dummy_socket_id;
 
  public:
   InlinedTreeMFSocketMap(const FunctionNodeTree &function_tree,
                          const MFNetwork &network,
-                         IdMultiMap socket_by_fsocket,
-                         Array<uint> fsocket_by_socket)
+                         IndexMap<const MFSocket *> dummy_socket_by_fsocket_id,
+                         IndexMap<const FSocket *> fsocket_by_dummy_socket_id)
       : m_function_tree(&function_tree),
         m_network(&network),
-        m_socket_by_fsocket(std::move(socket_by_fsocket)),
-        m_fsocket_by_socket(std::move(fsocket_by_socket))
+        m_dummy_socket_by_fsocket_id(std::move(dummy_socket_by_fsocket_id)),
+        m_fsocket_by_dummy_socket_id(std::move(fsocket_by_dummy_socket_id))
   {
   }
 
   bool is_mapped(const FSocket &fsocket) const
   {
-    return m_socket_by_fsocket.contains(fsocket.id());
+    return m_dummy_socket_by_fsocket_id.contains(fsocket.id());
   }
 
   bool is_mapped(const MFSocket &socket) const
   {
-    return m_fsocket_by_socket[socket.id()] != IdMultiMap_UNMAPPED;
+    return m_fsocket_by_dummy_socket_id.contains(socket.id());
   }
 
   const MFInputSocket &lookup_singly_mapped_input_socket(const FInputSocket &fsocket) const
   {
-    uint mapped_id = m_socket_by_fsocket.lookup_single(fsocket.id());
-    return m_network->socket_by_id(mapped_id).as_input();
+    return m_dummy_socket_by_fsocket_id.lookup(fsocket.id())->as_input();
   }
 
   const MFOutputSocket &lookup_socket(const FOutputSocket &fsocket) const
   {
-    uint mapped_id = m_socket_by_fsocket.lookup_single(fsocket.id());
-    return m_network->socket_by_id(mapped_id).as_output();
+    return m_dummy_socket_by_fsocket_id.lookup(fsocket.id())->as_output();
   }
 
   const FInputSocket &lookup_fsocket(const MFInputSocket &socket) const
   {
     BLI_assert(socket.node().is_dummy());
-    uint mapped_id = m_fsocket_by_socket[socket.id()];
-    return m_function_tree->socket_by_id(mapped_id).as_input();
+    return m_fsocket_by_dummy_socket_id.lookup(socket.id())->as_input();
   }
 
   const FOutputSocket &lookup_fsocket(const MFOutputSocket &socket) const
   {
     BLI_assert(socket.node().is_dummy());
-    uint mapped_id = m_fsocket_by_socket[socket.id()];
-    return m_function_tree->socket_by_id(mapped_id).as_output();
+    return m_fsocket_by_dummy_socket_id.lookup(socket.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 bee93531afb..ad982c8e6cf 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
@@ -203,18 +203,26 @@ std::unique_ptr<FunctionTreeMFNetwork> FunctionTreeMFNetworkBuilder::build()
 {
   // m_builder->to_dot__clipboard();
 
+  IndexMap<const MFSocket *> dummy_socket_by_fsocket_id(m_function_tree.socket_count(), nullptr);
+  IndexMap<const FSocket *> fsocket_by_dummy_socket_id(m_builder->sockets_by_id().size(), nullptr);
+
   auto network = BLI::make_unique<MFNetwork>(std::move(m_builder));
 
-  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 (const FSocket *fsocket : m_function_tree.all_sockets()) {
+    ArrayRef<uint> mapped_ids = m_socket_by_fsocket.lookup(fsocket->id());
     for (uint mapped_id : mapped_ids) {
-      fsocket_by_socket[mapped_id] = fsocket_id;
+      const MFSocket &socket = network->socket_by_id(mapped_id);
+      if (socket.node().is_dummy()) {
+        dummy_socket_by_fsocket_id.add_new(fsocket->id(), &socket);
+        fsocket_by_dummy_socket_id.add_new(socket.id(), fsocket);
+      }
     }
   }
 
-  InlinedTreeMFSocketMap socket_map(
-      m_function_tree, *network, std::move(m_socket_by_fsocket), std::move(fsocket_by_socket));
+  InlinedTreeMFSocketMap socket_map(m_function_tree,
+                                    *network,
+                                    std::move(dummy_socket_by_fsocket_id),
+                                    std::move(fsocket_by_dummy_socket_id));
 
   return BLI::make_unique<FunctionTreeMFNetwork>(
       m_function_tree, std::move(network), std::move(socket_map));