95 lines
3.7 KiB
Python
95 lines
3.7 KiB
Python
"""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}
|