web/lib/programRepr.ts

// Shared parser: program_repr (typed tree, opaque IDs) -> render model.
//
// Grammar (engine_v2 full DSL):
//   Vector := Reduce(Matrix, agg)
//           | Combine(Vector, Vector, op)
//           | Split(Vector, predicate)
//           | FitApply(Vector, target)
//   Matrix := M
//           | Select(Matrix, FeatureSet)
//           | Search(Matrix, k)
//   Scalar := Associate(Vector, target, kind)
//           | Effect(Vector, target, kind)
//   FeatureSet := [id1,id2,...]
//
// We also accept the legacy engine_v1 string shape
//   "Fit(Reduce(Select(M, [N ids]), mean), ... -> target)"
// by surfacing a "Fit" wrapper node — useful when the page is showing
// an older persisted run.

export type ReprNode =
  | { kind: "M" }
  | { kind: "Select"; matrix: ReprNode; features: string[] }
  | { kind: "Search"; matrix: ReprNode; k: number }
  | { kind: "Reduce"; matrix: ReprNode; agg: string }
  | { kind: "Combine"; left: ReprNode; right: ReprNode; op: string }
  | { kind: "Split"; inner: ReprNode; predicate: string }
  | { kind: "FitApply"; inner: ReprNode; target: string }
  | { kind: "Associate"; inner: ReprNode; target: string; assocKind: string }
  | { kind: "Effect"; inner: ReprNode; target: string; assocKind: string }
  | { kind: "Fit"; children: ReprNode[]; output: string }
  | { kind: "Unknown"; text: string };

class _Parser {
  s: string;
  i: number;
  constructor(s: string) {
    this.s = s;
    this.i = 0;
  }
  peek(): string {
    return this.s[this.i] ?? "";
  }
  eof(): boolean {
    return this.i >= this.s.length;
  }
  skip() {
    while (!this.eof() && /\s/.test(this.peek())) this.i++;
  }
  expect(ch: string) {
    this.skip();
    if (this.peek() !== ch) {
      throw new Error(
        `expected '${ch}' at offset ${this.i}, got '${this.peek()}' in ${this.s}`,
      );
    }
    this.i++;
  }
  ident(): string {
    this.skip();
    const start = this.i;
    while (!this.eof() && /[A-Za-z0-9_]/.test(this.peek())) this.i++;
    return this.s.slice(start, this.i);
  }
  parseFeatureSet(): string[] {
    this.expect("[");
    const out: string[] = [];
    this.skip();
    if (this.peek() === "]") {
      this.i++;
      return out;
    }
    while (true) {
      this.skip();
      const start = this.i;
      while (!this.eof() && /[^,\]\s]/.test(this.peek())) this.i++;
      out.push(this.s.slice(start, this.i).trim());
      this.skip();
      if (this.peek() === ",") {
        this.i++;
        continue;
      }
      if (this.peek() === "]") {
        this.i++;
        return out;
      }
      throw new Error(`bad FeatureSet at offset ${this.i}`);
    }
  }
  parseExpr(): ReprNode {
    this.skip();
    const head = this.ident();
    if (head === "M") return { kind: "M" };
    if (head === "Select") {
      this.expect("(");
      const matrix = this.parseExpr();
      this.expect(",");
      const features = this.parseFeatureSet();
      this.expect(")");
      return { kind: "Select", matrix, features };
    }
    if (head === "Reduce") {
      this.expect("(");
      const matrix = this.parseExpr();
      this.expect(",");
      const agg = this.ident();
      this.expect(")");
      return { kind: "Reduce", matrix, agg };
    }
    if (head === "Combine") {
      this.expect("(");
      const left = this.parseExpr();
      this.expect(",");
      const right = this.parseExpr();
      this.expect(",");
      const op = this.ident();
      this.expect(")");
      return { kind: "Combine", left, right, op };
    }
    if (head === "Search") {
      this.expect("(");
      const matrix = this.parseExpr();
      this.expect(",");
      this.skip();
      const kStart = this.i;
      while (!this.eof() && /[0-9]/.test(this.peek())) this.i++;
      const k = Number.parseInt(this.s.slice(kStart, this.i), 10) || 1;
      this.expect(")");
      return { kind: "Search", matrix, k };
    }
    if (head === "Split") {
      this.expect("(");
      const inner = this.parseExpr();
      this.expect(",");
      const predicate = this.ident();
      this.expect(")");
      return { kind: "Split", inner, predicate };
    }
    if (head === "FitApply") {
      this.expect("(");
      const inner = this.parseExpr();
      this.expect(",");
      const target = this.ident();
      this.expect(")");
      return { kind: "FitApply", inner, target };
    }
    if (head === "Associate") {
      this.expect("(");
      const inner = this.parseExpr();
      this.expect(",");
      const target = this.ident();
      this.expect(",");
      const assocKind = this.ident();
      this.expect(")");
      return { kind: "Associate", inner, target, assocKind };
    }
    if (head === "Effect") {
      this.expect("(");
      const inner = this.parseExpr();
      this.expect(",");
      const target = this.ident();
      this.expect(",");
      const assocKind = this.ident();
      this.expect(")");
      return { kind: "Effect", inner, target, assocKind };
    }
    if (head === "Fit") {
      return this.parseFitTail();
    }
    throw new Error(`unknown head '${head}' at offset ${this.i}`);
  }

  /** Tolerant parser for the legacy v1 shape:
   *    Fit(Reduce(Select(M, [3 ids]), mean), ... -> target)
   * It doesn't always serialise gene IDs (sometimes it's "[N ids]"). We
   * pull out the comma-separated children and an "-> output" tail; we
   * keep raw text for unknown bits. */
  parseFitTail(): ReprNode {
    this.expect("(");
    const children: ReprNode[] = [];
    let output = "target";
    while (true) {
      this.skip();
      // Detect "-> output" marker.
      if (this.s.startsWith("->", this.i)) {
        this.i += 2;
        this.skip();
        const start = this.i;
        let depth = 0;
        while (!this.eof()) {
          const ch = this.peek();
          if (ch === "(") depth++;
          else if (ch === ")") {
            if (depth === 0) break;
            depth--;
          }
          this.i++;
        }
        output = this.s.slice(start, this.i).trim();
        break;
      }
      // Try to parse a known child. If it fails, skip until comma/paren.
      const safeStart = this.i;
      try {
        const child = this.parseExpr();
        children.push(child);
      } catch {
        this.i = safeStart;
        let depth = 0;
        const start = this.i;
        while (!this.eof()) {
          const ch = this.peek();
          if (ch === "(") depth++;
          else if (ch === ")") {
            if (depth === 0) break;
            depth--;
          } else if (ch === "," && depth === 0) break;
          this.i++;
        }
        children.push({ kind: "Unknown", text: this.s.slice(start, this.i).trim() });
      }
      this.skip();
      if (this.peek() === ",") {
        this.i++;
        continue;
      }
      break;
    }
    this.expect(")");
    return { kind: "Fit", children, output };
  }
}

export function parseProgramRepr(repr: string): ReprNode {
  const p = new _Parser(repr.trim());
  p.skip();
  return p.parseExpr();
}

export function safeParseProgramRepr(repr: string): {
  ok: boolean;
  tree: ReprNode | null;
  error?: string;
} {
  try {
    return { ok: true, tree: parseProgramRepr(repr) };
  } catch (e) {
    return {
      ok: false,
      tree: null,
      error: e instanceof Error ? e.message : String(e),
    };
  }
}

// Walk a tree and collect every FeatureSet leaf's IDs (union, ordered).
export function geneIds(tree: ReprNode): string[] {
  const out: string[] = [];
  const seen = new Set<string>();
  function visit(n: ReprNode) {
    switch (n.kind) {
      case "Select":
        for (const id of n.features) {
          if (!seen.has(id)) {
            seen.add(id);
            out.push(id);
          }
        }
        visit(n.matrix);
        return;
      case "Search":
      case "Reduce":
        visit(n.matrix);
        return;
      case "Combine":
        visit(n.left);
        visit(n.right);
        return;
      case "Split":
      case "FitApply":
      case "Associate":
      case "Effect":
        visit(n.inner);
        return;
      case "Fit":
        for (const c of n.children) visit(c);
        return;
      case "M":
      case "Unknown":
        return;
    }
  }
  visit(tree);
  return out;
}