diff --git a/gui.py b/gui.py index 59fe862..2f468b4 100644 --- a/gui.py +++ b/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) diff --git a/nodes.py b/nodes.py index 9fb4f1e..2793282 100644 --- a/nodes.py +++ b/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,7 +375,9 @@ class AbstractRigNodesNode(bpy.types.Node): self.error = str(ex) raise finally: - self.has_run = True + self.iterations -= 1 + if self.iterations <= 0: + self.has_run = True def __str__(self) -> str: return f"{self.__class__.__name__}({self.name!r})" @@ -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, diff --git a/operators.py b/operators.py index 1951982..eea3a27 100644 --- a/operators.py +++ b/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) diff --git a/rignode-loops.blend b/rignode-loops.blend new file mode 100644 index 0000000..b3d7362 Binary files /dev/null and b/rignode-loops.blend differ