/**
 * Fold the single-use ANF temps numbl's inline pass left behind.
 *
 * numbl's own pass (`inlinePass`) only folds producers its C backend can fuse,
 * which leaves out tensor-producing calls — `sin(x)`, `exp(x)`, `tanh(x)`. The
 * WGSL emitter fuses all of those happily, so without this pass a source line
 * like
 *
 *   s = (amp * env) .* sin(2*pi*f*(t - t0)) .* exp(-(gx + gy));
 *
 * plans as half a dozen kernels rather than one, each writing a whole grid to
 * memory for the next one to read straight back.
 *
 * Same shape as numbl's pass, deliberately narrower where it matters: only
 * compiler temps (`_mtoc2_*`) are folded, so every variable the .m names keeps
 * its own buffer and one source line stays one kernel. The producer's RHS must
 * be something the emitter can evaluate per element, its result must be used
 * exactly once, and nothing between the two statements may write to anything
 * it reads.
 *
 * Adapted from math-webgpu-sandbox's src/mgpu/fuse.ts, trimmed to the
 * straight-line bodies this project compiles.
 */
import type { IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
import { isGpuFusableExpr } from './wgsl.ts';

const isTemp = (cName: string): boolean => cName.startsWith('_mtoc2_');

export function fuseTemps(stmts: IRStmt[]): IRStmt[] {
  let cur = stmts;
  for (let iter = 0; iter < 32; iter++) {
    const next = fuseOnePass(cur);
    if (next === cur) break;
    cur = next;
  }
  return cur;
}

/** One sweep, at most one fold. Returns the same array when nothing fired. */
function fuseOnePass(stmts: IRStmt[]): IRStmt[] {
  const uses = useCounts(stmts);
  for (let i = 0; i < stmts.length; i++) {
    const p = stmts[i];
    if (p.kind !== 'Assign' || !isTemp(p.cName)) continue;
    if (uses.get(p.cName) !== 1) continue;
    if (!isGpuFusableExpr(p.expr)) continue;

    const reads = new Set<string>();
    walkVars(p.expr, (c) => reads.add(c));

    for (let j = i + 1; j < stmts.length; j++) {
      const c = stmts[j];
      if (c.kind !== 'Assign') break; // anything else ends the safe window
      if (countIn(c.expr, p.cName) > 0) {
        // The one use. Fold into it if it is a kernel; if it is a stencil
        // call, or anything else that wants its argument in a buffer of its
        // own, leave the producer alone.
        if (countIn(c.expr, p.cName) === 1 && isGpuFusableExpr(c.expr)) {
          c.expr = substitute(c.expr, p.cName, p.expr);
          const out = stmts.slice();
          out.splice(i, 1);
          return out;
        }
        break;
      }
      // An intervening write to the temp itself, or to one of its operands,
      // invalidates the fold window. Statements that do neither are simply
      // skipped over — the stencil dispatch between the source term and the
      // update it feeds is exactly that case.
      if (c.cName === p.cName || reads.has(c.cName)) break;
    }
  }
  return stmts;
}

function useCounts(stmts: IRStmt[]): Map<string, number> {
  const counts = new Map<string, number>();
  const bump = (c: string): void => {
    counts.set(c, (counts.get(c) ?? 0) + 1);
  };
  for (const s of stmts) if (s.kind === 'Assign') walkVars(s.expr, bump);
  return counts;
}

function walkVars(e: IRExpr, visit: (cName: string) => void): void {
  const walk = (x: IRExpr): void => {
    switch (x.kind) {
      case 'Var':
        visit(x.cName);
        return;
      case 'Binary':
        walk(x.left);
        walk(x.right);
        return;
      case 'Unary':
        walk(x.operand);
        return;
      case 'Call':
        x.args.forEach(walk);
        return;
      default:
        return;
    }
  };
  walk(e);
}

function countIn(e: IRExpr, cName: string): number {
  let n = 0;
  walkVars(e, (c) => {
    if (c === cName) n++;
  });
  return n;
}

/** Replace the (single) `Var` read of `cName` with `replacement`. */
function substitute(e: IRExpr, cName: string, replacement: IRExpr): IRExpr {
  const sub = (x: IRExpr): IRExpr => {
    if (x.kind === 'Var' && x.cName === cName) return replacement;
    switch (x.kind) {
      case 'Binary':
        x.left = sub(x.left);
        x.right = sub(x.right);
        return x;
      case 'Unary':
        x.operand = sub(x.operand);
        return x;
      case 'Call':
        for (let i = 0; i < x.args.length; i++) x.args[i] = sub(x.args[i]);
        return x;
      default:
        return x;
    }
  };
  return sub(e);
}