lemmascript 0.5.7 → 0.5.9

package/tools/dist/resolve.js

@@ -5,7 +5,7 @@
  * No mutation — each let extends the chain, lookup walks it.
  */
 import { isBigInt } from "./typedir.js";
-import { parseTsType } from "./types.js";
+import { parseTsType, tyToCanonical } from "./types.js";
 import { parseExpr } from "./specparser.js";
 function lookup(env, name) {
     if (!env)
@@ -349,6 +349,14 @@ function isUnmodeledTy(ty, typeDecls) {
     }
     return false;
 }
+/** A `user` type that resolves to a string-union declare-type — runs as a plain
+ *  string at runtime, so it's a refinement of `string`, not an opaque blob. */
+function isStringUnionTy(ty, typeDecls) {
+    if (ty.kind !== "user")
+        return false;
+    const base = ty.name.includes("<") ? ty.name.slice(0, ty.name.indexOf("<")) : ty.name;
+    return typeDecls.some(d => d.name === base && d.kind === "string-union");
+}
 /** Infer quantifier variable type from usage in body.
  *  If the variable is used as a map/set key (e.g. map.has(k), map.get(k)),
  *  return the collection's key type. Otherwise return null (default to int). */
@@ -457,6 +465,16 @@ function tyToTsStr(ty) {
     return undefined;
 }
 function inferLambdaParamTypes(fn, rawArgs, ctx) {
+    // sort's comparator takes two params, both the element type.
+    if (fn.kind === "field" && fn.obj.ty.kind === "array" && fn.field === "sort" &&
+        rawArgs.length >= 1 && rawArgs[0].kind === "lambda" && rawArgs[0].params.length >= 1) {
+        const tsType = tyToTsStr(fn.obj.ty.elem);
+        if (tsType) {
+            const lam = rawArgs[0];
+            const updatedParams = lam.params.map(p => (p.tsType ? p : { ...p, tsType }));
+            return [{ ...lam, params: updatedParams }, ...rawArgs.slice(1)];
+        }
+    }
     if (fn.kind === "field" && fn.obj.ty.kind === "array" &&
         ["map", "filter", "every", "some", "find", "findLast", "findIndex"].includes(fn.field) &&
         rawArgs.length >= 1 && rawArgs[0].kind === "lambda" &&
@@ -573,7 +591,9 @@ function inferMethodReturnTy(fn, args, ctx) {
             return objTy.elem;
         if (fn.field === "pop")
             return { kind: "optional", inner: objTy.elem };
-        if (fn.field === "push" || fn.field === "concat")
+        if (fn.field === "push" || fn.field === "unshift" || fn.field === "concat")
+            return objTy;
+        if (fn.field === "sort")
             return objTy;
         if (fn.field === "filter")
             return objTy;
@@ -639,6 +659,106 @@ function lookupFieldTy(objTy, field, ctx) {
     return { ty: { kind: "unknown" }, isDiscriminant: false };
 }
 // ── Resolve expressions ──────────────────────────────────────
+// Fresh-binder counter for the someMatches synthesized by object-spread merge.
+let mergeBinder = 0;
+/** Expand `{ ...base, ...override }` into a faithful field-wise merge. Driven by
+ *  the result record type's fields: an override field wins when present, else the
+ *  base's field shows through. Optional fields decide presence at runtime
+ *  (`Some?`); an `Option`-typed override is the whole merge guarded by its tag. */
+function resolveRecordMerge(base, override, ctx) {
+    const tbase = resolveExpr(base, ctx);
+    const tover = resolveExpr(override, ctx);
+    const overInner = tover.ty.kind === "optional" ? tover.ty.inner : tover.ty;
+    // Result record type: prefer the override's (an optional override still merges
+    // into its inner type), else the base's.
+    const rTy = overInner.kind === "user" ? overInner
+        : tbase.ty.kind === "user" ? tbase.ty : null;
+    const decl = rTy ? ctx.typeDecls.find(d => d.name === rTy.name && d.kind === "record") : undefined;
+    if (!rTy || !decl?.fields) {
+        throw new Error(`object spread merge { ...a, ...b } needs a known record type for both operands ` +
+            `(base: ${tyToCanonical(tbase.ty)}, override: ${tyToCanonical(tover.ty)})`);
+    }
+    if (tbase.ty.kind === "optional") {
+        throw new Error(`object spread merge with an optional base operand is not supported (base: ${tyToCanonical(tbase.ty)})`);
+    }
+    const userTy = rTy;
+    const fields = decl.fields;
+    // Build the merged literal from concrete base/override values, both : userTy.
+    const merged = (bv, ov) => ({
+        kind: "record", spread: null, ty: userTy,
+        fields: fields.map(f => {
+            const ft = f.type;
+            const ovf = { kind: "field", obj: ov, field: f.name, ty: ft };
+            if (ft.kind !== "optional")
+                return { name: f.name, value: ovf }; // required: override always provides
+            // optional: override field wins iff present, else base's field
+            const bvf = { kind: "field", obj: bv, field: f.name, ty: ft };
+            const binder = `_m${mergeBinder++}`;
+            // someBody is the unwrapped present value; transform re-wraps each arm in
+            // the backend's Some constructor (Dafny `Some`, Lean `Option.some`).
+            return { name: f.name, value: {
+                    kind: "someMatch", scrutinee: ovf, binder, binderTy: ft.inner,
+                    someBody: { kind: "var", name: binder, ty: ft.inner }, noneBody: bvf, ty: ft,
+                } };
+        }),
+    });
+    if (tover.ty.kind === "optional") {
+        // override may be absent (undefined spreads nothing) → base unchanged
+        const binder = `_mo${mergeBinder++}`;
+        return {
+            kind: "someMatch", scrutinee: tover, binder, binderTy: userTy,
+            someBody: merged(tbase, { kind: "var", name: binder, ty: userTy }), noneBody: tbase, ty: userTy,
+        };
+    }
+    return merged(tbase, tover);
+}
+/** `rec[k]` where `rec` is a record and `k` an enum of its field names. A *named*
+ *  string-union key is a datatype → `match k { case f => rec.f }`; an *inline*
+ *  union (`"a" | "b"`, a bare string carrying its members) → an equality chain
+ *  `if k === "a" then rec.a else …`. Either way the chain/match covers exactly
+ *  the key's values, so a subset key stays sound. Returns null if the shape
+ *  doesn't apply (caller falls back to plain index). */
+function tryRecordIndexByEnum(obj, idx, ctx) {
+    const objTy = obj.ty, keyTy = idx.ty;
+    if (objTy.kind !== "user")
+        return null;
+    const rec = ctx.typeDecls.find(d => d.name === objTy.name && d.kind === "record");
+    if (!rec?.fields)
+        return null;
+    const fieldByName = new Map(rec.fields.map(f => [f.name, f]));
+    const fieldTy = (v) => fieldByName.get(v).type ?? { kind: "unknown" };
+    const field = (v) => ({ kind: "field", obj, field: v, ty: fieldTy(v) });
+    // The key's members, and whether it's a datatype (named) or a bare string (inline).
+    let values = null;
+    let datatype = null;
+    if (keyTy.kind === "user") {
+        const keyEnum = ctx.typeDecls.find(d => d.name === keyTy.name && d.kind === "string-union");
+        if (keyEnum?.values?.length) {
+            values = keyEnum.values;
+            datatype = keyEnum.name;
+        }
+    }
+    else if (keyTy.kind === "string" && keyTy.values?.length) {
+        values = keyTy.values;
+    }
+    if (!values || !values.every(v => fieldByName.has(v)))
+        return null; // key isn't a subset of fields
+    if (datatype) {
+        return {
+            kind: "tagMatch", scrutinee: idx, typeName: datatype,
+            cases: values.map(v => ({ variant: v, body: field(v) })), fallthrough: null, ty: fieldTy(values[0]),
+        };
+    }
+    // Inline union: fold right into an equality chain, last member as the bare else.
+    let expr = field(values[values.length - 1]);
+    for (let i = values.length - 2; i >= 0; i--) {
+        expr = {
+            kind: "conditional", ty: fieldTy(values[i]), then: field(values[i]), else: expr,
+            cond: { kind: "binop", op: "===", left: idx, right: { kind: "str", value: values[i], ty: { kind: "string" } }, ty: { kind: "bool" } },
+        };
+    }
+    return expr;
+}
 function resolveExpr(e, ctx) {
     switch (e.kind) {
         case "var":
@@ -812,6 +932,9 @@ function resolveExpr(e, ctx) {
                 idxTy = narrowed ? obj.ty.value : { kind: "optional", inner: obj.ty.value };
             }
             else {
+                const recIdx = tryRecordIndexByEnum(obj, idx, ctx);
+                if (recIdx)
+                    return recIdx;
                 idxTy = { kind: "unknown" };
             }
             return { kind: "index", obj, idx, ty: idxTy };
@@ -842,8 +965,10 @@ function resolveExpr(e, ctx) {
             // left ?? right — result type is left's inner (when left is optional)
             // or just left's type, unified with right's type.
             const left = resolveExpr(e.left, ctx);
-            const right = resolveExpr(e.right, ctx);
             const ty = left.ty.kind === "optional" ? left.ty.inner : left.ty;
+            // The default shares the result type, so coerce a string literal to a
+            // string-union enum (e.g. `availableLevels[0] ?? "off"`).
+            const right = coerceStr(resolveExpr(e.right, ctx), ty);
             return { kind: "nullish", left, right, ty };
         }
         case "optChain": {
@@ -944,6 +1069,8 @@ function resolveExpr(e, ctx) {
             });
             return { kind: "record", spread, fields, ty: recordTy ?? ty };
         }
+        case "recordMerge":
+            return resolveRecordMerge(e.base, e.override, ctx);
         case "result":
             // \result desugars to a regular var so all the variable-narrowing
             // machinery (env lookup, optional checks, path matching) just works.
@@ -967,8 +1094,14 @@ function resolveExpr(e, ctx) {
             // anonymous tuple (mirrors return-position and call-argument records, which
             // get their type via ctx.returnTy). Only narrow when the context type is an
             // array; otherwise leave ctx untouched.
-            const elemCtx = ctx.returnTy.kind === "array" ? { ...ctx, returnTy: ctx.returnTy.elem } : ctx;
-            const elems = e.elems.map(el => resolveExpr(el, elemCtx));
+            const expectedElem = ctx.returnTy.kind === "array" ? ctx.returnTy.elem : null;
+            const elemCtx = expectedElem ? { ...ctx, returnTy: expectedElem } : ctx;
+            const elems = e.elems.map(el => {
+                const r = resolveExpr(el, elemCtx);
+                // Coerce a bare string-literal element to a string-union enum (e.g.
+                // `["off", …]: ModelThinkingLevel[]`), like return/arg positions.
+                return expectedElem ? coerceStr(r, expectedElem) : r;
+            });
             const elemTy = elems.length > 0 ? elems[0].ty : { kind: "unknown" };
             return { kind: "arrayLiteral", elems, ty: { kind: "array", elem: elemTy } };
         }
@@ -1140,6 +1273,11 @@ function resolveStmt(s, ctx) {
                     ? { kind: "optional", inner: init.ty }
                     : init.ty;
             }
+            else if (declTy.kind === "string" && isStringUnionTy(init.ty, ctx.typeDecls) && !ctx.overrides.has(s.name)) {
+                // ts-morph widened a string-union to `string`; keep the initializer's
+                // datatype so `local === "lit"` lowers to a discriminant test.
+                ty = init.ty;
+            }
             else if ((declTy.kind === "int" || declTy.kind === "nat") && init.ty.kind === "real" && !ctx.overrides.has(s.name)) {
                 // TS infers `number` (→ int/nat) for an expression LS computes as `real`
                 // (e.g. `a / b`, now real division). `number` can't tell them apart, so
@@ -1332,6 +1470,10 @@ function collectCallsExpr(e, fns, out) {
             for (const f of e.fields)
                 collectCallsExpr(f.value, fns, out);
             return;
+        case "recordMerge":
+            collectCallsExpr(e.base, fns, out);
+            collectCallsExpr(e.override, fns, out);
+            return;
         case "arrayLiteral":
             for (const el of e.elems)
                 collectCallsExpr(el, fns, out);

package/tools/dist/transform.js

@@ -4,6 +4,7 @@
  * Consumes resolved types and classifications.
  * No type lookups, no string parsing, no re-inference.
  */
+import { anyExprInStmts } from "./ir.js";
 import { parseTsType } from "./types.js";
 // ── Generic IR walkers ──────────────────────────────────────
 /**
@@ -23,7 +24,8 @@ function mapExpr(e, f) {
         case "str":
         case "emptyMap":
         case "emptySet":
-        case "havoc": return e;
+        case "havoc":
+        case "default": return e;
         case "mapLiteral": return { ...e, entries: e.entries.map(en => ({ key: r(en.key), value: r(en.value) })) };
         case "constructor": return e.args ? { ...e, args: e.args.map(r) } : e;
         case "binop": return { ...e, left: r(e.left), right: r(e.right) };
@@ -630,6 +632,17 @@ function lowerExpr(e, binds) {
                     return { kind: "field", obj: transformExpr(e.obj), field: "true_?" };
                 }
             }
+            // Union destructor: `x.field` where x is a discriminated union and `field`
+            // is a data field of one of its variants. Dafny reads the destructor
+            // directly; Lean has no field projection on a multi-ctor inductive, so tag
+            // the node with the union's base name and let the Lean emitter `match`.
+            if (e.obj.ty.kind === "user") {
+                const baseName = e.obj.ty.name.includes("<") ? e.obj.ty.name.slice(0, e.obj.ty.name.indexOf("<")) : e.obj.ty.name;
+                const decl = _typeDecls.find(d => d.name === baseName && d.kind === "discriminated-union");
+                if (decl?.variants?.some(v => v.fields.some(f => f.name === e.field))) {
+                    return { kind: "field", obj: transformExpr(e.obj), field: e.field, fromUnion: baseName };
+                }
+            }
             return { kind: "field", obj: transformExpr(e.obj), field: e.field };
         case "index": {
             const idx = transformExpr(e.idx);
@@ -1106,6 +1119,188 @@ function buildPopLowering(arrName, arrTy) {
     const guardedTrunc = { kind: "if", cond: isNonEmpty, then: truncated, else: arrVar };
     return { optValue, guardedTrunc };
 }
+/**
+ * Velvet (Lean backend) rejects `return` inside a loop. For a function shaped as
+ * `…; while (…) { … early returns … }; return fallthrough`, hoist a mutable
+ * result variable: each in-loop `return e` becomes `_loopRet := e; break`, then
+ * `return _loopRet` after the loop. The trailing fallthrough return seeds
+ * `_loopRet`, so a normal loop exit returns it unchanged. Dafny is untouched —
+ * it keeps the native early returns.
+ */
+function stmtsContainReturn(stmts) {
+    for (const s of stmts) {
+        if (s.kind === "return")
+            return true;
+        if (s.kind === "if" && (stmtsContainReturn(s.then) || stmtsContainReturn(s.else)))
+            return true;
+        if (s.kind === "match" && s.arms.some(a => stmtsContainReturn(a.body)))
+            return true;
+        // Don't descend into a nested while — its returns would target that loop.
+    }
+    return false;
+}
+function replaceReturnsWithBreak(stmts, retVar) {
+    return stmts.flatMap((s) => {
+        if (s.kind === "return")
+            return [{ kind: "assign", target: retVar, value: s.value }, { kind: "break" }];
+        if (s.kind === "if")
+            return [{ ...s, then: replaceReturnsWithBreak(s.then, retVar), else: replaceReturnsWithBreak(s.else, retVar) }];
+        if (s.kind === "match")
+            return [{ ...s, arms: s.arms.map(a => ({ ...a, body: replaceReturnsWithBreak(a.body, retVar) })) }];
+        return [s];
+    });
+}
+// Seed value for `_loopRet` when a return-in-loop function has no trailing
+// fallthrough return. Readable literals for the primitives; everything else
+// (user types, arrays, maps, optionals) gets a typed `default : T` rather than a
+// type-incorrect `0` — these types derive `Inhabited`, so the default exists.
+function defaultExprForTy(ty) {
+    switch (ty.kind) {
+        case "bool": return { kind: "bool", value: false };
+        case "string": return { kind: "str", value: "" };
+        case "nat":
+        case "int":
+        case "real": return { kind: "num", value: 0 };
+        default: return { kind: "default", type: ty };
+    }
+}
+function eliminateReturnInLoops(stmts, retTy, resultInvariants) {
+    const out = [];
+    for (let i = 0; i < stmts.length; i++) {
+        const s = stmts[i];
+        if (s.kind === "while" && stmtsContainReturn(s.body)) {
+            const retVar = "_loopRet";
+            const next = stmts[i + 1];
+            const init = next && next.kind === "return" ? next.value : defaultExprForTy(retTy);
+            out.push({ kind: "let", name: retVar, type: retTy, mutable: true, value: init });
+            // The result variable carries the postcondition (the function's `ensures`
+            // with `\result` → `_loopRet`) as a loop invariant: it holds initially (the
+            // fallthrough seed) and after each `_loopRet := e; break`, so the
+            // postcondition is re-established from the invariant when the loop exits.
+            // Where Dafny discharged the postcondition at each `return` site, the
+            // break-rewrite discharges it once, after the loop.
+            out.push({ ...s, invariants: [...s.invariants, ...resultInvariants], body: replaceReturnsWithBreak(s.body, retVar) });
+            out.push({ kind: "return", value: { kind: "var", name: retVar } });
+            if (next && next.kind === "return")
+                i++; // consume the seeded fallthrough return
+        }
+        else {
+            out.push(s);
+        }
+    }
+    return out;
+}
+/** True if `name` is referenced as a variable anywhere in `stmts`. Conservative:
+ *  counts every occurrence and ignores shadowing — a spurious hit only costs an
+ *  unused `let`, whereas a miss would leave a binder unbound. */
+function stmtsUseVar(stmts, name) {
+    return anyExprInStmts(stmts, e => e.kind === "var" && e.name === name);
+}
+/**
+ * Velvet (Lean backend) cannot synthesize a `WPGen` for a monadic statement
+ * `match` on a user inductive in a method body (loom's matcher-WPGen drops a
+ * `sorry`). loom *can* handle method-body `if`s (cf. examples/arrayEquals), so
+ * lower such matches to discriminator `if`-chains: `match x with | .C f.. => B`
+ * becomes `if x.C? then (let f := x.f-destructor; B) else …`. Constructor field
+ * binders become `let`s bound to the (already-supported) union destructor. Option
+ * matches and other non-user matches are left alone (loom handles those). Dafny
+ * is untouched — it keeps the native match.
+ */
+function matchToIfChains(stmts) {
+    return stmts.flatMap((s) => {
+        if (s.kind === "if")
+            return [{ ...s, then: matchToIfChains(s.then), else: matchToIfChains(s.else) }];
+        if (s.kind === "while")
+            return [{ ...s, body: matchToIfChains(s.body) }];
+        if (s.kind === "forin")
+            return [{ ...s, body: matchToIfChains(s.body) }];
+        if (s.kind !== "match")
+            return [s];
+        const arms = s.arms.map(a => ({ ...a, body: matchToIfChains(a.body) }));
+        const ctorArms = arms.filter(a => a.pattern.trim() !== "_");
+        const firstCtor = ctorArms[0]?.pattern.trim().split(/\s+/)[0].replace(/^\./, "");
+        const decl = firstCtor
+            ? _typeDecls.find(d => (d.kind === "discriminated-union" || d.kind === "string-union") &&
+                ((d.variants?.some(v => v.name === firstCtor)) || (d.values?.includes(firstCtor))))
+            : undefined;
+        if (!decl)
+            return [{ ...s, arms }]; // not a user union (e.g. Option) — leave as match
+        const scrutExpr = typeof s.scrutinee === "string" ? { kind: "var", name: s.scrutinee } : s.scrutinee;
+        const defaultArm = arms.find(a => a.pattern.trim() === "_");
+        let elseBranch = defaultArm ? defaultArm.body : [];
+        for (let k = ctorArms.length - 1; k >= 0; k--) {
+            const armBody = ctorArms[k].body;
+            if (armBody.length === 0)
+                continue; // empty arm (no-op) — let it fall through to `else`
+            const toks = ctorArms[k].pattern.trim().split(/\s+/);
+            const ctor = toks[0].replace(/^\./, "");
+            const binders = toks.slice(1);
+            const variant = decl.variants?.find(v => v.name === ctor);
+            // Discriminator condition. A nullary discriminated-union constructor would
+            // need `DecidableEq` for `x = .Ctor` (which such unions don't derive), so
+            // use a match-bool instead. Multi-field ctors already lower to a match-bool;
+            // string-unions derive DecidableEq, so `=` is fine there.
+            const cond = decl.kind === "discriminated-union" && binders.length === 0
+                ? { kind: "match", scrutinee: scrutExpr, arms: [
+                        { pattern: `.${ctor}`, body: { kind: "bool", value: true } },
+                        { pattern: "_", body: { kind: "bool", value: false } }
+                    ] }
+                : { kind: "binop", op: "=", left: scrutExpr, right: { kind: "constructor", name: ctor, type: decl.name } };
+            // Bind only the constructor-field binders the body actually uses, pinning the
+            // owning ctor so the destructor doesn't guess (variants share field names).
+            const lets = [];
+            binders.forEach((b, i) => {
+                const f = variant?.fields[i];
+                if (b !== "_" && f && stmtsUseVar(armBody, b))
+                    lets.push({
+                        kind: "let", name: b, type: f.type ?? { kind: "unknown" }, mutable: false,
+                        value: { kind: "field", obj: scrutExpr, field: f.name, fromUnion: decl.name, ctor },
+                    });
+            });
+            elseBranch = [{ kind: "if", cond, then: [...lets, ...armBody], else: elseBranch }];
+        }
+        return elseBranch;
+    });
+}
+/** True if `stmts` contain a `break` not nested inside another while. */
+function stmtsContainBreak(stmts) {
+    for (const s of stmts) {
+        if (s.kind === "break")
+            return true;
+        if (s.kind === "if" && (stmtsContainBreak(s.then) || stmtsContainBreak(s.else)))
+            return true;
+        if (s.kind === "match" && s.arms.some(a => stmtsContainBreak(a.body)))
+            return true;
+    }
+    return false;
+}
+/**
+ * loom's default loop-exit fact is `¬guard`, which does not hold for a loop that
+ * `break`s (a break exits with the guard still true). Such loops need an explicit
+ * `//@ done_with` in the TS source (`//@ done_with true` when the loop invariant
+ * alone carries the exit facts). Rather than silently supplying one, reject —
+ * the exit fact is part of the spec and belongs beside the invariants.
+ * Lean-only; Dafny derives loop-exit facts from the break sites themselves.
+ */
+function requireDoneWithForBreaks(stmts, fnName) {
+    for (const s of stmts) {
+        if (s.kind === "if") {
+            requireDoneWithForBreaks(s.then, fnName);
+            requireDoneWithForBreaks(s.else, fnName);
+        }
+        if (s.kind === "match")
+            for (const a of s.arms)
+                requireDoneWithForBreaks(a.body, fnName);
+        if (s.kind === "while") {
+            if (!s.doneWith && stmtsContainBreak(s.body)) {
+                throw new Error(`${fnName}: loop with break needs //@ done_with on the Lean backend ` +
+                    `(an early return in a loop also lowers to a break). ` +
+                    `Use //@ done_with true when the loop invariants carry the exit facts.`);
+            }
+            requireDoneWithForBreaks(s.body, fnName);
+        }
+    }
+}
 function eliminateTopLevelContinue(stmts) {
     const out = [];
     for (let i = 0; i < stmts.length; i++) {
@@ -1359,7 +1554,7 @@ function transformStmt(s, typeDecls) {
                 const recv = s.expr.fn.obj;
                 const f = s.expr.fn.field;
                 const isMutating = ((recv.ty.kind === "map" || recv.ty.kind === "set") && (f === "set" || f === "add" || f === "delete")) ||
-                    (recv.ty.kind === "array" && f === "push");
+                    (recv.ty.kind === "array" && (f === "push" || f === "unshift" || f === "sort"));
                 if (isMutating && recv.kind === "var") {
                     const { binds, expr } = liftMethodCalls(s.expr);
                     return [...binds, { kind: "assign", target: recv.name, value: expr }];
@@ -1836,6 +2031,7 @@ function transformTypeDecl(d) {
     else {
         return {
             kind: "structure", name: d.name,
+            typeParams: d.typeParams,
             fields: d.fields.map(f => ({ name: f.name, type: f.type })),
             deriving: ["Repr", "Inhabited", "DecidableEq"],
         };
@@ -1854,7 +2050,7 @@ function findReassignedNames(stmts, names) {
             // Mutating collection calls: s.add(x), m.set(k,v), s.delete(x), arr.push(x)
             if (s.kind === "expr" && s.expr.kind === "call" && s.expr.fn.kind === "field" &&
                 s.expr.fn.obj.kind === "var" && names.has(s.expr.fn.obj.name) &&
-                ["add", "set", "delete", "push"].includes(s.expr.fn.field)) {
+                ["add", "set", "delete", "push", "unshift"].includes(s.expr.fn.field)) {
                 found.add(s.expr.fn.obj.name);
             }
             if (s.kind === "if") {
@@ -1913,11 +2109,11 @@ function replaceVar(e, name, replacement, narrowing) {
 }
 // ── Top-level transform ──────────────────────────────────────
 /** Transform for Lean backend — same logic, Lean options. */
-export function transformModuleLean(mod, specImport) {
+export function transformModuleLean(mod, specImport, moduleBase) {
     const prev = _opts;
     _opts = LEAN_OPTIONS;
     try {
-        return transformModule(mod, specImport);
+        return transformModule(mod, specImport, moduleBase);
     }
     finally {
         _opts = prev;
@@ -1934,7 +2130,7 @@ export function transformModuleDafny(mod) {
         _opts = prev;
     }
 }
-export function transformModule(mod, specImport) {
+export function transformModule(mod, specImport, moduleBaseOverride) {
     _forofCounters.clear();
     _liftCounter = 0;
     _typeDecls = mod.typeDecls;
@@ -1987,6 +2183,9 @@ export function transformModule(mod, specImport) {
         }
     }
     const base = mod.file.split("/").pop()?.replace(/\.ts$/, "") ?? "module";
+    // Lean module base — overridable via `//@ lean-module` (see lsc.ts). Only the
+    // def→types import below reads it; Dafny never passes an override.
+    const moduleBase = moduleBaseOverride ?? base;
     // Externs: emit as top-of-file `function {:axiom}` (Dafny) declarations.
     // Any `requires`/`ensures` from the source declaration come along so callers
     // see the same spec the source itself verified. Substitute `\result` with the
@@ -2010,12 +2209,18 @@ export function transformModule(mod, specImport) {
         ? [{ kind: "namespace", name: "Pure", decls: pureDefs }]
         : [];
     if (typeDecls.length > 0 || pureDefs.length > 0 || externDecls.length > 0) {
+        // Declaration order differs by backend. Dafny allows forward references, so
+        // externs go first to be in scope everywhere. Lean requires definition-before-use:
+        // an extern's signature may reference a declared type (e.g. `estimateTokens(m: AgentMessage)`),
+        // so types must precede externs, which in turn precede the pure mirrors that may call them.
+        const decls = _opts.backend === "lean"
+            ? [...typeDecls, ...externDecls, ...pureNamespace]
+            : [...externDecls, ...typeDecls, ...pureNamespace];
         typesFile = {
             comment: "  Generated by lsc — Lean types and pure function mirrors.",
             imports: typesImports,
             options: [],
-            // Externs come first so they're in scope for every later declaration.
-            decls: [...externDecls, ...typeDecls, ...pureNamespace],
+            decls,
         };
     }
     // Def file: Velvet methods
@@ -2035,6 +2240,19 @@ export function transformModule(mod, specImport) {
         let body = pureDefNames.has(fn.name)
             ? [{ kind: "return", value: { kind: "app", fn: `Pure.${fn.name}`, args: fn.params.map(p => ({ kind: "var", name: p.name })) } }]
             : promoteAssignedLets(transformStmts(fn.body, mod.typeDecls));
+        // Lean-only method-body rewrites (Velvet can't WP-synthesize monadic matches
+        // and forbids `return` in loops):
+        //   1. monadic statement-`match` on a user union → discriminator `if`-chains
+        //   2. `return` inside a loop → result var + break (postcondition as invariant)
+        // Dafny keeps the native forms. Breaking loops (including those synthesized
+        // by rewrite 2) must carry a `//@ done_with` in the TS source — enforced here,
+        // since loom's default loop-exit fact `¬guard` does not hold across a break.
+        if (_opts.backend === "lean" && !pureDefNames.has(fn.name)) {
+            body = matchToIfChains(body);
+            const resultInvariants = fn.ensures.map(e => replaceVar(transformExpr(e), "\\result", { kind: "var", name: "_loopRet" }));
+            body = eliminateReturnInLoops(body, fn.returnTy, resultInvariants);
+            requireDoneWithForBreaks(body, fn.name);
+        }
         // Shadow reassigned parameters with mutable locals
         const paramNames = new Set(fn.params.map(p => p.name));
         const reassigned = findReassignedNames(fn.body, paramNames);
@@ -2081,7 +2299,7 @@ export function transformModule(mod, specImport) {
             methods: classMethods,
         };
     });
-    const defImport = specImport ?? (typesFile ? `«${base}.types»` : null);
+    const defImport = specImport ?? (typesFile ? `«${moduleBase}.types»` : null);
     const defBaseImports = defImport ? [defImport] : ["LemmaScript"];
     const defFile = {
         comment: "  Generated by lsc from " + (mod.file.split("/").pop() ?? "") + "\n  Do not edit — re-run `lsc gen` to regenerate.",

package/tools/dist/types.js

@@ -55,6 +55,20 @@ function tyFromTypeNode(tn) {
         if (normalized.length === 1 && "syntheticBool" in normalized[0])
             return { kind: "bool" };
         const nonNullish = normalized.filter(a => "syntheticBool" in a || !isNullish(a.node));
+        // Inline string-literal union (`"a" | "b"`, not a //@ declare-type): no datatype
+        // to resolve against, so lower to plain string (the arms are strings; == holds).
+        const isStrLit = (a) => !("syntheticBool" in a) && Node.isLiteralTypeNode(a.node) && a.node.getLiteral().getKind() === SyntaxKind.StringLiteral;
+        if (nonNullish.length >= 2 && nonNullish.every(isStrLit)) {
+            // Keep the literal members so `rec[k]` can lower to an equality chain.
+            const values = nonNullish.map(a => {
+                const lit = a.node;
+                const inner = Node.isLiteralTypeNode(lit) ? lit.getLiteral() : lit;
+                return Node.isStringLiteral(inner) ? inner.getLiteralValue() : inner.getText();
+            });
+            return normalized.some(a => !("syntheticBool" in a) && isNullish(a.node))
+                ? { kind: "optional", inner: { kind: "string", values } }
+                : { kind: "string", values };
+        }
         if (nonNullish.length === 1 && normalized.length >= 2) {
             const sole = nonNullish[0];
             const inner = "syntheticBool" in sole ? { kind: "bool" } : tyFromTypeNode(sole.node);