vscode-cowriting-plugin/src/mermaidFlowchartDiff.ts

/**
 * mermaidFlowchartDiff — F7.1 (#22). Pure flowchart parser + diff/emit. Parses a
 * `graph`/`flowchart` source into nodes (by id) and edges (by declaration order),
 * diffs baseline vs current, and re-emits the CURRENT source augmented with
 * `classDef`/`class`/`linkStyle` directives coloring added/changed elements and
 * ghosting removed ones (INV-29/31). Deterministic; no vscode, no DOM.
 */
import { CW_COLORS } from "./mermaidDiff";

export interface FlowNode {
  id: string;
  label?: string;
  open?: string;
  close?: string;
  /** Verbatim declaration token, e.g. `A[Start]`, for ghost re-injection. */
  decl?: string;
}
export interface FlowEdge {
  from: string;
  to: string;
  label?: string;
  index: number;
  /** Verbatim edge statement for ghost re-injection. */
  raw: string;
}
export interface FlowGraph {
  header: string;
  nodes: Map<string, FlowNode>;
  edges: FlowEdge[];
}

const NODE_TOKEN = /^([A-Za-z0-9_]+)(\[\[[^\]]*\]\]|\[[^\]]*\]|\(\([^)]*\)\)|\([^)]*\)|\{[^}]*\}|>[^\]]*\])?/;
// link operators: -->, ---, -.->, -.-, ==>, ===, --x, --o, ==x, ==o, optionally |label|
const LINK = /^(-->|---|-\.->|-\.-|==>|===|--[xo]|==[xo])(\|([^|]*)\|)?/;

function shapeOf(bracket: string): { label: string; open: string; close: string } {
  if (bracket.startsWith("[[")) return { label: bracket.slice(2, -2), open: "[[", close: "]]" };
  if (bracket.startsWith("((")) return { label: bracket.slice(2, -2), open: "((", close: "))" };
  const open = bracket[0];
  const close = bracket[bracket.length - 1];
  return { label: bracket.slice(1, -1), open, close };
}

function recordNode(nodes: Map<string, FlowNode>, token: string): string {
  const m = token.match(NODE_TOKEN)!;
  const id = m[1];
  const existing = nodes.get(id) ?? { id };
  if (m[2]) {
    const s = shapeOf(m[2]);
    existing.label = s.label;
    existing.open = s.open;
    existing.close = s.close;
    existing.decl = `${id}${m[2]}`;
  }
  nodes.set(id, existing);
  return id;
}

export function parseFlowchart(source: string): FlowGraph {
  const lines = source.split(/\r?\n/);
  let header = "";
  const nodes = new Map<string, FlowNode>();
  const edges: FlowEdge[] = [];
  for (const rawLine of lines) {
    const line = rawLine.trim();
    if (line === "" || line.startsWith("%%")) continue;
    if (!header && /^(flowchart|graph)\b/.test(line)) {
      header = line;
      continue;
    }
    let rest = line;
    const firstM = rest.match(NODE_TOKEN);
    if (!firstM || firstM[0] === "") continue;
    let leftId = recordNode(nodes, firstM[0]);
    rest = rest.slice(firstM[0].length).trimStart();
    // Walk a (possibly chained) edge statement: A[..] --> B -->|x| C
    while (rest.length > 0) {
      const linkM = rest.match(LINK);
      if (!linkM) break;
      rest = rest.slice(linkM[0].length).trimStart();
      const rhsM = rest.match(NODE_TOKEN);
      if (!rhsM || rhsM[0] === "") break;
      const rightId = recordNode(nodes, rhsM[0]);
      edges.push({
        from: leftId,
        to: rightId,
        label: linkM[3] || undefined,
        index: edges.length,
        raw: `${leftId} ${linkM[1]}${linkM[2] ?? ""} ${rightId}`,
      });
      rest = rest.slice(rhsM[0].length).trimStart();
      leftId = rightId;
    }
  }
  return { header: header || "flowchart TD", nodes, edges };
}

function nodeChanged(a: FlowNode, b: FlowNode): boolean {
  return (a.label ?? "") !== (b.label ?? "") || (a.open ?? "") !== (b.open ?? "");
}
const edgeKey = (e: FlowEdge): string => `${e.from} ${e.to}`;

export function diffFlowchart(beforeSrc: string, currentSrc: string): string {
  const before = parseFlowchart(beforeSrc);
  const current = parseFlowchart(currentSrc);

  const addedNodes: string[] = [];
  const changedNodes: string[] = [];
  for (const [id, cur] of current.nodes) {
    const prev = before.nodes.get(id);
    if (!prev) addedNodes.push(id);
    else if (nodeChanged(prev, cur)) changedNodes.push(id);
  }
  const removedNodes: FlowNode[] = [];
  for (const [id, prev] of before.nodes) {
    if (!current.nodes.has(id)) removedNodes.push(prev);
  }

  const beforeEdgeKeys = new Set(before.edges.map(edgeKey));
  const currentEdgeKeys = new Set(current.edges.map(edgeKey));
  const addedEdgeIdx = current.edges.filter((e) => !beforeEdgeKeys.has(edgeKey(e))).map((e) => e.index);
  const removedEdges = before.edges.filter((e) => !currentEdgeKeys.has(edgeKey(e)));

  // Ghost-edge indices follow the current edge count, in deterministic order.
  let nextIdx = current.edges.length;
  const ghostEdgeLines: string[] = [];
  const ghostEdgeIdx: number[] = [];
  for (const e of removedEdges) {
    ghostEdgeLines.push(`  ${e.from} -.-> ${e.to}`);
    ghostEdgeIdx.push(nextIdx++);
  }

  const out: string[] = [currentSrc.replace(/\s+$/, "")];

  // Ghost removed nodes: re-inject their declaration (or bare id) so they appear.
  for (const n of removedNodes) out.push(`  ${n.decl ?? n.id}`);
  out.push(...ghostEdgeLines);

  // classDefs (always emit the three; harmless if a class is unused).
  out.push(
    `  classDef cwAdded fill:${CW_COLORS.added}22,stroke:${CW_COLORS.added},stroke-width:2px;`,
    `  classDef cwChanged fill:${CW_COLORS.changed}22,stroke:${CW_COLORS.changed},stroke-width:2px;`,
    `  classDef cwRemoved fill:${CW_COLORS.removed}11,stroke:${CW_COLORS.removed},stroke-width:1px,stroke-dasharray:5 3,color:${CW_COLORS.removed};`,
  );
  if (addedNodes.length) out.push(`  class ${addedNodes.join(",")} cwAdded;`);
  if (changedNodes.length) out.push(`  class ${changedNodes.join(",")} cwChanged;`);
  if (removedNodes.length) out.push(`  class ${removedNodes.map((n) => n.id).join(",")} cwRemoved;`);

  for (const i of addedEdgeIdx) out.push(`  linkStyle ${i} stroke:${CW_COLORS.added},stroke-width:2px;`);
  for (const i of ghostEdgeIdx)
    out.push(`  linkStyle ${i} stroke:${CW_COLORS.removed},stroke-width:1px,stroke-dasharray:5 3;`);

  return out.join("\n");
}