AnyImage/tools/node_assets/node_layout.py

"""Shared presentation rules for authored geometry-node layouts."""


def arrange_nodes(group, positions, input_regions):
    """Place authored branches and rebuild their local Group Inputs."""
    controls = [
        (link.to_socket, link.from_socket.name)
        for link in group.links
        if link.from_node.bl_idname == "NodeGroupInput"
    ]
    for node in list(group.nodes):
        if node.bl_idname == "NodeGroupInput":
            group.nodes.remove(node)
    assert set(positions) == {node.name for node in group.nodes}, group.name
    for name, location in positions.items():
        group.nodes[name].location = location
    for targets, location in input_regions:
        node = group.nodes.new("NodeGroupInput")
        node.location = location
        for socket, name in controls:
            if socket.node.name in targets:
                group.links.new(node.outputs[name], socket)
    assert all(socket.is_linked for socket, _ in controls), group.name
    for node in group.nodes:
        node.select = False
        node.width = (
            200
            if "Attribute" in node.bl_idname or node.bl_idname == "GeometryNodeGroup"
            else 160
        )
        node.hide = node.bl_idname in {
            "ShaderNodeMath",
            "ShaderNodeVectorMath",
        } and all(
            socket.is_linked
            for socket in node.inputs
            if socket.enabled and not socket.hide
        )
        for socket in node.outputs:
            socket.hide = not socket.is_linked


def validate_node_layout(group):
    """Check visible bounds and sockets, allowing direct backward references."""
    nodes = [n for n in group.nodes if n.bl_idname not in {"NodeReroute", "NodeFrame"}]
    assert all(not n.label for n in nodes), group.name
    assert not any(n.bl_idname == "FunctionNodeInputString" for n in nodes), group.name
    assert all(s.is_linked or s.hide for n in nodes for s in n.outputs), group.name
    validate_peer_inputs(group)
    for link in group.links:
        if (
            link.from_node.bl_idname != "NodeReroute"
            or link.to_node.bl_idname == "NodeReroute"
        ):
            continue
        source = link.from_node
        while source.bl_idname == "NodeReroute":
            source = source.inputs[0].links[0].from_node
        assert source.location.x <= link.to_node.location.x, (
            group.name,
            source.name,
            link.to_node.name,
        )
    for index, node in enumerate(nodes):
        width = node.dimensions.x or node.width
        height = node.dimensions.y or (30 if node.hide else node.height)
        for other in nodes[index + 1 :]:
            other_width = other.dimensions.x or other.width
            other_height = other.dimensions.y or (30 if other.hide else other.height)
            overlap = (
                node.location.x < other.location.x + other_width
                and node.location.x + width > other.location.x
                and node.location.y - height < other.location.y
                and node.location.y > other.location.y - other_height
            )
            assert not overlap, (group.name, node.name, other.name)
        if node.bl_idname == "GeometryNodeRepeatInput":
            assert (
                node.paired_output and node.paired_output.location.x > node.location.x
            )


def validate_peer_inputs(group):
    """Align peer branch ends while leaving shared upstream sources in place."""
    parents = {node: set() for node in group.nodes}
    consumers = {node: set() for node in group.nodes}
    for link in group.links:
        parents[link.to_node].add(link.from_node)
        consumers[link.from_node].add(link.to_node)

    def ancestors(node):
        found = set()
        pending = list(parents[node])
        while pending:
            parent = pending.pop()
            if parent in found:
                continue
            found.add(parent)
            pending.extend(parents[parent])
        return found

    for node in group.nodes:
        if node.bl_idname not in {
            "GeometryNodeJoinGeometry",
            "GeometryNodeMenuSwitch",
            "GeometryNodeSwitch",
            "ShaderNodeCombineXYZ",
        }:
            continue
        combine = node.bl_idname == "ShaderNodeCombineXYZ"
        sources = set()
        for socket in node.inputs:
            if not combine and socket.type != "GEOMETRY":
                continue
            for link in socket.links:
                source = link.from_node
                while source.bl_idname == "NodeReroute":
                    source = source.inputs[0].links[0].from_node
                if combine or source.bl_idname != "NodeGroupInput":
                    sources.add(source)
        upstream = set().union(*(ancestors(source) for source in sources))
        peers = sources - upstream
        if combine:
            peers = {source for source in peers if consumers[source] == {node}}
        edges = [source.location.x + source.width for source in peers]
        assert not edges or max(edges) - min(edges) < 1.0, (
            group.name,
            node.name,
            [source.name for source in peers],
        )
        if node.bl_idname == "GeometryNodeMenuSwitch":
            rows = [
                link.from_node.location.y
                for socket in node.inputs
                for link in socket.links
                if link.from_node in peers
            ]
            assert rows == sorted(rows, reverse=True), (group.name, node.name)