gui.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: GPL-2.0-or-later
 
 import bpy
+from typing import List
 
 
 class RIGNODES_MT_snap_pie(bpy.types.Menu):
@@ -54,12 +55,13 @@ 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")
 
 
-addon_keymaps = []
+addon_keymaps: List[bpy.types.KeyMap] = []
 classes = (RIGNODES_MT_snap_pie,)
 _register, _unregister = bpy.utils.register_classes_factory(classes)
 

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,86 @@ 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 = (node for node in self.nodes if isinstance(node, LoopOutputNode))
+
+        # Set iteration count foreach (nested) loop.
+        for output_node in output_nodes:
+            input_node = output_node._get_connection_node()
+            for node in output_node._get_innern_loop_nodes(input_node):
+                node.iterations *= output_node.num_socks
+
+    def _run_loop(
+        self,
+        node: "LoopInputNode",
+        depsgraph: bpy.types.Depsgraph,
+        to_visit: deque["AbstractRigNodesNode"],
+        depth: int = 1,
+    ) -> deque["AbstractRigNodesNode"]:
+        """Runs (nested) loops.
+
+        :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()
+        visited: Set["AbstractRigNodesNode"] = set()
+
+        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_innern_loop_nodes(node))
+                nested_loop.appendleft(node)
+                nested_loop.append(nested_output_node)
+
+                # Don't revisit nested nodes in the parent loop.
+                for nested_node in nested_loop:
+                    visited.add(nested_node)
+
+                # Run the nested loop using the current inputs.
+                for _ in range(node.num_socks):
+                    self._run_loop(node, depsgraph, nested_loop.copy(), depth + 1)
+                continue
+
+            print(f"{indent}\tRunning: {node}, remaining runs: {node.iterations}")
+            visited.add(node)
+            node.run(depsgraph)
+
+            if node == output_node:
+                break
+
+        return to_visit
+
     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- Visiting {node}")
 
             nodes_before = list(node.exec_order_prerequisites())
             if nodes_before:
@@ -88,6 +153,25 @@ class RigNodesNodeTree(bpy.types.NodeTree):
                 to_visit.extendleft(reversed(nodes_before))
                 continue
 
+            if isinstance(node, LoopInputNode):
+                if node in visited:
+                    to_visit.popleft()
+                    continue
+
+                # Find the loop.
+                output_node = node._get_connection_node()
+                innern_loop = deque(output_node._get_innern_loop_nodes(node))
+
+                # Add (nested) loop nodes to visited.
+                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())
+
             # Everything that this node depends on has run, so time to run it
             # itself. It can be taken off the queue.
             to_visit.popleft()
@@ -99,6 +183,7 @@ class RigNodesNodeTree(bpy.types.NodeTree):
                 continue
             visited.add(node)
 
+            print(f"\tRunning: {node}")
             node.run(depsgraph)
 
             # Queue up the next nodes.
@@ -164,6 +249,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 +265,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 +288,37 @@ 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_innern_loop_nodes(
+        self: "LoopOutputNode",
+        input_node: "LoopInputNode",
+    ) -> Iterable["AbstractRigNodesNode"]:
+        """Generator, yields nodes between the loop input and the loop output node"""
+        visited: Set["AbstractRigNodesNode"] = set()
+        to_visit = [self]
+
+        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."""
         # 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 +329,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 +375,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 +385,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 +534,238 @@ class AbstractRigNodesNode(bpy.types.Node):
                 tree.links.new(from_socket, to_socket)
 
 
+class ObjectInputNode(AbstractRigNodesNode):
+    bl_idname = "ObjectInputNode"
+    bl_label = "Object Input Node"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketObject", f"Object")
+        self.outputs.new("NodeSocketObject", f"Object")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        obj = self._get_input_value(f"Object", bpy.types.Object)
+        self.outputs["Object"].default_value = obj
+
+
+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")
+
+    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.")
+
+    # 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 reset_run(self):
+        super().reset_run()
+        self.index = 0
+        self.iterations = self.num_socks
+
+    @AbstractRigNodesNode._outputs.getter  # type: ignore
+    def _outputs(self):
+        return self.outputs[1:]
+
+    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")
+
+    @AbstractRigNodesNode._inputs.getter  # type: ignore
+    def _inputs(self):
+        return self.inputs[1:]
+
+    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.")
+
+    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
+
+    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
+
+
 class AbstractRigNodesEventNode(AbstractRigNodesNode):
     """Node that can only exist once in a tree."""
 
@@ -1232,10 +1583,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 +1593,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 +1654,15 @@ node_categories = [
         "Flow",
         items=[
             nodeitems_utils.NodeItem("SequenceNode"),
+            nodeitems_utils.NodeItem("LoopInputNode"),
+            nodeitems_utils.NodeItem("LoopOutputNode"),
+        ],
+    ),
+    RigNodesNodeCategory(
+        "INPUTS",
+        "Inputs",
+        items=[
+            nodeitems_utils.NodeItem("ObjectInputNode"),
         ],
     ),
     RigNodesNodeCategory(
@@ -1371,6 +1727,9 @@ classes = (
     TwoBoneIKNode,
     SetCursorNode,
     SequenceNode,
+    LoopInputNode,
+    LoopOutputNode,
+    ObjectInputNode,
     # 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 = "Reset 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)