Nodes: cache socket identifier to index mapping

While this preprocessing does take some time upfront,
it avoids longer lookup later on, especially as nodes get
more sockets.

It's probably possible to make this more efficient in some cases
but this is good enough for now.

3 files changed, 120 insertions, 0 deletions

diff --git a/source/blender/nodes/NOD_derived_node_tree.hh b/source/blender/nodes/NOD_derived_node_tree.hh
index e12e0bde975..de9e4c8c812 100644
--- a/source/blender/nodes/NOD_derived_node_tree.hh
+++ b/source/blender/nodes/NOD_derived_node_tree.hh
@@ -95,6 +95,9 @@ class DNode {
 
   DInputSocket input(int index) const;
   DOutputSocket output(int index) const;
+
+  DInputSocket input_by_identifier(StringRef identifier) const;
+  DOutputSocket output_by_identifier(StringRef identifier) const;
 };
 
 /* A (nullable) reference to a socket and the context it is in. It is unique within an entire
@@ -288,6 +291,16 @@ inline DOutputSocket DNode::output(int index) const
   return {context_, &node_ref_->output(index)};
 }
 
+inline DInputSocket DNode::input_by_identifier(StringRef identifier) const
+{
+  return {context_, &node_ref_->input_by_identifier(identifier)};
+}
+
+inline DOutputSocket DNode::output_by_identifier(StringRef identifier) const
+{
+  return {context_, &node_ref_->output_by_identifier(identifier)};
+}
+
 /* --------------------------------------------------------------------
  * DSocket inline methods.
  */
diff --git a/source/blender/nodes/NOD_node_tree_ref.hh b/source/blender/nodes/NOD_node_tree_ref.hh
index d4805daf8f5..b028fc28bbc 100644
--- a/source/blender/nodes/NOD_node_tree_ref.hh
+++ b/source/blender/nodes/NOD_node_tree_ref.hh
@@ -70,6 +70,8 @@ class NodeTreeRef;
 class LinkRef;
 class InternalLinkRef;
 
+using SocketIndexByIdentifierMap = Map<std::string, int>;
+
 class SocketRef : NonCopyable, NonMovable {
  protected:
   NodeRef *node_;
@@ -169,6 +171,8 @@ class NodeRef : NonCopyable, NonMovable {
   Vector<InputSocketRef *> inputs_;
   Vector<OutputSocketRef *> outputs_;
   Vector<InternalLinkRef *> internal_links_;
+  SocketIndexByIdentifierMap *input_index_by_identifier_;
+  SocketIndexByIdentifierMap *output_index_by_identifier_;
 
   friend NodeTreeRef;
 
@@ -182,6 +186,9 @@ class NodeRef : NonCopyable, NonMovable {
   const InputSocketRef &input(int index) const;
   const OutputSocketRef &output(int index) const;
 
+  const InputSocketRef &input_by_identifier(StringRef identifier) const;
+  const OutputSocketRef &output_by_identifier(StringRef identifier) const;
+
   bNode *bnode() const;
   bNodeTree *btree() const;
 
@@ -246,6 +253,7 @@ class NodeTreeRef : NonCopyable, NonMovable {
   Vector<OutputSocketRef *> output_sockets_;
   Vector<LinkRef *> links_;
   MultiValueMap<const bNodeType *, NodeRef *> nodes_by_type_;
+  Vector<std::unique_ptr<SocketIndexByIdentifierMap>> owned_identifier_maps_;
 
  public:
   NodeTreeRef(bNodeTree *btree);
@@ -279,6 +287,7 @@ class NodeTreeRef : NonCopyable, NonMovable {
                                       bNodeSocket *bsocket);
 
   void create_linked_socket_caches();
+  void create_socket_identifier_maps();
 };
 
 using NodeTreeRefMap = Map<bNodeTree *, std::unique_ptr<const NodeTreeRef>>;
@@ -502,6 +511,18 @@ inline const OutputSocketRef &NodeRef::output(int index) const
   return *outputs_[index];
 }
 
+inline const InputSocketRef &NodeRef::input_by_identifier(StringRef identifier) const
+{
+  const int index = input_index_by_identifier_->lookup_as(identifier);
+  return this->input(index);
+}
+
+inline const OutputSocketRef &NodeRef::output_by_identifier(StringRef identifier) const
+{
+  const int index = output_index_by_identifier_->lookup_as(identifier);
+  return this->output(index);
+}
+
 inline bNode *NodeRef::bnode() const
 {
   return bnode_;
diff --git a/source/blender/nodes/intern/node_tree_ref.cc b/source/blender/nodes/intern/node_tree_ref.cc
index 8aa582e250c..154ee716153 100644
--- a/source/blender/nodes/intern/node_tree_ref.cc
+++ b/source/blender/nodes/intern/node_tree_ref.cc
@@ -14,6 +14,8 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include <mutex>
+
 #include "NOD_node_tree_ref.hh"
 
 #include "BLI_dot_export.hh"
@@ -105,6 +107,7 @@ NodeTreeRef::NodeTreeRef(bNodeTree *btree) : btree_(btree)
     }
   }
 
+  this->create_socket_identifier_maps();
   this->create_linked_socket_caches();
 
   for (NodeRef *node : nodes_by_id_) {
@@ -316,6 +319,89 @@ void OutputSocketRef::foreach_logical_target(
   }
 }
 
+namespace {
+struct SocketByIdentifierMap {
+  SocketIndexByIdentifierMap *map = nullptr;
+  std::unique_ptr<SocketIndexByIdentifierMap> owned_map;
+};
+}  // namespace
+
+static std::unique_ptr<SocketIndexByIdentifierMap> create_identifier_map(const ListBase &sockets)
+{
+  std::unique_ptr<SocketIndexByIdentifierMap> map = std::make_unique<SocketIndexByIdentifierMap>();
+  int index;
+  LISTBASE_FOREACH_INDEX (bNodeSocket *, socket, &sockets, index) {
+    map->add_new(socket->identifier, index);
+  }
+  return map;
+}
+
+/* This function is not threadsafe.  */
+static SocketByIdentifierMap get_or_create_identifier_map(
+    const bNode &node, const ListBase &sockets, const bNodeSocketTemplate *sockets_template)
+{
+  SocketByIdentifierMap map;
+  if (sockets_template == nullptr) {
+    if (BLI_listbase_is_empty(&sockets)) {
+      static SocketIndexByIdentifierMap empty_map;
+      map.map = &empty_map;
+    }
+    else if (node.type == NODE_REROUTE) {
+      if (&node.inputs == &sockets) {
+        static SocketIndexByIdentifierMap reroute_input_map = [] {
+          SocketIndexByIdentifierMap map;
+          map.add_new("Input", 0);
+          return map;
+        }();
+        map.map = &reroute_input_map;
+      }
+      else {
+        static SocketIndexByIdentifierMap reroute_output_map = [] {
+          SocketIndexByIdentifierMap map;
+          map.add_new("Output", 0);
+          return map;
+        }();
+        map.map = &reroute_output_map;
+      }
+    }
+    else {
+      /* The node has a dynamic amount of sockets. Therefore we need to create a new map. */
+      map.owned_map = create_identifier_map(sockets);
+      map.map = &*map.owned_map;
+    }
+  }
+  else {
+    /* Cache only one map for nodes that have the same sockets. */
+    static Map<const bNodeSocketTemplate *, std::unique_ptr<SocketIndexByIdentifierMap>> maps;
+    map.map = &*maps.lookup_or_add_cb(sockets_template,
+                                      [&]() { return create_identifier_map(sockets); });
+  }
+  return map;
+}
+
+void NodeTreeRef::create_socket_identifier_maps()
+{
+  /* `get_or_create_identifier_map` is not threadsafe, therefore we have to hold a lock here. */
+  static std::mutex mutex;
+  std::lock_guard lock{mutex};
+
+  for (NodeRef *node : nodes_by_id_) {
+    bNode &bnode = *node->bnode_;
+    SocketByIdentifierMap inputs_map = get_or_create_identifier_map(
+        bnode, bnode.inputs, bnode.typeinfo->inputs);
+    SocketByIdentifierMap outputs_map = get_or_create_identifier_map(
+        bnode, bnode.outputs, bnode.typeinfo->outputs);
+    node->input_index_by_identifier_ = inputs_map.map;
+    node->output_index_by_identifier_ = outputs_map.map;
+    if (inputs_map.owned_map) {
+      owned_identifier_maps_.append(std::move(inputs_map.owned_map));
+    }
+    if (outputs_map.owned_map) {
+      owned_identifier_maps_.append(std::move(outputs_map.owned_map));
+    }
+  }
+}
+
 static bool has_link_cycles_recursive(const NodeRef &node,
                                       MutableSpan<bool> visited,
                                       MutableSpan<bool> is_in_stack)