gui.py

@@ -54,6 +54,7 @@ def draw_nodetree_header_menu(self, context):
     layout = self.layout
     row = layout.row(align=False)
     row.operator("rignodes.execute_tree", icon="PLAY")
+    row.operator("rignodes.prepare_tree", icon="NODETREE")
     row.prop(tree, "autorun")
     row.prop(context.scene, "rignodes_mode")
     row.operator("rignodes.rebuild_node")

nodes.py

@@ -4,7 +4,7 @@ import functools
 from collections import deque
 from copy import copy
 from math import degrees, radians
-from typing import TypeVar, Callable, Optional, Iterable
+from typing import TypeVar, Callable, Optional, Iterable, Set
 
 import bpy
 import nodeitems_utils
@@ -65,21 +65,78 @@ class RigNodesNodeTree(bpy.types.NodeTree):
                 return node
         return None
 
-    def _prepare_nodes(self) -> None:
+    def _reset_nodes(self) -> None:
         for node in self.nodes:
             node.reset_run()
 
+    def _prepare_nodes(self) -> None:
+        self._reset_nodes()
+
+        # Find all loops in the tree.
+        output_nodes = list()
+        for node in self.nodes:
+            if isinstance(node, LoopOutputNode):
+                output_nodes.append(node)
+
+        # Set iteration count foreach (nested) loop.
+        for output_node in output_nodes:
+            input_node = output_node._get_connection_node()
+            nodes = list(output_node._get_nodes_between(input_node))
+            for node in nodes:
+                if node == output_node:
+                    continue
+                node.iterations *= output_node.num_socks
+
+    def _run_loop(
+        self,
+        node: "LoopInputNode",
+        depsgraph: bpy.types.Depsgraph,
+        to_visit: deque["AbstractRigNodesNode"],
+        depth: int = 1,
+    ) -> None: 
+        """Run one iteration of a loop.
+
+        :return: Node(s) to visit next."""
+
+        indent = '\t' * depth
+        print(f"{indent}\033[38;5;214mRunning {node}\033[0m")
+
+        input_node = node
+        output_node = node._get_connection_node()
+
+        while to_visit:
+            node = to_visit.popleft()
+
+            nested_loop: Optional[deque[AbstractRigNodesNode]] = None
+            if isinstance(node, LoopInputNode) and node != input_node:
+                # Run nested loop.
+                nested_output_node = node._get_connection_node()
+                nested_loop = deque(nested_output_node._get_nodes_between(node))
+                nested_loop.append(nested_output_node)
+                self._run_loop(node, depsgraph, nested_loop, depth+1)
+                # continue
+
+            print(f"{indent}\t-{node}")
+            node.run(depsgraph)
+            # if node.iterations > 0:
+            #     # Add output node for future iterations.
+            #     to_visit.append(node)
+
+            if node == output_node:
+                break
+
+
     def _run_from_node(
         self, depsgraph: bpy.types.Depsgraph, start_node: "AbstractRigNodesNode"
     ) -> None:
         # Execute the ForwardSolveNode and traverse its output connections
-        to_visit = deque([start_node])
+        to_visit: deque[AbstractRigNodesNode] = deque([start_node])
         visited = set()
 
         print(f"\033[38;5;214mRunning tree {self.name}\033[0m")
         while to_visit:
             node: AbstractRigNodesNode = to_visit[0]
-            print(f"    - {node}")
+            print(f"\t- {node}")
 
             nodes_before = list(node.exec_order_prerequisites())
             if nodes_before:
@@ -88,6 +145,11 @@ class RigNodesNodeTree(bpy.types.NodeTree):
                 to_visit.extendleft(reversed(nodes_before))
                 continue
 
+            if isinstance(node, LoopInputNode):
+                # Run one iteration if the found loop (recursive if nested).
+                self._run_loop(node, depsgraph, to_visit)
+                continue
+
             # Everything that this node depends on has run, so time to run it
             # itself. It can be taken off the queue.
             to_visit.popleft()
@@ -164,6 +226,7 @@ class AbstractRigNodesNode(bpy.types.Node):
         description="Error message; if empty, everything is fine",
         update=_on_node_error_updated,
     )
+    iterations: bpy.props.IntProperty(name="Iterations", default=1)  # type: ignore
 
     def init(self, context: bpy.types.Context) -> None:
         pass
@@ -179,6 +242,9 @@ class AbstractRigNodesNode(bpy.types.Node):
         else:
             layout.label(text="did not run", icon="PAUSE")
 
+        if context.preferences.view.show_developer_ui:
+            layout.prop(self, "iterations")
+
         if self.error:
             layout.label(text=self.error, icon="ERROR")
 
@@ -199,10 +265,29 @@ class AbstractRigNodesNode(bpy.types.Node):
         sock.hide_value = True
         return sock
 
+    @property
+    def _inputs(self):
+        return self.inputs
+
+    @property
+    def _outputs(self):
+        return self.outputs
+
+    def _get_nodes_between(
+        self, to_node: "AbstractRigNodesNode"
+    ) -> Iterable["AbstractRigNodesNode"]:
+        # TODO: the naming is bad here I think. Check that it's actually from -> to or to -> from
+        nodes = list(self._connected_nodes(self._inputs, "from_socket"))
+
+        if not to_node in nodes and len(nodes) > 0:
+            from_node = nodes[-1]
+            yield from from_node._get_nodes_between(to_node)
+        yield from nodes
+
     def exec_order_prerequisites(self) -> Iterable["AbstractRigNodesNode"]:
         """Generator, yields the nodes that should be executed before this one."""
         # For input execution order, consider all connected nodes regardless of socket types.
-        return self._connected_nodes(self.inputs, "from_socket")
+        return self._connected_nodes(self._inputs, "from_socket")
 
     def exec_order_successors(self) -> Iterable["AbstractRigNodesNode"]:
         """Generator, yields the nodes that should be executed after this one."""
@@ -213,7 +298,7 @@ class AbstractRigNodesNode(bpy.types.Node):
                 socket.node, AbstractAlwaysExecuteNode
             )
 
-        return self._connected_nodes(self.outputs, "to_socket", follow_socket)
+        return self._connected_nodes(self._outputs, "to_socket", follow_socket)
 
     def _connected_nodes(
         self,
@@ -259,6 +344,8 @@ class AbstractRigNodesNode(bpy.types.Node):
             self.error = str(ex)
             raise
         finally:
+            self.iterations -= 1
+            if self.iterations == 0:
                 self.has_run = True
 
     def __str__(self) -> str:
@@ -267,6 +354,7 @@ class AbstractRigNodesNode(bpy.types.Node):
     def reset_run(self) -> None:
         self.has_run = False
         self.error = ""
+        self.iterations = 1
 
     def _get_input_socket_value(self, input_socket: bpy.types.NodeSocket) -> object:
         for link in input_socket.links:
@@ -415,6 +503,222 @@ class AbstractRigNodesNode(bpy.types.Node):
                 tree.links.new(from_socket, to_socket)
 
 
+def _on_update_recreate_node(self: "SequenceNode", context: bpy.types.Context) -> None:
+    self.recreate(context)
+
+
+_enum_loop_dtype = (
+    ("OBJECT", "Object", ""),
+    ("VECTOR", "Vector", ""),
+    ("QUATERNION", "Rotation", ""),
+    ("VALUE", "Value", ""),
+)
+
+
+class LoopInputNode(AbstractRigNodesNode):
+    bl_idname = "LoopInputNode"
+    bl_label = "Loop Input Node"
+    bl_icon = "LOOP_FORWARDS"
+
+    dtype: bpy.props.EnumProperty(  # type: ignore
+        name="Type",
+        items=_enum_loop_dtype,
+        default="OBJECT",
+        update=_on_update_recreate_node,
+    )
+
+    num_socks: bpy.props.IntProperty(  # type: ignore
+        name="Number of Sockets",
+        default=2,
+        update=_on_update_recreate_node,
+    )
+
+    index: bpy.props.IntProperty(name="Index", default=0)  # type: ignore
+
+    def draw_buttons(
+        self, context: bpy.types.Context, layout: bpy.types.UILayout
+    ) -> None:
+        super().draw_buttons(context, layout)
+        layout.prop(self, "dtype")
+        layout.prop(self, "num_socks")
+        layout.prop(self, "index")
+
+    # 67649: The socket_value_update fn does not work when changing values in a node tree.
+    # Therefore using update for prototying.
+    def update(self):
+        if "Loop Connection" in self.outputs:
+            setattr(self.outputs["Loop Connection"], "default_value", self.num_socks)
+
+    def _get_connection_node(self) -> "LoopOutputNode":
+        # The 'Loop Connection' should be connected to a 'Loop Input Node',
+        # otherwise the loop won't be executed.
+        for socket in self.outputs[:1]:
+            for link in socket.links[:1]:
+                connected_socket = getattr(link, "to_socket")
+                if not hasattr(connected_socket, "node"):
+                    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.")
+
+    def reset_run(self):
+        super().reset_run()
+        self.index = 0
+
+    def init(self, context: bpy.types.Context) -> None:
+        for index in range(self.num_socks):
+            match self.dtype:
+                case "OBJECT":
+                    self.inputs.new("NodeSocketObject", f"Object {index+1}")
+                case "VECTOR":
+                    self.inputs.new("NodeSocketVector", f"Vector {index+1}")
+                case "QUATERNION":
+                    self.inputs.new("NodeSocketQuaternion", f"Rotation {index+1}")
+                case "VALUE":
+                    self.inputs.new("NodeSocketFloat", f"Value {index+1}")
+
+        self.outputs.new("NodeSocketInt", "Loop Connection")
+
+        match self.dtype:
+            case "OBJECT":
+                self.outputs.new("NodeSocketObject", f"Object")
+            case "VECTOR":
+                self.outputs.new("NodeSocketVector", f"Vector")
+            case "QUATERNION":
+                self.outputs.new("NodeSocketQuaternion", f"Rotation")
+            case "VALUE":
+                self.outputs.new("NodeSocketFloat", f"Value")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        if self.index == self.num_socks:
+            self.index = 0
+
+        self.index += 1
+        index = self.index
+
+        match self.dtype:
+            case "OBJECT":
+                obj = self._get_input_value(f"Object {index}", bpy.types.Object)
+                self.outputs["Object"].default_value = obj
+            case "VECTOR":
+                value = self._get_input_value(f"Vector {index}", Vector)
+                self.outputs["Vector"].default_value = value
+            case "QUATERNION":
+                value = self._get_input_value(f"Quaternion {index}", Quaternion)
+                self.outputs["Rotation"].default_value = value
+            case "VALUE":
+                value = self._get_input_value(f"Value {index}", float)
+                self.outputs["Value"].default_value = value
+
+
+class LoopOutputNode(AbstractRigNodesNode):
+    bl_idname = "LoopOutputNode"
+    bl_label = "Loop Output Node"
+    bl_icon = "LOOP_BACK"
+
+    dtype: bpy.props.EnumProperty(  # type: ignore
+        name="Type",
+        items=_enum_loop_dtype,
+        default="OBJECT",
+        update=_on_update_recreate_node,
+    )
+
+    num_socks: bpy.props.IntProperty(  # type: ignore
+        name="Number of Sockets",
+        default=2,
+        update=_on_update_recreate_node,
+    )
+
+    index: bpy.props.IntProperty(name="Index", default=0)  # type: ignore
+
+    def draw_buttons(
+        self, context: bpy.types.Context, layout: bpy.types.UILayout
+    ) -> None:
+        super().draw_buttons(context, layout)
+        layout.prop(self, "dtype")
+        layout.prop(self, "index")
+
+    def _get_connection_node(self) -> "LoopInputNode":
+        # The 'Loop Connection' should be connected to a 'Loop Input Node',
+        # otherwise the loop won't be executed.
+        for socket in self.inputs[:1]:
+            for link in socket.links[:1]:
+                connected_socket = getattr(link, "from_socket")
+                if not hasattr(connected_socket, "node"):
+                    raise RuntimeError("Loop Input Node is required to execute a loop.")
+
+                if isinstance(connected_socket.node, LoopInputNode):
+                    return connected_socket.node
+        raise RuntimeError("Loop Input Node is required to execute a loop.")
+
+    @AbstractRigNodesNode._inputs.getter  # type: ignore
+    def _inputs(self):
+        return self.inputs[1:]
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketInt", "Loop Connection")
+        match self.dtype:
+            case "OBJECT":
+                self.inputs.new("NodeSocketObject", f"Object")
+            case "VECTOR":
+                self.inputs.new("NodeSocketVector", f"Vector")
+            case "QUATERNION":
+                self.inputs.new("NodeSocketQuaternion", f"Rotation")
+            case "VALUE":
+                self.inputs.new("NodeSocketFloat", f"Value")
+
+        for index in range(self.num_socks):
+            match self.dtype:
+                case "OBJECT":
+                    self.outputs.new("NodeSocketObject", f"Object {index+1}")
+                case "VECTOR":
+                    self.outputs.new("NodeSocketVector", f"Vector {index+1}")
+                case "QUATERNION":
+                    self.outputs.new("NodeSocketQuaternion", f"Rotation {index+1}")
+                case "VALUE":
+                    self.outputs.new("NodeSocketFloat", f"Value {index+1}")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        if self.index == self.num_socks:
+            self.index = 0
+
+        self.index += 1
+        index = self.index
+
+        match self.dtype:
+            case "OBJECT":
+                obj = self._get_input_value(f"Object", bpy.types.Object)
+                self.outputs[f"Object {index}"].default_value = obj
+            case "VECTOR":
+                obj = self._get_input_value(f"Vector", Vector)
+                self.outputs[f"Vector {index}"].default_value = obj
+            case "QUATERNION":
+                obj = self._get_input_value(f"Rotation", Quaternion)
+                self.outputs[f"Rotation {index}"].default_value = obj
+            case "VALUE":
+                obj = self._get_input_value(f"Value", float)
+                self.outputs[f"Value {index}"].default_value = obj
+
+    def reset_run(self) -> None:
+        super().reset_run()
+        self.index = 0
+        self.iterations = self.num_socks
+
+    # 67649: The socket_value_update fn does not work when changing values in a node tree.
+    # Therefore using update for prototying the loop function.
+    # This allows to dynamically change the sockets based loop input socket amount.
+    def update(self):
+        if "Loop Connection" not in self.inputs:
+            return
+
+        val = self._get_optional_input_value("Loop Connection", int)
+        if val != self.num_socks and isinstance(val, int):
+            self.num_socks = val
+
+
 class AbstractRigNodesEventNode(AbstractRigNodesNode):
     """Node that can only exist once in a tree."""
 
@@ -1232,10 +1536,6 @@ class ClampNode(AbstractRigNodesNode):
         self.outputs["Result"].default_value = clamped
 
 
-def _on_num_sockets_change(self: "SequenceNode", context: bpy.types.Context) -> None:
-    self.recreate(context)
-
-
 class SequenceNode(AbstractRigNodesNode):
     """Multiple 'Execute' node sockets."""
 
@@ -1246,7 +1546,7 @@ class SequenceNode(AbstractRigNodesNode):
     num_socks: bpy.props.IntProperty(  # type: ignore
         name="Number of Sockets",
         default=2,
-        update=_on_num_sockets_change,
+        update=_on_update_recreate_node,
     )
 
     def draw_buttons(
@@ -1307,6 +1607,8 @@ node_categories = [
         "Flow",
         items=[
             nodeitems_utils.NodeItem("SequenceNode"),
+            nodeitems_utils.NodeItem("LoopInputNode"),
+            nodeitems_utils.NodeItem("LoopOutputNode"),
         ],
     ),
     RigNodesNodeCategory(
@@ -1371,6 +1673,8 @@ classes = (
     TwoBoneIKNode,
     SetCursorNode,
     SequenceNode,
+    LoopInputNode,
+    LoopOutputNode,
     # Math Nodes
     RotateTowards,
     AngleFromVectors,

operators.py

@@ -56,6 +56,27 @@ class RigNodes_OT_snap_to_bone(bpy.types.Operator):
         snap_to_bone(context.object, context.view_layer.depsgraph)
         return {"FINISHED"}
 
+class RigNodes_OT_prepare_tree(bpy.types.Operator):
+    """Prepare node tree manually."""
 
-classes = (RigNodes_OT_snap_to_bone,)
+    bl_idname = "rignodes.prepare_tree"
+    bl_label = "Prepare Tree"
+    bl_options = {"REGISTER", "UNDO"}
+
+    @classmethod
+    def poll(cls, context: bpy.types.Context) -> bool:
+        return bool(context.object) and context.preferences.view.show_developer_ui
+    
+    def execute(self, context: bpy.types.Context) -> set[str]:
+        for tree in bpy.data.node_groups:
+            if tree.bl_idname != "RigNodesNodeTree":
+                continue
+            if tree.is_running:
+                continue
+            tree._prepare_nodes()
+
+        return {"FINISHED"}
+
+
+classes = (RigNodes_OT_snap_to_bone, RigNodes_OT_prepare_tree,)
 register, unregister = bpy.utils.register_classes_factory(classes)