/**
 * Statement list -> a replayable sequence of GPU operations.
 *
 * Everything expensive happens once, here: pipeline compilation, buffer
 * allocation, bind-group construction. Because numbl fixes every type and
 * shape at lowering time, the resulting op sequence is fully static — so
 * `encodeSteps` is pure synchronous command recording, with no allocation, no
 * pipeline lookup and no readback. That is what lets a whole batch of
 * timesteps be encoded into one submit and keeps the CPU out of the loop.
 */
import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
import type { Assign, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
import type { CompiledFunction } from './compile.ts';
import { EXTERNAL_OPS } from './externals.ts';
import { OpPlan, type OpKind } from './ops.ts';
import { kernelOperandBudget } from '../device.ts';
import {
  buildKernel,
  checkShapes,
  UnsupportedOnGpu,
  WORKGROUP_SIZE,
  type KernelInputs,
} from './wgsl.ts';

const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);

interface Slot {
  buffer: GPUBuffer;
  count: number;
}

const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
  device.createBuffer({
    label,
    size: 4 * count,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
  });

/**
 * Buffers for host-bound variables, shared across plans.
 *
 * A model is two programs — `init` and `step` — compiled separately but
 * operating on the same state. `p` in the step must be the very buffer `init`
 * wrote, so the buffers for host bindings live here rather than inside either
 * plan.
 */
export class HostBuffers {
  #device: GPUDevice;
  #slots = new Map<string, Slot>();

  constructor(device: GPUDevice) {
    this.#device = device;
  }

  ensure(name: string, count: number): Slot {
    const existing = this.#slots.get(name);
    if (existing) {
      if (existing.count !== count) {
        throw new UnsupportedOnGpu(
          `'${name}' is ${existing.count} elements in one program and ` +
            `${count} in another`,
        );
      }
      return existing;
    }
    const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
    this.#slots.set(name, slot);
    return slot;
  }

  get(name: string): Slot | undefined {
    return this.#slots.get(name);
  }

  /** Upload initial data for a host binding. */
  upload(name: string, data: Float32Array): void {
    const slot = this.#slots.get(name);
    if (!slot) throw new Error(`upload: no buffer named '${name}'`);
    if (data.length !== slot.count) {
      throw new Error(
        `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
      );
    }
    this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
  }

  destroy(): void {
    for (const s of this.#slots.values()) s.buffer.destroy();
    this.#slots.clear();
  }
}

type Op =
  | {
      kind: 'kernel';
      pipeline: GPUComputePipeline;
      bindGroup: GPUBindGroup;
      count: number;
      label: string;
      /** Set when the kernel had to write to scratch because its output
       *  aliases one of its inputs; copied back after the dispatch. */
      copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
    }
  | {
      kind: 'external';
      pipeline: GPUComputePipeline;
      bindGroup: GPUBindGroup;
      workgroups: number;
      label: string;
    }
  | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };

export interface PlanSpec {
  /** The specialized function this plan executes. */
  fn: CompiledFunction;
  /** Output index -> host binding name to copy the result into after the run,
   *  so the next call reads it (the new field feeds the old). */
  feedback: (string | null)[];
}

/**
 * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
 * buffers after it, then the params buffer.
 *
 * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
 * only contains the bindings the shader actually references — so a kernel that
 * happens to use no parameters would drop the params binding and no longer
 * match the bind group. An explicit layout may carry bindings the shader
 * ignores.
 */
function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
  const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
    binding,
    visibility: GPUShaderStage.COMPUTE,
    buffer: { type: 'read-only-storage' },
  });
  return device.createBindGroupLayout({
    entries: [
      {
        binding: 0,
        visibility: GPUShaderStage.COMPUTE,
        buffer: { type: 'storage' },
      },
      ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
      readOnly(inputs + 1),
    ],
  });
}

/**
 * Compile one shader into a pipeline.
 *
 * No validation error scope around it: `createComputePipelineAsync` already
 * rejects on a shader that will not compile or a layout that does not match,
 * which is the whole of what a scope here would have caught, and the scope
 * costs an extra device round trip per pipeline. `getCompilationInfo`, which
 * has the line and column within the generated WGSL, is asked for only once
 * something has gone wrong, and defensively even then.
 *
 * In practice the WGSL here is generated, so a shader that fails to compile is
 * this project's bug rather than the user's; a mistake in a .m is caught
 * earlier, by the emitter, with a position in the MATLAB source.
 */
async function makePipeline(
  device: GPUDevice,
  code: string,
  label: string,
  bindGroupLayout: GPUBindGroupLayout,
): Promise<GPUComputePipeline> {
  const module = device.createShaderModule({ code, label });
  try {
    return await device.createComputePipelineAsync({
      layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
      compute: { module, entryPoint: 'main' },
      label,
    });
  } catch (e) {
    throw new UnsupportedOnGpu(
      `generated WGSL failed to compile for '${label}':\n` +
        `${await shaderErrors(module, e)}\n--- shader ---\n${code}`,
    );
  }
}

/** Per-line compile errors, if the browser will hand them over. */
async function shaderErrors(module: GPUShaderModule, cause: unknown): Promise<string> {
  const fallback = cause instanceof Error ? cause.message : String(cause);
  try {
    const info = await module.getCompilationInfo();
    const errors = info.messages.filter((m) => m.type === 'error');
    if (!errors.length) return fallback;
    return errors.map((e) => `  ${e.lineNum}:${e.linePos} ${e.message}`).join('\n');
  } catch {
    return fallback;
  }
}

/** A compiled .m function, ready to run on the GPU. */
export class ModelPlan {
  /** Scalar parameter names, in the order the params buffer expects them. */
  readonly paramNames: string[];

  #device: GPUDevice;
  #ops: Op[];
  #owned: GPUBuffer[];
  #paramBuf: GPUBuffer;
  #paramData: Float32Array;
  /** Public name -> buffer, for uploading initial state and reading results. */
  #byName: Map<string, Slot>;

  private constructor(init: {
    device: GPUDevice;
    ops: Op[];
    byName: Map<string, Slot>;
    owned: GPUBuffer[];
    paramBuf: GPUBuffer;
    paramData: Float32Array;
    paramNames: string[];
  }) {
    this.#device = init.device;
    this.#ops = init.ops;
    this.#byName = init.byName;
    this.#owned = init.owned;
    this.#paramBuf = init.paramBuf;
    this.#paramData = init.paramData;
    this.paramNames = init.paramNames;
  }

  static async create(
    device: GPUDevice,
    external: OpPlan,
    spec: PlanSpec,
    host: HostBuffers,
    /** Overrides what the device allows; only tests pass it. */
    operandBudget?: number,
  ): Promise<ModelPlan> {
    const { fn } = spec;

    const slots = new Map<string, Slot>();
    const byName = new Map<string, Slot>();
    const owned: GPUBuffer[] = [];
    /** Scalars the .m computes from its parameters, by cName. */
    const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
    /** Grid fields one kernel may read on this device (see fitToBudget). */
    const budget = operandBudget ?? kernelOperandBudget(device);
    /** How many kernels a line has been split into, for naming the pieces. */
    let splits = 0;

    const alloc = (label: string, count: number): Slot => {
      const buffer = makeBuffer(device, label, count);
      owned.push(buffer);
      return { buffer, count };
    };

    // Arguments, bound by what the function's signature declares. Array
    // arguments come from the shared pool, so a value one function returns is
    // the same buffer the next one reads. Scalar parameters share one small
    // storage buffer, in signature order.
    const paramNames: string[] = [];
    const paramSlots = new Map<string, number>();
    for (const p of fn.params) {
      if (p.binding.kind === 'tensor') {
        const count = p.binding.shape.reduce((x, y) => x * y, 1);
        const slot = host.ensure(p.name, count);
        slots.set(p.cName, slot);
        byName.set(p.name, slot);
      } else if (p.binding.kind === 'param') {
        paramSlots.set(p.cName, paramNames.length);
        paramNames.push(p.name);
      }
      // `const` arguments are exact in the IR and fold into the kernels.
    }
    const paramData = new Float32Array(Math.max(1, paramNames.length));
    const paramBuf = device.createBuffer({
      label: 'mgpu-params',
      size: 4 * paramData.length,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    });

    const ops: Op[] = [];
    for (const stmt of fn.body) {
      await planStatement(stmt);
    }

    planFeedback();

    return new ModelPlan({ device, ops, byName, owned, paramBuf, paramData, paramNames });

    async function planStatement(stmt: IRStmt): Promise<void> {
      if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
      if (stmt.kind !== 'Assign') {
        throw new UnsupportedOnGpu(
          `a model function body may only contain assignments ` +
            `(found '${stmt.kind}')`,
          stmt.span,
        );
      }
      if (!isNumeric(stmt.ty)) {
        throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric value`, stmt.span);
      }
      if (!isTensor(stmt.ty)) {
        // A scalar the model derives from its parameters (`om = 2*pi*f`). It
        // gets no buffer and no dispatch: the kernels that read it bind it as
        // a `let` in their prologue.
        derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
        return;
      }
      const count = numel(stmt.ty);

      // Reuse the destination buffer across steps: the same cName always maps
      // to the same buffer, so a step allocates nothing.
      let dest = slots.get(stmt.cName);
      if (!dest) {
        dest = alloc(`mgpu-${stmt.name}`, count);
        slots.set(stmt.cName, dest);
      } else if (dest.count !== count) {
        throw new UnsupportedOnGpu(
          `'${stmt.name}' changes size between assignments`,
          stmt.span,
        );
      }
      byName.set(stmt.name, dest);

      const ext = externalCall(stmt);
      if (ext) return planExternal(stmt, ext, dest);

      // Element-wise. Checked against the whole line first, so a broadcasting
      // mistake is reported against what was written rather than against a
      // fragment of it.
      checkShapes(stmt.expr, stmt.ty, stmt.name);
      const expr = await fitToBudget(stmt.expr, stmt.name, count, stmt.span);
      await emitElementwise(stmt.name, expr, stmt.ty, stmt.span, dest, count, stmt.cName);
    }

    /**
     * Emit one element-wise kernel: `expr` evaluated at every index into
     * `dest`. `selfCName` is the variable being assigned, if any, so an
     * in-place update can be spotted.
     */
    async function emitElementwise(
      name: string,
      expr: IRExpr,
      ty: NumericType,
      span: unknown,
      dest: Slot,
      count: number,
      selfCName?: string,
    ): Promise<void> {
      // Collect the distinct tensor operands and give them dense binding slots.
      const tensors = new Map<string, number>();
      collectTensorVars(expr, (cName) => {
        if (!tensors.has(cName)) tensors.set(cName, tensors.size);
      });

      const label = `${name} = <${count} elements, element-wise>`;
      const kernel = buildKernel(
        { kind: 'Assign', name, ty, expr, span } as unknown as Assign,
        {
          tensors,
          params: paramSlots,
          scalars: derivedScalars,
        } satisfies KernelInputs,
        count,
        label,
      );

      const bindGroupLayout = kernelLayout(device, tensors.size);
      const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);

      // WebGPU forbids aliasing a writable storage binding with another
      // binding in the same group, so an in-place update (`p = p + 1`) writes
      // to scratch and copies back. Element-wise kernels only ever touch
      // their own index, so the copy is the only cost.
      const aliased = selfCName !== undefined && tensors.has(selfCName);
      const target = aliased ? alloc(`mgpu-${name}-scratch`, count) : dest;

      const entries: GPUBindGroupEntry[] = [
        { binding: 0, resource: { buffer: target.buffer } },
      ];
      for (const [cName, i] of tensors) {
        const s = slots.get(cName);
        if (!s) {
          throw new UnsupportedOnGpu(`'${name}' reads a value with no buffer`, span);
        }
        entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
      }
      entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });

      ops.push({
        kind: 'kernel',
        pipeline,
        bindGroup: device.createBindGroup({ layout: bindGroupLayout, entries }),
        count,
        label,
        copyBack: aliased
          ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
          : undefined,
      });
    }

    /**
     * Split an expression that reads more grid fields than one kernel may bind.
     *
     * A kernel binds one storage buffer per distinct field it reads, plus its
     * output and the parameter block, and WebGPU guarantees only eight per
     * compute stage — fewer in compatibility mode. numbl's inline pass, which
     * is what makes one source line become one kernel, does not know about
     * that limit, and a model has no way to ask it for less: a temporary used
     * once is exactly what it folds away.
     *
     * So the budget is enforced here instead. Any child subtree that reads
     * more than one field is evaluated into its own buffer and replaced by a
     * reference to it, which leaves the parent reading at most one field per
     * child. The result is the same arithmetic in a few more passes over
     * memory, and it only happens on a line that would not otherwise compile.
     */
    async function fitToBudget(
      expr: IRExpr,
      hint: string,
      count: number,
      span: unknown,
    ): Promise<IRExpr> {
      if (tensorCount(expr) <= budget) return expr;

      const fit = async (e: IRExpr): Promise<IRExpr> => {
        if (tensorCount(e) <= budget) return e;
        const kids = children(e);
        if (!kids.length) return e;
        const out: IRExpr[] = [];
        for (const kid of kids) {
          const fitted = await fit(kid);
          out.push(tensorCount(fitted) > 1 ? await hoist(fitted) : fitted);
        }
        return withChildren(e, out);
      };

      /** Evaluate a subtree into its own buffer and hand back a reference. */
      const hoist = async (e: IRExpr): Promise<IRExpr> => {
        if (!isNumeric(e.ty) || !isTensor(e.ty)) return e;
        const name = `${hint}_part${++splits}`;
        const cName = `mgpu_split_${splits}`;
        const slot = alloc(`mgpu-${name}`, count);
        slots.set(cName, slot);
        await emitElementwise(name, e, e.ty, e.span, slot, count);
        return { kind: 'Var', name, cName, ty: e.ty, span: e.span } as IRExpr;
      };

      const fitted = await fit(expr);
      if (tensorCount(fitted) > budget) {
        throw new UnsupportedOnGpu(
          `'${hint}' reads ${tensorCount(fitted)} grid fields at once, and this ` +
            `device allows ${budget} per kernel. Compute part of it into a ` +
            `named field on a line of its own.`,
          span,
        );
      }
      return fitted;
    }

    /** `lp = lapw(p, wall)`: one dispatch, sources and dst distinct. */
    async function planExternal(
      stmt: Assign,
      ext: { name: OpKind; args: (IRExpr & { kind: 'Var' })[] },
      dest: Slot,
    ): Promise<void> {
      const contract = external.spec(ext.name);
      const argSlots = ext.args.map((arg, i) => {
        const s = slots.get(arg.cName);
        if (!s) {
          throw new UnsupportedOnGpu(
            `'${ext.name}' reads '${arg.name}', which has no buffer`,
            stmt.span,
          );
        }
        if (s.count !== contract.argCounts[i]) {
          throw new UnsupportedOnGpu(
            `'${ext.name}' wants ${contract.argCounts[i]} elements for its ` +
              `argument ${i + 1}, but '${arg.name}' holds ${s.count}`,
            stmt.span,
          );
        }
        // An op reads its neighbours, so unlike an element-wise kernel it
        // cannot be routed through scratch and copied back — the neighbours
        // would already have been overwritten. WebGPU forbids the aliasing
        // outright anyway; refuse rather than silently reroute.
        if (s.buffer === dest.buffer) {
          throw new UnsupportedOnGpu(
            `'${stmt.name} = ${ext.name}(...)' reads and writes the same ` +
              `buffer through '${arg.name}'; assign to a new name instead`,
            stmt.span,
          );
        }
        return s;
      });
      if (dest.count !== contract.outCount) {
        throw new UnsupportedOnGpu(
          `'${ext.name}' produces a ${contract.outCount}-point field, but ` +
            `'${stmt.name}' holds ${dest.count}`,
          stmt.span,
        );
      }
      const built = await external.pipeline(ext.name);
      ops.push({
        kind: 'external',
        pipeline: built.pipeline,
        bindGroup: external.createBinding(
          ext.name,
          argSlots.map((s) => s.buffer),
          dest.buffer,
        ),
        workgroups: external.workgroups(contract.outCount),
        label: `${stmt.name} = ${ext.name}(${ext.args.map((a) => a.name).join(', ')})`,
      });
    }

    /**
     * Feed declared outputs back into the argument buffers they replace, so
     * the next call reads what this one produced.
     *
     * The copies are not independent: a model whose new history field is the
     * old current one (`function [pn, pold] = step(p, pm, ...)`) has an output
     * whose *source* is another output's *destination*. Doing them in order
     * would then copy the new value where the old one was wanted, silently.
     * So any source that a previous copy overwrites is staged through scratch
     * first — normally none, since a model that writes `pold = p;` gets its
     * own buffer from the copy kernel that line plans to.
     */
    function planFeedback(): void {
      const copies: { from: Slot; to: Slot; label: string }[] = [];
      fn.outputs.forEach((out, i) => {
        const to = spec.feedback[i];
        if (!to) return;
        const src = slots.get(out.cName);
        const dst = host.get(to);
        if (!src) {
          throw new UnsupportedOnGpu(
            `'${fn.name}' declares the output '${out.name}' but never assigns it`,
          );
        }
        if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
        if (src.count !== dst.count) {
          throw new UnsupportedOnGpu(
            `'${out.name}' (${src.count} elements) cannot feed '${to}' (${dst.count})`,
          );
        }
        copies.push({ from: src, to: dst, label: `${out.name} -> ${to}` });
      });

      const written = new Set<GPUBuffer>();
      for (const c of copies) written.add(c.to.buffer);
      for (const c of copies) {
        // Only a source another copy overwrites needs staging, and only if it
        // is not that same copy's own destination (which is a no-op anyway).
        if (c.from.buffer !== c.to.buffer && written.has(c.from.buffer)) {
          const scratch = alloc(`mgpu-feedback-scratch`, c.from.count);
          ops.push({
            kind: 'copy',
            from: c.from.buffer,
            to: scratch.buffer,
            bytes: 4 * c.from.count,
            label: `${c.label} (staged)`,
          });
          c.from = scratch;
        }
      }
      for (const c of copies) {
        if (c.from.buffer === c.to.buffer) continue; // already in place
        ops.push({
          kind: 'copy',
          from: c.from.buffer,
          to: c.to.buffer,
          bytes: 4 * c.from.count,
          label: c.label,
        });
      }
    }
  }

  /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
  setParams(values: Record<string, number>): void {
    this.paramNames.forEach((name, i) => {
      const v = values[name];
      this.#paramData[i] = Number.isFinite(v) ? v : 0;
    });
    this.#device.queue.writeBuffer(
      this.#paramBuf,
      0,
      this.#paramData as Float32Array<ArrayBuffer>,
    );
  }

  /** Buffer holding the named value, or undefined if the .m never binds it. */
  buffer(name: string): GPUBuffer | undefined {
    return this.#byName.get(name)?.buffer;
  }

  elementCount(name: string): number | undefined {
    return this.#byName.get(name)?.count;
  }

  /**
   * Record `steps` passes of this plan. Synchronous: no awaits, no readback.
   * The dispatches share one compute pass, which WebGPU executes in submission
   * order with a barrier between them.
   *
   * `after` runs once per step, inside the same submission — which is what
   * lets the microphone sample every timestep rather than every frame.
   */
  encodeSteps(
    encoder: GPUCommandEncoder,
    steps: number,
    after?: (encoder: GPUCommandEncoder) => void,
  ): void {
    for (let s = 0; s < steps; s++) {
      this.#encodeOps(encoder);
      after?.(encoder);
    }
  }

  /** Record one pass over the op sequence into `encoder`. */
  #encodeOps(encoder: GPUCommandEncoder): void {
    let pass: GPUComputePassEncoder | null = null;
    const inPass = (): GPUComputePassEncoder => {
      if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
      return pass;
    };
    const endPass = (): void => {
      if (pass) {
        pass.end();
        pass = null;
      }
    };
    for (const op of this.#ops) {
      switch (op.kind) {
        case 'kernel': {
          const p = inPass();
          p.setPipeline(op.pipeline);
          p.setBindGroup(0, op.bindGroup);
          p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
          if (op.copyBack) {
            endPass();
            encoder.copyBufferToBuffer(
              op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
            );
          }
          break;
        }
        case 'external': {
          const p = inPass();
          p.setPipeline(op.pipeline);
          p.setBindGroup(0, op.bindGroup);
          p.dispatchWorkgroups(op.workgroups);
          break;
        }
        case 'copy':
          endPass();
          encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
          break;
      }
    }
    endPass();
  }

  /** Human-readable op sequence — what the .m actually compiled to. */
  describe(): string[] {
    return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
  }

  destroy(): void {
    for (const b of this.#owned) b.destroy();
    this.#paramBuf.destroy();
    this.#owned.length = 0;
  }
}

/** `lp = lapw(p, wall)` -> the op's name and its arguments. */
function externalCall(
  stmt: Assign,
): { name: OpKind; args: (IRExpr & { kind: 'Var' })[] } | null {
  const e = stmt.expr;
  if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
  const spec = EXTERNAL_OPS.get(e.name)!;
  if (e.args.length !== spec.args.length) {
    throw new UnsupportedOnGpu(
      `'${e.name}' takes ${spec.args.length} argument(s), got ${e.args.length}`,
      stmt.span,
    );
  }
  const args = e.args.map((arg) => {
    if (arg.kind !== 'Var') {
      throw new UnsupportedOnGpu(
        `'${e.name}' must be applied to variables, not expressions — ` +
          `name the field first`,
        stmt.span,
      );
    }
    return arg;
  });
  return { name: e.name as OpKind, args };
}

/** Distinct grid fields an expression reads — its storage-buffer cost. */
function tensorCount(e: IRExpr): number {
  const seen = new Set<string>();
  collectTensorVars(e, (c) => seen.add(c));
  return seen.size;
}

/** The subexpressions of a node, in evaluation order. Leaves have none. */
function children(e: IRExpr): IRExpr[] {
  switch (e.kind) {
    case 'Binary':
      return [e.left, e.right];
    case 'Unary':
      return [e.operand];
    case 'Call':
      return e.args;
    default:
      return [];
  }
}

/** The same node with its subexpressions replaced. */
function withChildren(e: IRExpr, kids: IRExpr[]): IRExpr {
  switch (e.kind) {
    case 'Binary':
      return { ...e, left: kids[0], right: kids[1] };
    case 'Unary':
      return { ...e, operand: kids[0] };
    case 'Call':
      return { ...e, args: kids };
    default:
      return e;
  }
}

function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
  const walk = (x: IRExpr): void => {
    switch (x.kind) {
      case 'Var':
        if (isTensor(x.ty)) 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);
}