js/layout.js

// layout.js — nested, layered (top-down) layout derived entirely from the IR.
//
// Hierarchy is rendered as containment boxes; coarse dataflow is rendered as
// arrows overlaid on top. Repeats are collapsible: a collapsed repeat is a
// single block, an expanded repeat is a group box holding its body's layout.
//
// The layout is computed bottom-up (measure child sizes, then place them in
// layers), producing absolute rectangles for every visible node plus a
// deduplicated set of edges mapped to their nearest visible representative.

const PAD = 18; // inner padding of a container box
const HGAP = 34; // horizontal gap between siblings in a layer
const VGAP = 46; // vertical gap between layers
const REPEAT_HEADER_H = 34;
const CONTAINER_HEADER_H = 26;
const LEAF_H = 54;
const REPEAT_LEAF_H = 66; // collapsed repeat shows a count subtitle + stacked look
const SCHEDULE_STRIP_H = 16; // per-layer attention-schedule strip on repeat nodes
const MIN_W = 132;

// Edge kinds that imply a top-down ordering (everything except residual,
// which loops back inside a block and would create cycles).
const LAYERING_KINDS = new Set(["data", "cross_attention", "position", "mask"]);

class Layout {
  constructor(ir, opts) {
    this.ir = ir;
    this.expanded = opts.expanded; // Set<repeatId>
    this.fields = opts.fields; // resolved config fields
    this.showInfo = opts.showInfo !== false; // shape/attribute captions on nodes
    this.tied = opts.tied !== false; // word embeddings tied → hide the separate LM head

    this.rectById = new Map(); // id -> {x,y,w,h,...} absolute
    this.placed = []; // draw order (containers before children)
    this.width = 0;
    this.height = 0;

    this._run();
  }

  isRepeatOpen(id) {
    return this.expanded.has(id);
  }

  isOpen(id) {
    const n = this.ir.node(id);
    if (!n) return false;
    if (n.nodeType === "repeat") return this.isRepeatOpen(id);
    return this.ir.layoutChildrenOf(id).length > 0;
  }

  // Direct visible layout-children of a container (root also gets input nodes).
  childrenOf(id) {
    let kids = this.ir.layoutChildrenOf(id).slice();
    if (id === this.ir.rootId) {
      kids = kids.concat([...this.ir.pseudoIds]);
      // The synthetic LM head is only a distinct node when embeddings are untied.
      if (this.tied && this.ir.lmHeadId) kids = kids.filter((k) => k !== this.ir.lmHeadId);
    }
    return kids;
  }

  // The direct child of `container` whose subtree contains `x` (or x itself).
  childContaining(x, container, childSet) {
    if (this.ir.isPseudo(x)) {
      return childSet.has(x) ? x : null;
    }
    let cur = x;
    let guard = 0;
    while (cur && guard++ < 64) {
      if (childSet.has(cur)) return cur;
      const n = this.ir.node(cur);
      if (!n || cur === container) return null;
      cur = n.parent;
    }
    return null;
  }

  // Nearest drawn representative of any id under the current expand state.
  representative(id) {
    if (this.ir.isPseudo(id)) return id;
    let cur = id;
    let highestCollapsed = null;
    let guard = 0;
    while (cur && guard++ < 64) {
      const n = this.ir.node(cur);
      if (!n) break;
      if (n.nodeType === "repeat" && !this.isRepeatOpen(cur)) highestCollapsed = cur;
      cur = n.parent;
    }
    let base = highestCollapsed || id;
    if (this.ir.bodyToRepeat.has(base)) base = this.ir.bodyToRepeat.get(base);
    return base;
  }

  // --- Measurement (bottom-up) ---------------------------------------------

  measure(id) {
    if (!this.isOpen(id)) return this._leaf(id);

    const childIds = this.childrenOf(id);
    const childBoxes = childIds.map((cid) => this.measure(cid));
    const boxByChild = new Map(childIds.map((cid, i) => [cid, childBoxes[i]]));

    const rows = this._layerize(id, childIds);
    this._placeRows(rows, boxByChild);

    // Include each child's deck extent so the block's stacked depth doesn't eat
    // into padding / overlap the next node.
    let contentW = Math.max(...childBoxes.map((b) => b.rx + b.w + (b.deckW || 0)), MIN_W);
    const contentH = Math.max(...childBoxes.map((b) => b.ry + b.h + (b.deckH || 0)), LEAF_H);

    const n = this.ir.node(id);
    const isRoot = id === this.ir.rootId;
    const isRepeat = n && n.nodeType === "repeat";
    // Expanded repeat with a schedule reserves header room for the strip.
    const hasSchedule = isRepeat && !!this.ir.scheduleForRepeat(n, this.fields);
    const headerH =
      (isRoot ? 0 : isRepeat ? REPEAT_HEADER_H : CONTAINER_HEADER_H) +
      (hasSchedule ? SCHEDULE_STRIP_H : 0);
    const drawFrame = !isRoot;

    // Ensure the box is wide enough for its header label + collapse toggle, so
    // e.g. "Llama Decoder Layer × 32" is never clipped by the toggle button.
    let extraX = 0;
    if (drawFrame) {
      const headerLabel = isRepeat
        ? this.ir.label(id, this.fields)
        : `${this.ir.label(id, this.fields)}  ·  ${this.ir.kindLabel(n && n.kind)}`;
      const headerNeed = Math.round(headerLabel.length * 6.9) + 24 + (isRepeat ? 34 : 8);
      const innerNeed = headerNeed - PAD * 2;
      if (innerNeed > contentW) {
        extraX = (innerNeed - contentW) / 2; // keep children centred under the header
        contentW = innerNeed;
      }
    }

    // Offset children inside padding + header (plus any header-driven widening).
    const ox = drawFrame ? PAD + extraX : 0;
    const oy = drawFrame ? PAD + headerH : 0;
    childBoxes.forEach((b) => {
      b.rx += ox;
      b.ry += oy;
    });

    const deck = this._deckExtent(n);
    return {
      id,
      node: n,
      kind: "container",
      isRepeat,
      isRoot,
      drawFrame,
      headerH,
      hasSchedule,
      deckW: deck.dw,
      deckH: deck.dh,
      children: childBoxes,
      rx: 0,
      ry: 0,
      w: drawFrame ? contentW + PAD * 2 : contentW,
      h: drawFrame ? contentH + PAD * 2 + headerH : contentH,
    };
  }

  _leaf(id) {
    const n = this.ir.node(id);
    const isRepeat = n && n.nodeType === "repeat";
    const label = this.ir.label(id, this.fields);
    const isPseudo = !n;
    const info = this.showInfo && !isPseudo ? this.ir.nodeInfo(id, this.fields) : null;
    let w;
    let h;
    if (isPseudo) {
      // Small centred pill (11px text, no icon).
      w = Math.max(88, Math.min(240, Math.round(label.length * 6.2 + 34)));
      h = 40;
    } else {
      // Box must fit its title: bold 13px starting at the x+28 icon inset, plus
      // right padding (extra for the collapse toggle on repeats). The caption
      // line (11px mono) may need more. Cap high enough for long class names.
      const LEFT = 28;
      const rightPad = isRepeat ? 42 : 20;
      let need = LEFT + Math.round(label.length * 8.0) + rightPad;
      if (info) need = Math.max(need, LEFT + Math.round(info.length * 6.6) + 18);
      // Reserve room for the kernel org avatar(s) / bolt badge (top-right).
      if (n && n.attributes && n.attributes.kernel) need += 52;
      // Reserve room for the "🔗 tied" badge on the word embedding.
      if (this.tied && this.ir.isWordEmbedding(n)) need += 62;
      w = Math.max(MIN_W, Math.min(400, need));
      h = (isRepeat ? REPEAT_LEAF_H : LEAF_H) + (info ? 16 : 0);
    }
    // Per-layer attention schedule strip (collapsed repeat block).
    const hasSchedule = isRepeat && !!this.ir.scheduleForRepeat(n, this.fields);
    if (hasSchedule) h += SCHEDULE_STRIP_H;
    const deck = this._deckExtent(n);
    return {
      id,
      node: n,
      kind: isRepeat ? "repeat" : n ? "leaf" : "input",
      isRepeat,
      info,
      hasSchedule,
      deckW: deck.dw,
      deckH: deck.dh,
      children: [],
      rx: 0,
      ry: 0,
      w,
      h,
    };
  }

  // The extent a repeat's stacked deck adds to the bottom-right — from the same
  // geometry the renderer uses, so reserved space matches what's drawn.
  _deckExtent(node) {
    const geo = this.ir.deckGeometry(node, this.fields);
    return geo ? { dw: geo.dw, dh: geo.dh } : { dw: 0, dh: 0 };
  }

  // Assign each child to a layer via longest-path over layering edges.
  _layerize(containerId, childIds) {
    const childSet = new Set(childIds);
    const adj = []; // [a, b] directed
    for (const e of this.ir.edges) {
      if (!LAYERING_KINDS.has(e.kind)) continue;
      const a = this.childContaining(e.source, containerId, childSet);
      const b = this.childContaining(e.target, containerId, childSet);
      if (a && b && a !== b) adj.push([a, b]);
    }
    const layer = new Map(childIds.map((c) => [c, 0]));
    // Relaxation bounded by node count handles accidental cycles safely.
    for (let it = 0; it < childIds.length; it++) {
      let changed = false;
      for (const [a, b] of adj) {
        const cand = layer.get(a) + 1;
        if (cand > layer.get(b)) {
          layer.set(b, cand);
          changed = true;
        }
      }
      if (!changed) break;
    }
    const rows = new Map();
    childIds.forEach((c) => {
      const l = layer.get(c);
      if (!rows.has(l)) rows.set(l, []);
      rows.get(l).push(c);
    });
    return [...rows.keys()].sort((a, b) => a - b).map((k) => rows.get(k));
  }

  _placeRows(rows, boxByChild) {
    const fw = (b) => b.w + (b.deckW || 0); // footprint incl. deck depth
    const fh = (b) => b.h + (b.deckH || 0);
    const rowWidths = rows.map((row) =>
      row.reduce((s, c) => s + fw(boxByChild.get(c)), 0) + HGAP * Math.max(0, row.length - 1)
    );
    const maxW = Math.max(...rowWidths, MIN_W);
    let y = 0;
    rows.forEach((row, ri) => {
      const rowH = Math.max(...row.map((c) => fh(boxByChild.get(c))));
      let x = (maxW - rowWidths[ri]) / 2;
      row.forEach((c) => {
        const b = boxByChild.get(c);
        b.rx = x;
        b.ry = y + (rowH - fh(b)) / 2;
        x += fw(b) + HGAP;
      });
      y += rowH + VGAP;
    });
  }

  // --- Flatten to absolute coordinates -------------------------------------

  _flatten(box, px, py, parentId) {
    const x = px + box.rx;
    const y = py + box.ry;
    const rect = {
      id: box.id,
      node: box.node,
      kind: box.kind,
      isRepeat: box.isRepeat,
      isRoot: box.isRoot,
      drawFrame: box.drawFrame,
      headerH: box.headerH || 0,
      info: box.info || null,
      hasSchedule: !!box.hasSchedule,
      deckW: box.deckW || 0,
      deckH: box.deckH || 0,
      parentId: parentId || null,
      x,
      y,
      w: box.w,
      h: box.h,
    };
    this.rectById.set(box.id, rect);
    this.placed.push(rect); // containers pushed before their children
    box.children.forEach((c) => this._flatten(c, x, y, box.id));
  }

  _run() {
    if (!this.ir.rootId) return;
    const tree = this.measure(this.ir.rootId);
    this._flatten(tree, 0, 0, null);
    this.width = tree.w;
    this.height = tree.h;
  }

  // --- Visible edges --------------------------------------------------------

  // Map every IR edge to its visible representatives, drop internal/self
  // edges, and dedupe. `kindFilter` is a Set of enabled edge kinds.
  visibleEdges(kindFilter) {
    const seen = new Set();
    const out = [];
    for (const e of this.ir.edges) {
      if (kindFilter && !kindFilter.has(e.kind)) continue;
      const s = this.representative(e.source);
      const t = this.representative(e.target);
      if (s === t) continue;
      if (!this.rectById.has(s) || !this.rectById.has(t)) continue;
      // Skip only block-level *residuals* that point to an enclosed node (the
      // arrow from a big group box across to a descendant is unreadable). Keep
      // containment DATA edges — those are the fan-out into a block's children
      // (e.g. self_attn → q/k/v, mlp → gate/up), which show the real flow.
      if (e.kind === "residual" && (this._encloses(s, t) || this._encloses(t, s))) continue;
      const key = `${s}${t}${e.kind}`;
      if (seen.has(key)) continue;
      seen.add(key);
      out.push({ source: s, target: t, kind: e.kind, raw: e });
    }
    return out;
  }

  // Is `ancestor` a drawn-tree ancestor of `descendant`?
  _encloses(ancestor, descendant) {
    let cur = this.rectById.get(descendant);
    let guard = 0;
    while (cur && cur.parentId && guard++ < 64) {
      if (cur.parentId === ancestor) return true;
      cur = this.rectById.get(cur.parentId);
    }
    return false;
  }
}

export { Layout };