Jelajahi Sumber

optimize init_sockets a bit

This commit makes code uglier but it saves a few instructions in a hot path
what it does is remove the string path for init sockets
so that I can avoid doing several comparisons
it also makes some duplicate code, which is annoying, but it prevents
redundant comparisons, looping, and function calls
overall I expect a very small speedup, but everything counts.
Joseph Brandenburg 2 minggu lalu
induk
melakukan
1b9ad4ec8b
6 mengubah file dengan 30 tambahan dan 32 penghapusan
  1. 15 11
      base_definitions.py
  2. 1 8
      execution/readtree.py
  3. 2 1
      execution/schema_solve.py
  4. 4 3
      internal_nodes.py
  5. 3 1
      misc_nodes.py
  6. 5 8
      schema_nodes.py

+ 15 - 11
base_definitions.py

@@ -787,8 +787,18 @@ class MantisNode:
         self.execution_prepared = False
 
     def init_sockets(self) -> None:
-        self.inputs.init_sockets(self.socket_templates)
-        self.outputs.init_sockets(self.socket_templates)
+        # Formerly I was doing this:
+        # self.inputs.init_sockets(self.socket_templates)
+        # self.outputs.init_sockets(self.socket_templates)
+        # Since doing each one causes some redundancy, we'll just reimplement it
+        # This is a highly used and therefore expensive function so we have to make small
+        # optimizations like this, even tho it is ugly.
+        for socket in self.socket_templates:
+            if socket.ui_only == True: continue # used in UI but not in generation
+            if socket.is_input:
+                self.inputs[socket.name] = NodeSocket(is_input=True, name=socket.name, node=self.node)
+            else:
+                self.inputs[socket.name] = NodeSocket(is_input=False, name=socket.name, node=self.node)
 
     def init_parameters(self, additional_parameters = {}) -> None:
         for socket in self.inputs:
@@ -1222,15 +1232,9 @@ class MantisNodeSocketCollection(dict):
 
     def init_sockets(self, sockets):
         for socket in sockets:
-            if isinstance(socket, str):
-                self[socket] = NodeSocket(is_input=self.is_input, name=socket, node=self.node)
-            elif isinstance(socket, MantisSocketTemplate):
-                if socket.ui_only == True: continue # used in UI but not in generation
-                if socket.is_input != self.is_input: continue
-                self[socket.name] = NodeSocket(is_input=self.is_input, name=socket.name, node=self.node)
-            else:
-                raise RuntimeError(f"NodeSocketCollection keys must be str or MantisSocketTemplate, not {type(socket)}")
-
+            if socket.ui_only == True: continue # used in UI but not in generation
+            if socket.is_input != self.is_input: continue
+            self[socket.name] = NodeSocket(is_input=self.is_input, name=socket.name, node=self.node)
     def __delitem__(self, key):
         """Deletes a node socket by name, and all its links."""
         socket = self[key]

+ 1 - 8
execution/readtree.py

@@ -45,7 +45,7 @@ def autogen_node(base_tree, ui_socket, signature, mContext):
     from ..internal_nodes import AutoGenNode
     mantis_node = AutoGenNode(signature, base_tree)
     mantis_node.mContext = mContext
-    mantis_node.outputs.init_sockets([ui_socket.name])
+    mantis_node.outputs[ui_socket.name] = NodeSocket(is_input=False, name=ui_socket.name, node=mantis_node)
     mantis_node.ui_signature = None # does not exist in the UI
     return mantis_node
 
@@ -267,13 +267,6 @@ def node_scan_overlays(node, socket_id,):
     return result # because sometimes I need to get the result, too
 
 from ..utilities.utilities import check_and_add_root, init_connections, init_dependencies, init_schema_dependencies
-
-def is_signature_in_other_signature(parent_signature, child_signature):
-    # If the other signature is shorter, it isn't a child node
-    if len(parent_signature) > len(child_signature):
-        return False
-    return parent_signature[0:] == child_signature[:len(parent_signature)]
-
 def solve_schema_to_tree(mantis_node, all_mantis_nodes, roots=[], error_popups=False):
     from ..utilities.utilities import get_ui_node
     np = get_ui_node(mantis_node.signature, mantis_node.base_tree)

+ 2 - 1
execution/schema_solve.py

@@ -179,7 +179,8 @@ class SchemaSolver:
                     if ui_node.bl_idname in ["MantisNodeGroup"]:
                         mantis_node.prepared = False
                         mantis_node.node_type = 'DUMMY_SCHEMA' # we promote it to a schema for now
-                        mantis_node.inputs.init_sockets(['Schema Length']) # add a Schema Length socket
+                        from ..base_definitions import NodeSocket
+                        mantis_node.inputs['Schema Length'] = NodeSocket(is_input=True, node=mantis_node, name='Schema Length')
                         mantis_node.parameters['Schema Length'] = 1 # set the length to 1 since it is a single group instance
                         # we'll make the autogenerated nodes for constant inputs. It doesn't matter that there is technically
                         #  a ui_node available for each one  -- these are cheap and I want this to be easy.

+ 4 - 3
internal_nodes.py

@@ -148,11 +148,12 @@ class GroupInterface(MantisNode):
     node_type = 'INTERFACE'
     def __init__(self, signature, base_tree, ui_node, in_out):
         super().__init__(signature, base_tree)
-        self.prepared, self.executed = True, True; sockets = []
+        self.prepared, self.executed = True, True
         self.in_out = in_out
         # init the sockets based on in/out, then set up traversal
         collection = ui_node.inputs if in_out == 'INPUT' else ui_node.outputs
-        for socket in collection: sockets.append(socket.identifier)
-        self.inputs.init_sockets(sockets); self.outputs.init_sockets(sockets)
+        for socket in collection:
+            self.inputs[socket.identifier] = NodeSocket(is_input=True, node=self, name=socket.identifier)
+            self.outputs[socket.identifier] = NodeSocket(is_input=False, node=self, name=socket.identifier)
         for socket in self.inputs.keys(): self.set_traverse( [(socket, socket)] )
         self.execution_prepared=True

+ 3 - 1
misc_nodes.py

@@ -164,7 +164,9 @@ def array_choose_data(node, data, output):
         Used to choose the correct data to send out of an array-choose node.
     """
     # We need to make new outputs and link from each one based on the data in the array...
-    node.outputs.init_sockets([output+"."+str(i).zfill(4) for i in range(len(data)) ])
+    for i in range(len(data)):
+        name = output+"."+str(i).zfill(4)
+        node.outputs[name] = NodeSocket(node=node, name=name, is_input=False)
     for i, data_item in enumerate(data):
         node.parameters[output+"."+str(i).zfill(4)] = data_item
     for link in node.outputs[output].links:

+ 5 - 8
schema_nodes.py

@@ -29,7 +29,7 @@ def schema_init_sockets(mantis_node, is_input = True, in_out='INPUT', category='
             if item.item_type == 'PANEL': continue
             if item.parent and read_schema_type(item) == category:
                 if item.in_out == in_out:
-                    sockets.init_sockets([item.name])
+                    sockets[item.name] = NodeSocket(is_input=is_input, name=item.name, node=mantis_node)
     mantis_node.init_parameters()
 
 
@@ -91,7 +91,7 @@ class SchemaConstInput(SchemaNode):
             for item in parent_tree.interface.items_tree:
                 if item.item_type == 'PANEL': continue
                 if item.in_out == 'INPUT':
-                    sockets.init_sockets([item.name])
+                    sockets[item.name] = NodeSocket(is_input=False, node=self, name=item.name)
             self.init_parameters()
             
         else:
@@ -102,19 +102,16 @@ class SchemaConstOutput(SchemaNode):
     __slots__ = ()
     def __init__(self, signature, base_tree, parent_schema_node=None):
         super().__init__(signature, base_tree)
-        inputs = [
-            "Expose at Index",
-        ]
-        self.inputs.init_sockets(inputs)
+        sockets=self.inputs
+        sockets["Expose at Index"] = NodeSocket(is_input=False, node=self, name="Expose at Index")
         if parent_schema_node:
             # this allows us to generate the Constant Input from a normal Node Group
             # and treat the node group as a schema
             parent_tree = parent_schema_node.ui_node.node_tree
-            sockets=self.inputs
             for item in parent_tree.interface.items_tree:
                 if item.item_type == 'PANEL': continue
                 if item.in_out == 'OUTPUT':
-                    sockets.init_sockets([item.name, "Expose at Index",])
+                    sockets[item.name] = NodeSocket(is_input=False, node=self, name=item.name)
             self.init_parameters()
             self.parameters['Expose at Index"']=1