AnyImage/tools/nodes/arrangement/arrange.py

"""Arrange geometry paths, local fields, and shared references after construction."""

from collections import defaultdict
from dataclasses import dataclass, field
from graphlib import TopologicalSorter

import numpy as np
from scipy.optimize import LinearConstraint, minimize

from .dimensions import estimate_node_dimensions, visible_socket
from .zones import node_zones
from .flow import graph, geometry_paths, origin_socket, peer_links, shared_geometry_outputs, socket_links
from .connections import curve_footprint, link_endpoints
from .fields import field_layout


GEOMETRY_GAP = 100
FIELD_GAP = 40
BRANCH_GAP = 160
SECTION_GAP = 200
FIELD_HEIGHT = 720
ROUTE_GAP = 40
ROUTE_RUN = 120


def _replace_link(node_group, link, source):
    """Replace a source without changing a multi-input socket's branch order."""
    multi = link.to_socket.is_multi_input
    replacement = node_group.links.new(source, link.to_socket)
    if multi:
        replacement.swap_multi_input_sort_id(link)
        node_group.links.remove(link)


def _remove_routes(node_group):
    previous = [node for node in node_group.nodes if node.bl_idname == "NodeReroute"]

    for node in [node for node in node_group.nodes if node not in previous]:
        for socket in node.inputs:
            links = socket_links(socket)
            if not any(link.from_node in previous for link in links):
                continue
            if not socket.is_multi_input:
                _replace_link(node_group, links[0], origin_socket(links[0].from_socket))
                continue
            sources = [origin_socket(link.from_socket) for link in links]
            for link in links:
                node_group.links.remove(link)
            seen = set()
            for source in reversed(sources):
                original = source
                if source in seen:
                    # Blender coalesces identical direct links; retain repeated inputs.
                    duplicate = node_group.nodes.new("NodeReroute")
                    node_group.links.new(source, duplicate.inputs[0])
                    source = duplicate.outputs[0]
                seen.add(original)
                node_group.links.new(source, socket)
    for node in previous:
        node_group.nodes.remove(node)


def _peer_inputs(node, ancestors):
    sources = list(dict.fromkeys(
        link.from_node for link in peer_links(node)
        if link.from_node.bl_idname != "NodeGroupInput"
    ))
    upstream = set().union(*(ancestors[source] for source in sources))
    return [source for source in sources if source not in upstream]


def _localize_peer_inputs(node_group, ancestors, children):
    """Give shared inputs local ports without moving their public source."""
    for node in list(node_group.nodes):
        links = peer_links(node)
        if len(links) < 2:
            continue
        sources = [link.from_node for link in links]
        upstream = set().union(*(ancestors[source] for source in sources))
        for link in links:
            source = link.from_node
            if source.bl_idname == "NodeReroute":
                source["arrangement_input"] = True
                continue
            if source.bl_idname != "NodeGroupInput" and len(children[source]) <= 1 and sources.count(source) == 1 and source not in upstream:
                continue
            port = node_group.nodes.new("NodeReroute")
            port["arrangement_input"] = True
            node_group.links.new(link.from_socket, port.inputs[0])
            _replace_link(node_group, link, port.outputs[0])


def _columns(nodes, parents, peers, sizes, gap, spacing=None):
    """Allocate columns with equal right edges for independent peer inputs."""
    roots = {node: node for node in nodes}

    def root(node):
        while roots[node] != node:
            node = roots[node]
        return node

    for sources in peers:
        members = [source for source in sources if source in roots]
        for source in members[1:]:
            roots[root(source)] = root(members[0])
    dependencies = {root(node): [] for node in nodes}
    for node in nodes:
        for parent in parents[node]:
            if parent in roots and root(parent) not in dependencies[root(node)]:
                dependencies[root(node)].append(root(parent))
    levels = {}
    for component in TopologicalSorter(dependencies).static_order():
        levels[component] = max((levels[p] + 1 for p in dependencies[component]), default=0)
    columns = defaultdict(list)
    for node in nodes:
        columns[levels[root(node)]].append(node)
    right, positions = 0, {}
    for level in sorted(columns):
        right += max(((spacing or {}).get(node, 0) for node in columns[level]), default=0)
        right += max(sizes[node].width for node in columns[level])
        for node in columns[level]:
            positions[node] = right - sizes[node].width
        right += gap
    return positions, columns


def _field_regions(nodes, core, children, order):
    """Group calculations by their first geometry consumers and shared scope."""
    targets = {}
    fields = set(nodes) - core
    for node in reversed(order):
        targets[node] = frozenset({node}) if node in core else frozenset().union(
            *(targets[child] for child in children[node])
        )
    for node in order:
        if node.get("arrangement_exit"):
            targets[node] = targets[origin_socket(node.outputs[0]).node]
    regions = defaultdict(list)
    for node in order:
        if node in fields:
            regions[targets[node]].append(node)
    return [(region, owners) for owners, region in regions.items()]


def _localize_module_outputs(node_group, core, children):
    """Expose reused results at the boundary of their calculation module."""
    shared = set(shared_geometry_outputs(node_group.nodes))
    for link in list(node_group.links):
        source = link.from_node
        if link.to_node not in core or source in core or source.bl_idname == "NodeGroupInput":
            continue
        if link.from_socket in shared or not any(node not in core for node in children[source]):
            continue
        port = node_group.nodes.new("NodeReroute")
        port["arrangement_exit"] = True
        node_group.links.new(link.from_socket, port.inputs[0])
        _replace_link(node_group, link, port.outputs[0])


def _localize_inputs(node_group):
    """Give each operation local interface reads, separating geometry from values."""
    controls = [node for node in node_group.nodes if node.bl_idname == "NodeGroupInput"]
    uses = defaultdict(list)
    for link in node_group.links:
        if link.from_node in controls:
            uses[link.to_node, link.from_socket.type == "GEOMETRY"].append((link, link.from_socket.identifier))
    for inputs in uses.values():
        local = node_group.nodes.new("NodeGroupInput")
        for link, identifier in inputs:
            source = next(socket for socket in local.outputs if socket.identifier == identifier)
            _replace_link(node_group, link, source)
    for node in controls:
        node_group.nodes.remove(node)


def _presentation(node_group):
    for node in node_group.nodes:
        node.select = False
        for socket in node.outputs:
            socket.hide = not socket.is_linked
        if node.bl_idname in {"ShaderNodeMath", "ShaderNodeVectorMath"}:
            node.hide = all(socket.is_linked for socket in node.inputs if visible_socket(socket))


def _place_geometry_paths(paths, geometry_x, sizes, order):
    """Reuse fixed horizontal lanes; field sizes never participate in lane choice."""
    paths = [list(path) for path in paths]
    for path in list(paths):
        for node in path:
            paired = getattr(node, "paired_output", None)
            if paired is None or paired not in geometry_x:
                continue
            first = next(nodes for nodes in paths if node in nodes)
            last = next(nodes for nodes in paths if paired in nodes)
            if first is not last:
                first.extend(last)
                paths.remove(last)
    pitch = max(sizes[node].height for path in paths for node in path) + SECTION_GAP
    path_of = {node: index for index, path in enumerate(paths) for node in path}
    preceding = {index: set() for index in range(len(paths))}
    for node in order:
        inputs = list(dict.fromkeys(path_of[link.from_node] for link in peer_links(node)
                                  if link.from_node in path_of))
        for first, second in zip(inputs, inputs[1:]):
            pending, seen = [first], set()
            while pending:
                current = pending.pop()
                if current not in seen:
                    seen.add(current)
                    pending.extend(preceding[current])
            if second not in seen:
                preceding[second].add(first)
    lanes, positions, rows = defaultdict(list), {}, {}
    pending = set(preceding)
    while pending:
        ready = [index for index in pending if not preceding[index] & pending]
        index = min(ready, key=lambda index: (-len(paths[index]), geometry_x[paths[index][0]], index))
        path = paths[index]
        left = min(geometry_x[node] for node in path)
        right = max(geometry_x[node] + sizes[node].width for node in path)
        minimum = max((rows[previous] + 1 for previous in preceding[index]), default=-len(paths))
        for lane in sorted(range(minimum, len(paths) * 2), key=lambda value: (abs(value), value)):
            if all(right + FIELD_GAP <= a or left >= b + FIELD_GAP for a, b in lanes[lane]):
                break
        lanes[lane].append((left, right))
        rows[index] = lane
        pending.remove(index)
        for node in path:
            positions[node] = (geometry_x[node], -lane * pitch)
    # Solve the lane heights together so successive merges cannot create a staircase.
    equalities, centers, inequalities, separations = [], [], [], []
    for node in order:
        sources = [link.from_socket for link in peer_links(node)]
        if node not in path_of or len(sources) < 2 or any(socket.node not in path_of for socket in sources):
            continue
        output = next(socket for socket in node.outputs if socket in sizes[node].sockets)
        first, last = sources[0], sources[-1]
        row = np.zeros(len(paths))
        row[path_of[node]] += 2
        row[path_of[first.node]] -= 1
        row[path_of[last.node]] -= 1
        equalities.append(row)
        centers.append(2 * sizes[node].sockets[output] - sizes[first.node].sockets[first] - sizes[last.node].sockets[last])
    spans = [(min(geometry_x[node] for node in path), max(geometry_x[node] + sizes[node].width for node in path))
             for path in paths]
    local_ports = {index for index, path in enumerate(paths)
                   if all(node.get("arrangement_input") for node in path)}
    pairs = {(first, second) for second, sources in preceding.items() for first in sources}
    for first in range(len(paths)):
        for second in range(first + 1, len(paths)):
            if first in local_ports or second in local_ports:
                continue
            if spans[first][0] < spans[second][1] + FIELD_GAP and spans[first][1] + FIELD_GAP > spans[second][0]:
                pairs.add((first, second) if rows[first] < rows[second] else (second, first))
    for first, second in sorted(pairs):
        row = np.zeros(len(paths))
        row[first], row[second] = 1, -1
        inequalities.append(row)
        gap = FIELD_GAP if first in local_ports or second in local_ports else SECTION_GAP
        separations.append(max(sizes[node].height for node in paths[first]) + gap)
    weights = np.array([max(1, len(path)) ** 2 for path in paths], dtype=float)
    weights /= weights.max()
    constraints = []
    if equalities:
        constraints.append(LinearConstraint(np.array(equalities), centers, centers))
    if inequalities:
        constraints.append(LinearConstraint(np.array(inequalities), separations, np.inf))
    compact = []
    for node in order:
        if node not in path_of:
            continue
        for link in peer_links(node):
            if link.from_node in path_of and path_of[link.from_node] in local_ports:
                row = np.zeros(len(paths))
                row[path_of[link.from_node]], row[path_of[node]] = 1, -1
                compact.append(row)
    metric = np.diag(weights)
    if compact:
        matrix = np.array(compact)
        metric += matrix.T @ matrix * 4
    solution = minimize(lambda values: values @ metric @ values,
                        np.array([-rows[index] * pitch for index in range(len(paths))], dtype=float),
                        jac=lambda values: 2 * metric @ values, constraints=constraints,
                        method="SLSQP", options={"ftol": 1e-5, "maxiter": 300})
    if not solution.success:
        raise ValueError(f"Geometry lane constraints could not be placed: {solution.message}")
    for index, path in enumerate(paths):
        for node in path:
            positions[node] = (positions[node][0], float(solution.x[index]))
    for node in order:
        if node not in positions:
            continue
        links = [link for socket in node.inputs if socket.type == "GEOMETRY" for link in socket.links
                 if link.from_node in positions]
        if not links:
            continue
        x = positions[node][0]
        required = []
        for link in links:
            source = link.from_socket
            source_y = positions[source.node][1] - sizes[source.node].sockets[source]
            target_y = positions[node][1] - sizes[node].sockets[link.to_socket]
            required.append(positions[source.node][0] + sizes[source.node].width
                            + max(GEOMETRY_GAP, abs(target_y - source_y) * 0.8))
        extra = max(0, max(required) - x)
        if extra:
            for member, (left, y) in positions.items():
                if left >= x:
                    positions[member] = (left + extra, y)
    return positions


def _geometry_corridors(paths, positions, sizes):
    """Reserve continuous breathing room along every horizontal geometry path."""
    corridors = []
    for path in paths:
        if all(node.get("arrangement_input") for node in path):
            continue
        boxes = [_node_box(node, positions, sizes) for node in path]
        corridors.append((min(box[0] for box in boxes) - GEOMETRY_GAP / 2,
                          max(box[1] for box in boxes) + GEOMETRY_GAP / 2,
                          max(box[2] for box in boxes), min(box[3] for box in boxes)))
    for port in positions:
        if not port.get("arrangement_input") or port.outputs[0].type != "GEOMETRY":
            continue
        source = origin_socket(port.outputs[0]).node
        if source not in positions:
            continue
        left = positions[source][0] + sizes[source].width
        right, level = positions[port][0] + 8, positions[port][1] - 8
        if right > left:
            corridors.append((left, right, level + ROUTE_GAP, level - ROUTE_GAP))
    return corridors


def _node_box(node, positions, sizes):
    x, y = positions[node]
    return x, x + sizes[node].width, y, y - sizes[node].height


def _overlap(first, second, gap):
    a, b, top, bottom = first
    c, d, upper, lower = second
    return a < d + gap and b > c - gap and top > lower - gap and bottom < upper + gap


@dataclass
class BranchSpace:
    """Keep module bodies and their exit channels in one placement reservation."""

    bodies: list = field(default_factory=list)
    exits: list = field(default_factory=list)
    routes: list = field(default_factory=list)

    def connections(self, outlets, placed, sizes):
        exits, routes = [], []
        for links in outlets:
            for link in links:
                start, end = link_endpoints(link, placed, sizes)
                corners = []
                run = max(ROUTE_RUN, abs(end[1] - start[1]) * 0.8)
                if end[0] - start[0] > max(ROUTE_RUN * 4, run + ROUTE_RUN * 2):
                    corners = [(start[0] + ROUTE_RUN, start[1]),
                               (max(start[0] + ROUTE_RUN, end[0] - run), start[1])]
                points = [start, *dict.fromkeys(corners), end]
                boxes = [box for a, b in zip(points, points[1:]) for box in curve_footprint(a, b)]
                exits.append((link, boxes))
                routes.append((link, list(dict.fromkeys(corners))))
        return exits, routes

    def fits(self, body, exits):
        if any(_overlap(body, box, ROUTE_GAP) for _, boxes in self.exits for box in boxes):
            return False
        for link, boxes in exits:
            if any(_overlap(box, rectangle, ROUTE_GAP) for nodes, rectangle in self.bodies
                   if link.to_node not in nodes for box in boxes):
                return False
        return True

    def crossings(self, exits):
        return sum(any(_overlap(box, previous, 8) for box in boxes for previous in channel)
                   for link, boxes in exits for other, channel in self.exits
                   if link.to_node != other.to_node and link.from_socket != other.from_socket)


def _place_field_region(local, width, height, destinations, owners, positions, sizes, occupied, limits, core, trunks, space, horizontal=None):
    """Place an upstream module together with its forward-facing exit channels."""
    right = min(positions[node][0] for node in destinations) - FIELD_GAP
    baseline = sum(positions[node][1] for node in owners) / len(owners)
    output = next(node for node in core if node.bl_idname == "NodeGroupOutput" and node.is_active_output)
    ceiling = max(positions[node][1] for node in owners)
    floor = min(positions[node][1] - sizes[node].height for node in owners)
    axis_top = positions[output][1]
    axis_bottom = axis_top - sizes[output].height
    core_left = min(positions[node][0] for node in core)
    side = 1 if floor > axis_top else -1 if ceiling < axis_bottom else 0
    desired = right - width
    front_allowed = not horizontal and all(link.from_node in local
                    for node in local for socket in node.inputs for link in socket.links)
    shared = [] if horizontal else shared_geometry_outputs(local)
    references = set(shared_geometry_outputs(local))
    outlets, constraints, entries = [], [], []
    for node, (x, y) in local.items():
        if horizontal:
            for socket in node.inputs:
                for link in socket.links:
                    if link.from_node in positions:
                        source_x, source_y = positions[link.from_node]
                        entries.append((source_x + sizes[link.from_node].width - x,
                                        source_y - sizes[link.from_node].sockets[link.from_socket],
                                        y - sizes[node].sockets[socket]))
        for socket in node.outputs:
            links = list(socket.links)
            if socket not in references:
                returns = [link for link in links if link.to_node in core]
                if returns:
                    outlets.append(returns)
            for link in links:
                if link.to_node in positions and link.to_node not in local:
                    target_x, target_y = positions[link.to_node]
                    constraints.append((target_x - x - sizes[node].width,
                                        y - sizes[node].sockets[socket],
                                        target_y - sizes[link.to_node].sockets[link.to_socket]))
    trunk_end = max(positions[node][0] for node in {*owners, *destinations})
    candidates = {desired, *(left - BRANCH_GAP - width for left, _, _, _ in occupied),
                  *(end + BRANCH_GAP for _, end, _, _ in occupied)}
    if horizontal:
        candidates = {left for left in candidates if horizontal[0] <= left and left + width <= horizontal[1]}
        candidates.add(horizontal[0])
    best, seen = None, set()
    for left in sorted(candidates, key=lambda x: (abs(x - desired), x)):
        tops = {ceiling + BRANCH_GAP + height, floor - BRANCH_GAP,
                limits[0], limits[1] + height,
                *(top + BRANCH_GAP + height for _, _, top, _ in occupied),
                *(bottom - BRANCH_GAP for _, _, _, bottom in occupied)}
        if front_allowed:
            tops.add(baseline + height / 2)
        for top in sorted(tops, key=lambda y: abs(y - height / 2 - baseline)):
            if top > limits[0] or top - height < limits[1]:
                continue
            minimum = max([horizontal[0], *(start + max(GEOMETRY_GAP, abs(top + target_y - source_y) * 0.8)
                                            for start, source_y, target_y in entries)]) if horizontal else -float("inf")
            placed_left = max(left, minimum) if horizontal else min(left, desired)
            placed_left = min([placed_left, *(end - max(GEOMETRY_GAP, abs(target_y - top - source_y) * 0.8)
                                              for end, source_y, target_y in constraints)])
            if horizontal and not minimum <= placed_left <= horizontal[1] - width:
                continue
            in_front = placed_left + width + BRANCH_GAP <= core_left
            if not (in_front and front_allowed) and ((side == 1 and top - height < ceiling + BRANCH_GAP - 1e-5)
                                 or (side == -1 and top > floor - BRANCH_GAP + 1e-5)):
                continue
            point = (placed_left, top)
            if point in seen:
                continue
            seen.add(point)
            distance = abs(placed_left - desired)
            base_score = distance + abs(top - height / 2 - baseline) * 1.5
            if best is not None and base_score > best[0]:
                continue
            overlapping = [box for box in occupied
                           if placed_left < box[1] + BRANCH_GAP and placed_left + width > box[0] - BRANCH_GAP]
            if any(top > bottom - BRANCH_GAP + 1e-5 and top - height < upper + BRANCH_GAP - 1e-5
                   for _, _, upper, bottom in overlapping):
                continue
            channels = [(placed_left + local[socket.node][0] + sizes[socket.node].width,
                         top + local[socket.node][1] - sizes[socket.node].sockets[socket]) for socket in shared]
            if any(a < trunk_end and b > start and bottom - ROUTE_GAP < level < upper + ROUTE_GAP
                   for start, level in channels for a, b, upper, bottom in occupied):
                continue
            placed = positions | {node: (x + placed_left, y + top) for node, (x, y) in local.items()}
            exits, routes = space.connections(outlets, placed, sizes)
            body = (placed_left, placed_left + width, top, top - height)
            if not space.fits(body, exits):
                continue
            score = base_score + space.crossings(exits) * BRANCH_GAP * 4
            candidate = (score, distance, -top, placed_left, top)
            if best is None or candidate < best:
                best = candidate
                best_exits, best_routes = exits, routes
    if best is None:
        raise ValueError("The field region does not fit its reserved scope")
    _, _, _, left, top = best
    for node, (x, y) in local.items():
        positions[node] = (x + left, y + top)
    body = (left, left + width, top, top - height)
    occupied.append(body)
    space.bodies.append((set(local), body))
    space.exits.extend(best_exits)
    space.routes.extend(best_routes)
    for socket in shared:
        start = positions[socket.node][0] + sizes[socket.node].width
        level = positions[socket.node][1] - sizes[socket.node].sockets[socket]
        trunks[socket] = level
        occupied.append((start, trunk_end, level + ROUTE_GAP, level - ROUTE_GAP))


def _route_field_links(node_group, layouts, positions):
    """Materialize the passages reserved between expression columns."""
    protected = set()
    for layout in layouts:
        anchor = next(iter(layout.positions))
        dx = positions[anchor][0] - layout.positions[anchor][0]
        dy = positions[anchor][1] - layout.positions[anchor][1]
        for link, points in layout.routes:
            source = link.from_socket
            protected.add((source, link.to_socket))
            for x, y in points:
                node = node_group.nodes.new("NodeReroute")
                node["arrangement_branch"] = True
                node.location = (x + dx - 8, y + dy + 8)
                node_group.links.new(source, node.inputs[0])
                source = node.outputs[0]
            _replace_link(node_group, link, source)
    return protected


def _route_branch_exits(node_group, routes):
    """Materialize the channels reserved during placement, sharing trunk corners."""
    deliveries, protected = defaultdict(list), set()
    for link, points in routes:
        protected.add((link.from_socket, link.to_socket))
        deliveries[link.from_socket].append((link, points))
    for source, branches in deliveries.items():
        corners = {}
        for point in sorted({point for _, points in branches for point in points}):
            node = node_group.nodes.new("NodeReroute")
            node["arrangement_branch"] = True
            node.location = (point[0] - 8, point[1] + 8)
            node_group.links.new(source, node.inputs[0])
            source = corners[point] = node.outputs[0]
        for link, points in branches:
            if points:
                _replace_link(node_group, link, corners[points[-1]])
    return protected


def _route_peer_bypasses(node_group, sizes, positions, ancestors):
    """Use an entry and an exit corner for a branch that bypasses processing."""
    for port in list(node_group.nodes):
        if not port.get("arrangement_input") or port.outputs[0].type != "GEOMETRY":
            continue
        source = origin_socket(port.outputs[0])
        junction = port.outputs[0].links[0].to_node
        processed = [origin_socket(link.from_socket).node for link in peer_links(junction)
                     if source.node in ancestors[origin_socket(link.from_socket).node]]
        if not processed:
            continue
        left = positions[source.node][0] + sizes[source.node].width
        right, level = positions[port][0] + 8, positions[port][1] - 8
        if right - left < ROUTE_GAP * 2:
            continue
        branch_nodes = {node for end in processed for node in ancestors[end] | {end}
                        if source.node in ancestors[node] and any(socket.type == "GEOMETRY" for socket in node.outputs)
                        and node.bl_idname != "NodeReroute"}
        first = min(positions[node][0] for node in branch_nodes)
        run = max(ROUTE_RUN, abs(level - positions[source.node][1] + sizes[source.node].sockets[source]) * 0.8)
        x = min(right - ROUTE_GAP, max(left + min(run, (right - left) / 2), first - ROUTE_GAP))
        entry = node_group.nodes.new("NodeReroute")
        entry["arrangement_bypass"] = True
        node_group.links.new(source, entry.inputs[0])
        _replace_link(node_group, port.inputs[0].links[0], entry.outputs[0])
        positions[entry] = (x - 8, level + 8)
        sizes[entry] = estimate_node_dimensions(entry)


def _route_references(node_group, sizes, positions, regions, paths, trunks, protected):
    """Share forward trunks and branch next to each consumer for every socket type."""
    references = defaultdict(list)
    path_of = {node: index for index, path in enumerate(paths) for node in path}
    peer_edges = {link for node in node_group.nodes if len(links := peer_links(node)) > 1 for link in links}
    for link in list(node_group.links):
        if (link.from_socket, link.to_socket) in protected or link.from_node.get("arrangement_branch") or link.to_node.get("arrangement_branch"):
            continue
        if link.from_node.get("arrangement_bypass") or link.to_node.get("arrangement_bypass"):
            continue
        if (link.from_node in path_of and link.to_node in path_of
                and path_of[link.from_node] == path_of[link.to_node] and len(link.from_socket.links) == 1):
            continue
        if link in peer_edges:
            continue
        if positions[link.to_node][0] - positions[link.from_node][0] - sizes[link.from_node].width > ROUTE_RUN * 4:
            references[link.from_socket].append(link)
    boxes = {node: _node_box(node, positions, sizes) for node in positions}
    tracks = []

    def route(source, x, y):
        node = node_group.nodes.new("NodeReroute")
        node.location = (x - 8, y + 8)
        node_group.links.new(source, node.inputs[0])
        return node.outputs[0]

    for source, links in references.items():
        if len(links) < 2 and source.type != "GEOMETRY":
            continue
        left = positions[source.node][0] + sizes[source.node].width
        right = max(positions[link.to_node][0] for link in links)
        start_y = positions[source.node][1] - sizes[source.node].sockets[source]
        destinations = {link.to_node for link in links}
        obstacles = [(bottom - ROUTE_GAP, top + ROUTE_GAP) for node, (a, b, top, bottom) in boxes.items()
                     if node not in destinations and a < right and b > left]
        for a, b, top, bottom in regions:
            if a >= right or b <= left:
                continue
            endpoints = destinations | {source.node}
            if any(a <= positions[node][0] <= b and bottom <= positions[node][1] <= top for node in endpoints):
                continue
            obstacles.append((bottom - ROUTE_GAP, top + ROUTE_GAP))
        obstacles += [(y - ROUTE_GAP, y + ROUTE_GAP) for a, b, y in tracks if a < right and b > left]
        if len(links) == 1:
            target = links[0]
            target_y = positions[target.to_node][1] - sizes[target.to_node].sockets[target.to_socket]
            if abs(target_y - start_y) < ROUTE_GAP and all(not lower < start_y < upper for lower, upper in obstacles):
                continue
        first_target = min(positions[link.to_node][0] for link in links)
        reach = (first_target - left) / 2.4
        lower, upper = start_y - reach, start_y + reach
        for link in links:
            y = positions[link.to_node][1] - sizes[link.to_node].sockets[link.to_socket]
            reach = (positions[link.to_node][0] - left) / 2.4
            lower, upper = max(lower, y - reach), min(upper, y + reach)
        candidates = [start_y, lower, upper, *(edge for interval in obstacles for edge in interval)]
        available = [y for y in candidates if lower <= y <= upper and all(not a < y < b for a, b in obstacles)]
        if source in trunks:
            level = trunks[source]
        elif not available:
            continue
        else:
            level = min(available, key=lambda y: (abs(y - start_y) + sum(
                abs(y - positions[link.to_node][1] + sizes[link.to_node].sockets[link.to_socket])
                for link in links) / len(links), -y))
        entry_run = min(max(ROUTE_RUN, abs(level - start_y) * 0.8), (first_target - left) / 3)
        trunk = route(source, left + entry_run, level)
        tracks.append((left + entry_run, right, level))
        previous_x = left + entry_run
        branches = []
        for link in links:
            target_x = positions[link.to_node][0]
            target_y = positions[link.to_node][1] - sizes[link.to_node].sockets[link.to_socket]
            exit_run = min(max(ROUTE_RUN, abs(target_y - level) * 0.8), (target_x - left) / 3)
            branches.append((target_x - exit_run, link))
        for x, link in sorted(branches, key=lambda item: item[0]):
            if x != previous_x:
                trunk = route(trunk, x, level)
                previous_x = x
            _replace_link(node_group, link, trunk)


def _zone_width(entry, interior, layout, sizes):
    """Reserve the body and the actual socket runs at each end of a paired scope."""
    body_height = max(sizes[node].height for node in interior)
    spans = []
    for top in (layout.height + BRANCH_GAP, -body_height - BRANCH_GAP):
        left, right = FIELD_GAP, layout.width + FIELD_GAP
        for node, (x, y) in layout.positions.items():
            for socket in node.inputs:
                for link in socket.links:
                    if link.from_node == entry:
                        delta = top + y - sizes[node].sockets[socket] + sizes[entry].sockets[link.from_socket]
                        left = max(left, max(GEOMETRY_GAP, abs(delta) * 0.8) - x)
            for socket in node.outputs:
                for link in socket.links:
                    if link.to_node in interior:
                        delta = top + y - sizes[node].sockets[socket] + sizes[link.to_node].sockets[link.to_socket]
                        right = max(right, x + sizes[node].width + max(GEOMETRY_GAP, abs(delta) * 0.8))
        spans.append(left + right)
    return max(spans)


def arrange_nodes(node_group):
    """Fix geometry lanes, pack shallow field regions, and route shared references."""
    _remove_routes(node_group)
    parents, children, order, ancestors = graph(node_group)
    _localize_peer_inputs(node_group, ancestors, children)
    parents, children, order, ancestors = graph(node_group)
    geometry_parents, paths = geometry_paths(node_group)
    core = set(geometry_parents)
    _localize_module_outputs(node_group, core, children)
    _localize_inputs(node_group)
    _presentation(node_group)
    parents, children, order, ancestors = graph(node_group)
    sizes = {node: estimate_node_dimensions(node) for node in order}
    for node, size in sizes.items():
        node.width = size.width
    peers = [_peer_inputs(node, ancestors) for node in order]
    regions = _field_regions(node_group.nodes, core, children, order)
    for region, _ in regions:
        outlets = [link.from_node for node in order if node in core for socket in node.inputs
                   for link in socket_links(socket) if link.from_node in region]
        peers.append(list(dict.fromkeys(outlets)))
    spacing, zones, field_layouts = {}, [], []
    scoped = set()
    for entry, end, members in reversed(node_zones(node_group)):
        if entry not in core or end not in core:
            continue
        interior = (members & core) - {entry}
        fields = [node for node in order if node in members - core - scoped]
        if fields:
            scoped.update(fields)
            layout = field_layout(fields, parents, peers, sizes, gap=FIELD_GAP, target_height=FIELD_HEIGHT)
            field_layouts.append(layout)
            local, width, height = layout.positions, layout.width, layout.height
            span = _zone_width(entry, interior, layout, sizes)
            for node in core:
                if entry in geometry_parents[node]:
                    spacing[node] = max(spacing.get(node, 0), span)
            zones.append((entry, end, interior, fields, local, width, height))
    column_parents = {node: list(sources) for node, sources in geometry_parents.items()}
    for node in core:
        if node.get("arrangement_input"):
            source = origin_socket(node.outputs[0]).node
            if source in core:
                column_parents[node] = [source]
    geometry_x, _ = _columns([node for node in order if node in core], column_parents, peers, sizes, GEOMETRY_GAP, spacing)
    positions = _place_geometry_paths(paths, geometry_x, sizes, order)
    occupied = _geometry_corridors(paths, positions, sizes)
    trunks = {}
    space = BranchSpace()
    layouts = []
    for region, targets in regions:
        remaining = [node for node in region if node not in scoped]
        if remaining:
            layout = field_layout(remaining, parents, peers, sizes, gap=FIELD_GAP, target_height=FIELD_HEIGHT)
            field_layouts.append(layout)
            layouts.append((remaining, targets, layout.positions, layout.width, layout.height))
    reserve = max([FIELD_HEIGHT, *(height for _, _, _, _, height in layouts),
                   *(height for _, _, _, _, _, _, height in zones)]) * 2 + BRANCH_GAP * 2
    limits = (max(y for _, y in positions.values()) + reserve,
              min(y - sizes[node].height for node, (_, y) in positions.items()) - reserve)
    for entry, end, interior, fields, local, width, height in zones:
        destinations = [min((node for node in interior if node.bl_idname != "NodeReroute"),
                            key=lambda node: (positions[node][0], node.name))]
        horizontal = (positions[entry][0] + sizes[entry].width + FIELD_GAP, positions[end][0] - FIELD_GAP)
        _place_field_region(local, width, height, destinations, destinations, positions, sizes, occupied, limits, core, trunks, space, horizontal)
        members = interior | {entry} | set(fields)
        boxes = [_node_box(node, positions, sizes) for node in members]
        occupied.append((min(box[0] for box in boxes), max(box[1] for box in boxes),
                         max(box[2] for box in boxes), min(box[3] for box in boxes)))
    pending = list(layouts)
    shared_nodes = {socket.node for socket in shared_geometry_outputs(node_group.nodes)}
    def delivery_order(layout):
        ports = [link.to_socket for node in layout[0] for socket in node.outputs
                 for link in socket.links if link.to_node in core]
        return min(((positions[socket.node][0], list(socket.node.inputs).index(socket))
                    for socket in ports), default=(float("inf"), 0))

    while pending:
        ready = [layout for layout in pending if all(
            link.to_node in positions or link.to_node in layout[0]
            for node in layout[0] for socket in node.outputs for link in socket.links)]
        layout = min(ready, key=lambda item: (any(node in shared_nodes for node in item[0]),
                     delivery_order(item), order.index(item[0][0])))
        pending.remove(layout)
        region, targets, local, width, height = layout
        consumers = list(dict.fromkeys(link.to_node for node in region for socket in node.outputs
                                      for link in socket.links if link.to_node not in region and link.to_node in positions))
        destinations = consumers or [node for node in order if node in targets] or [paths[0][-1]]
        owners = [node for node in order if node in targets] or [paths[0][-1]]
        _place_field_region(local, width, height, destinations, owners, positions, sizes, occupied, limits, core, trunks, space)
    protected = _route_field_links(node_group, field_layouts, positions)
    protected.update(_route_branch_exits(node_group, space.routes))
    _route_peer_bypasses(node_group, sizes, positions, ancestors)
    for node, position in positions.items():
        node.location = position
    _route_references(node_group, sizes, positions, occupied, paths, trunks, protected)
    while unused := [node for node in node_group.nodes
                     if node.bl_idname == "NodeReroute" and not node.outputs[0].is_linked]:
        for node in unused:
            node_group.nodes.remove(node)
    _presentation(node_group)