unreal-engine-mcp-system-plugin/ue_side/graph_layout.py

"""Small deterministic graph layout helpers shared by MCP graph tools.

The goal is not to replace Unreal's graph layout. It is a guard rail for
generated graphs: keep columns readable and prevent heavy nodes from being
spawned on top of each other.
"""
from __future__ import annotations


DEFAULT_NODE_W = 260
DEFAULT_NODE_H = 150
DEFAULT_GAP_X = 180
DEFAULT_GAP_Y = 80


def estimate_size(spec: dict, domain: str = "blueprint") -> tuple[int, int]:
    if domain == "pcg":
        cls = str(spec.get("class") or spec.get("type") or "").lower()
        # I/O nodes and reroutes are slim; spawners/samplers carry big detail panels.
        if cls in ("input", "output"):
            return 220, 120
        if "spawner" in cls or "sampler" in cls or "subgraph" in cls:
            return 340, 240
        if "filter" in cls or "transform" in cls or "createpoints" in cls:
            return 320, 200
        return 300, 170

    if domain == "material":
        cls = str(spec.get("class") or spec.get("target") or "").lower()
        if "texturesample" in cls:
            return 310, 360
        if "vectorparameter" in cls or "constant3vector" in cls or "constant4vector" in cls:
            return 300, 300
        if "scalarparameter" in cls:
            return 260, 150
        if "panner" in cls:
            return 280, 190
        if "linearinterpolate" in cls or "multiply" in cls or "add" in cls:
            return 260, 170
        if "texturecoordinate" in cls:
            return 260, 130
        return 280, 180

    kind = str(spec.get("kind") or "").lower()
    params = spec.get("params") or {}
    if kind == "comment":
        return int(params.get("width", 400)), int(params.get("height", 220))
    if kind in ("event", "custom_event", "branch", "sequence", "operator", "knot"):
        return 240, 120
    if kind in ("call_function", "spawn_actor", "format_text"):
        return 340, 190
    if kind in ("make_struct", "break_struct", "set_fields_in_struct"):
        return 360, 260
    if kind in ("variable_get", "variable_set", "self"):
        return 260, 130
    return DEFAULT_NODE_W, DEFAULT_NODE_H


def normalize_specs(nodes: list[dict], edges: list[dict] | None = None, domain: str = "blueprint",
                    gap_x: int = DEFAULT_GAP_X, gap_y: int = DEFAULT_GAP_Y,
                    snap_x: int = 360) -> list[dict]:
    """Return copied nodes with collision-resistant positions.

    The algorithm preserves the user's left-to-right intent from positions, but
    snaps nearby x values into columns and stacks each column by estimated node
    height. This keeps generated graphs stable and readable.
    """
    copied = [dict(n) for n in (nodes or [])]
    if len(copied) < 2:
        return copied

    with_meta = []
    for index, node in enumerate(copied):
        pos = node.get("position") or [0, 0]
        x = float(pos[0]) if len(pos) > 0 else 0.0
        y = float(pos[1]) if len(pos) > 1 else 0.0
        w, h = estimate_size(node, domain)
        with_meta.append({"index": index, "node": node, "x": x, "y": y, "w": w, "h": h})

    sorted_by_x = sorted(with_meta, key=lambda item: (item["x"], item["y"], item["index"]))
    columns: list[list[dict]] = []
    for item in sorted_by_x:
        if not columns:
            columns.append([item])
            continue
        avg_x = sum(i["x"] for i in columns[-1]) / len(columns[-1])
        if abs(item["x"] - avg_x) <= max(120, snap_x * 0.45):
            columns[-1].append(item)
        else:
            columns.append([item])

    min_x = min(item["x"] for item in with_meta)
    for column_index, column in enumerate(columns):
        column_x = min_x + column_index * (max(item["w"] for item in column) + gap_x)
        column.sort(key=lambda item: (item["y"], item["index"]))
        cursor_y = min(item["y"] for item in column)
        for item in column:
            desired_y = item["y"]
            y = max(desired_y, cursor_y)
            item["node"]["position"] = [round(column_x), round(y)]
            cursor_y = y + item["h"] + gap_y

    return copied


def tidy_existing_nodes(nodes: list[dict], edges: list[dict] | None = None, domain: str = "material") -> list[dict]:
    """Return [{guid/name,x,y}] for existing graph dump nodes."""
    if edges:
        return _tidy_existing_by_edges(nodes, edges, domain)

    specs = []
    for node in nodes or []:
        cls = node.get("class") or ""
        x = node.get("x", (node.get("position") or [0, 0])[0])
        y = node.get("y", (node.get("position") or [0, 0])[1])
        specs.append({
            "id": node.get("guid") or node.get("name"),
            "class": cls,
            "position": [float(x), float(y)],
        })
    normalized = normalize_specs(specs, domain=domain)
    moves = []
    for spec in normalized:
        x, y = spec.get("position") or [0, 0]
        moves.append({"guid": spec.get("id"), "x": x, "y": y})
    return moves


def _tidy_existing_by_edges(nodes: list[dict], edges: list[dict], domain: str) -> list[dict]:
    by_id = {}
    for node in nodes or []:
        node_id = node.get("guid") or node.get("name")
        if node_id:
            by_id[node_id] = node
    if not by_id:
        return []

    incoming = {node_id: set() for node_id in by_id}
    outgoing = {node_id: set() for node_id in by_id}
    for edge in edges or []:
        try:
            src = edge["from"][0]
            dst = edge["to"][0]
        except Exception:
            continue
        if src in by_id and dst in by_id:
            outgoing[src].add(dst)
            incoming[dst].add(src)

    depths = {node_id: 0 for node_id in by_id}
    changed = True
    guard = 0
    while changed and guard < max(4, len(by_id) * 3):
        changed = False
        guard += 1
        for src, dsts in outgoing.items():
            for dst in dsts:
                candidate = depths[src] + 1
                if candidate > depths[dst]:
                    depths[dst] = candidate
                    changed = True

    columns: dict[int, list[dict]] = {}
    for node_id, node in by_id.items():
        columns.setdefault(depths.get(node_id, 0), []).append(node)

    all_x = []
    all_y = []
    for node in by_id.values():
        pos = node.get("position")
        all_x.append(float(node.get("x", pos[0] if pos else 0)))
        all_y.append(float(node.get("y", pos[1] if pos else 0)))
    min_x = min(all_x) if all_x else 0
    min_y = min(all_y) if all_y else 0

    moves = []
    sorted_depths = sorted(columns)
    col_x = min_x
    for depth in sorted_depths:
        column = columns[depth]
        column.sort(key=lambda node: float(node.get("y", (node.get("position") or [0, 0])[1])))
        max_w = 0
        cursor_y = min_y
        for node in column:
            spec = {"class": node.get("class"), "kind": node.get("kind")}
            w, h = estimate_size(spec, domain)
            max_w = max(max_w, w)
            node_id = node.get("guid") or node.get("name")
            moves.append({"guid": node_id, "x": round(col_x), "y": round(cursor_y)})
            cursor_y += h + DEFAULT_GAP_Y
        col_x += max_w + DEFAULT_GAP_X
    return moves