nodes.py

@@ -81,10 +81,7 @@ class RigNodesNodeTree(bpy.types.NodeTree):
         # Set iteration count foreach (nested) loop.
         for output_node in output_nodes:
             input_node = output_node._get_connection_node()
-            nodes = list(output_node._get_loop_nodes(input_node))
-            for node in nodes:
-                if node == output_node:
-                    continue
+            for node in output_node._get_innern_loop_nodes(input_node):
                 node.iterations *= output_node.num_socks
 
     def _run_loop(
@@ -108,10 +105,17 @@ class RigNodesNodeTree(bpy.types.NodeTree):
         while to_visit:
             node = to_visit.popleft()
 
+            if node in visited:
+                # Nested nodes are added to visited nodes as they get executed in advance.
+                # A Node can also be visited when a node has inputs from two different
+                # nodes; both of those will list this node as 'successor'.
+                continue
+
             if isinstance(node, LoopInputNode) and node != input_node:
                 # Find the nested loop.
                 nested_output_node = node._get_connection_node()
-                nested_loop = deque(nested_output_node._get_loop_nodes(node))
+                nested_loop = deque(nested_output_node._get_innern_loop_nodes(node))
+                nested_loop.appendleft(node)
                 nested_loop.append(nested_output_node)
 
                 # Don't revisit nested nodes in the parent loop.
@@ -123,12 +127,6 @@ class RigNodesNodeTree(bpy.types.NodeTree):
                     self._run_loop(node, depsgraph, nested_loop.copy(), depth + 1)
                 continue
 
-            if node in visited:
-                # Nested nodes are added to visited nodes as they get executed in advance.
-                # A Node can also be visited when a node has inputs from two different
-                # nodes; both of those will list this node as 'successor'.
-                continue
-
             print(f"{indent}\tRunning: {node}, remaining runs: {node.iterations}")
             visited.add(node)
             node.run(depsgraph)
@@ -164,17 +162,17 @@ class RigNodesNodeTree(bpy.types.NodeTree):
 
                 # Find the loop.
                 output_node = node._get_connection_node()
-                loop = deque(output_node._get_loop_nodes(node))
-                loop.append(output_node)
+                innern_loop = deque(output_node._get_innern_loop_nodes(node))
 
                 # Add (nested) loop nodes to visited.
-                for loop_node in loop:
+                visited.add(node)
+                visited.add(output_node)
+                for loop_node in innern_loop:
                     visited.add(loop_node)
 
                 # Run all iterations of the found loop (recursive if nested).
                 for _ in range(0, output_node.num_socks):
                     self._run_loop(node, depsgraph, to_visit.copy())
-                continue
 
             # Everything that this node depends on has run, so time to run it
             # itself. It can be taken off the queue.
@@ -300,16 +298,23 @@ class AbstractRigNodesNode(bpy.types.Node):
     def _outputs(self):
         return self.outputs
 
-    def _get_loop_nodes(
-        self: "LoopOutputNode", input_node: "LoopInputNode"
+    def _get_innern_loop_nodes(
+        self: "LoopOutputNode", input_node: "LoopInputNode",
     ) -> Iterable["AbstractRigNodesNode"]:
-        """Generator, yields nodes that should be executed before the loop output node."""
-        nodes = list(self._connected_nodes(self._inputs, "from_socket"))
+        """Generator, yields nodes between the loop input and the loop output node"""
+        visited: Set["AbstractRigNodesNode"] = set()
+        to_visit = [self]
 
-        if not input_node in nodes and len(nodes) > 0:
-            from_node = nodes[-1]
-            yield from from_node._get_loop_nodes(input_node)
-        yield from nodes
+        while to_visit:
+            from_node = to_visit.pop()
+            for node in from_node.exec_order_prerequisites():
+                if node in visited:
+                    continue
+
+                visited.add(node)
+                if node != input_node:
+                    to_visit.append(node)
+                    yield node
 
     def exec_order_prerequisites(self) -> Iterable["AbstractRigNodesNode"]:
         """Generator, yields the nodes that should be executed before this one."""
@@ -593,7 +598,6 @@ class LoopInputNode(AbstractRigNodesNode):
                     raise RuntimeError(
                         "Loop Output Node is required to execute a loop."
                     )
-
                 if isinstance(connected_socket.node, LoopOutputNode):
                     return connected_socket.node
         raise RuntimeError("Loop Output Node is required to execute a loop.")
@@ -607,6 +611,7 @@ class LoopInputNode(AbstractRigNodesNode):
     def reset_run(self):
         super().reset_run()
         self.index = 0
+        self.iterations = self.num_socks
 
     def init(self, context: bpy.types.Context) -> None:
         for index in range(self.num_socks):