Initial public release

tools/nodes/arrangement/flow.py

@@ -0,0 +1,94 @@
+"""Infer geometry flow and socket relationships for layout and preview."""
+
+from collections import defaultdict
+
+from .branches import calculation_blocks
+
+
+SHARED_GEOMETRY_USES = 3
+
+
+def socket_links(socket):
+    return sorted(socket.links, key=lambda link: link.multi_input_sort_id, reverse=True)
+
+
+def origin_socket(socket):
+    while socket.node.bl_idname == "NodeReroute":
+        socket = socket.node.inputs[0].links[0].from_socket
+    return socket
+
+
+def geometry_paths(node_group):
+    """Trace every geometry input and split paths at forks and merges."""
+    output = next((node for node in node_group.nodes if node.bl_idname == "NodeGroupOutput" and node.is_active_output), None)
+    parents = {}
+    uses = defaultdict(set)
+    for link in node_group.links:
+        if link.from_socket.type == "GEOMETRY" and link.to_node.bl_idname != "NodeReroute":
+            uses[origin_socket(link.from_socket).node].add(link.to_node)
+    paired = {member for node in node_group.nodes if getattr(node, "paired_output", None) is not None
+              for member in (node, node.paired_output)}
+    result_nodes, pending = set(), [*paired, *([output] if output else [])]
+    while pending:
+        node = pending.pop()
+        if node in result_nodes:
+            continue
+        result_nodes.add(node)
+        pending.extend(origin_socket(link.from_socket).node
+                       for socket in node.inputs if socket.type == "GEOMETRY" for link in socket.links)
+    shared = {node for node, targets in uses.items()
+              if len(targets) >= SHARED_GEOMETRY_USES and node not in paired
+              and (not any(socket.type == "GEOMETRY" and socket.is_linked for socket in node.inputs)
+                   or len(targets & result_nodes) >= SHARED_GEOMETRY_USES)}
+
+    def source_node(link):
+        return link.from_node if link.from_node.get("arrangement_input") else origin_socket(link.from_socket).node
+
+    def collect(node):
+        if node in parents:
+            return
+        parents[node] = list(dict.fromkeys(
+            source_node(link) for socket in node.inputs if socket.type == "GEOMETRY"
+            for link in socket_links(socket)
+            if source_node(link).bl_idname != "NodeGroupInput" and source_node(link) not in shared
+        )) if not node.get("arrangement_input") else []
+        for parent in parents[node]:
+            collect(parent)
+
+    roots = []
+    if output is not None:
+        collect(output)
+        roots.append(output)
+    for node in node_group.nodes:
+        if node in paired and node not in parents:
+            collect(node)
+            roots.append(node)
+    children = {node: set() for node in parents}
+    for node, sources in parents.items():
+        for source in sources:
+            children[source].add(node)
+    paths, assigned, pending = [], set(), roots
+    while pending:
+        node = pending.pop(0)
+        path = []
+        while node not in assigned:
+            path.append(node)
+            assigned.add(node)
+            sources = parents[node]
+            if len(sources) != 1 or len(children[sources[0]]) != 1:
+                pending.extend(source for source in sources if source not in assigned)
+                break
+            node = sources[0]
+        if path:
+            paths.append(list(reversed(path)))
+    return parents, paths
+
+
+def node_roles(node_group):
+    """Classify by contribution to the result, not by available socket types."""
+    if node_group.bl_idname == "ShaderNodeTree":
+        nodes = list(node_group.nodes)
+        _, core = calculation_blocks(nodes, list(node_group.links))
+        return {node: "shader" if i in core else "field" for i, node in enumerate(nodes)}
+    core, _ = geometry_paths(node_group)
+    return {node: "geometry" if node in core else "field" for node in node_group.nodes}