AnyImage/tools/nodes/arrangement/flow.py

"""Infer geometry flow and socket relationships for layout and preview."""

import json
from collections import defaultdict
from graphlib import TopologicalSorter
from pathlib import Path

from .dimensions import visible_socket

SEMANTICS = json.loads(Path(__file__).with_name("semantics.json").read_text(encoding="utf-8"))
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 graph(node_group):
    def key(node):
        if node.bl_idname in {"NodeGroupInput", "NodeReroute"}:
            return (1, tuple(sorted((link.to_node.name, link.to_socket.identifier,
                                     socket_links(link.to_socket).index(link))
                                    for socket in node.outputs for link in socket.links)))
        return (0, node.name)

    nodes = sorted(node_group.nodes, key=key)
    parents = {node: [] for node in nodes}
    children = {node: set() for node in nodes}
    for link in node_group.links:
        if link.from_node not in parents[link.to_node]:
            parents[link.to_node].append(link.from_node)
        children[link.from_node].add(link.to_node)
    for sources in parents.values():
        sources.sort(key=key)
    order = list(TopologicalSorter(parents).static_order())
    ancestors = {}
    for node in order:
        ancestors[node] = set(parents[node]) | set().union(*(ancestors[p] for p in parents[node]))
    return parents, children, order, ancestors



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 shared_geometry_outputs(nodes):
    """Find geometry outputs reused by several distinct operations."""
    shared = []
    for node in nodes:
        for socket in node.outputs:
            if socket.type != "GEOMETRY":
                continue
            targets, pending = set(), [link.to_node for link in socket.links]
            while pending:
                target = pending.pop()
                if target.bl_idname == "NodeReroute":
                    pending.extend(link.to_node for link in target.outputs[0].links)
                else:
                    targets.add(target)
            if len(targets) >= SHARED_GEOMETRY_USES:
                shared.append(socket)
    return shared



def peer_links(node):
    """Read ordered data operands from the semantic table, excluding controls."""
    rule = SEMANTICS.get(node.bl_idname)
    if rule is None or ("operations" in rule and node.operation not in rule["operations"]):
        return []
    return [link for socket in list(node.inputs)[rule["peer_inputs_from"]:]
            if visible_socket(socket) and socket.type != "CUSTOM"
            for link in socket_links(socket)]


def node_roles(node_group):
    """Classify by contribution to the result, not by available socket types."""
    core, _ = geometry_paths(node_group)
    return {node: "geometry" if node in core else "field" for node in node_group.nodes}