lemmascript 0.1.0 → 0.2.0

package/tools/dist/transform.js

@@ -1,56 +1,78 @@
 /**
- * Transform — Typed IR → Lean IR.
+ * Transform — Typed IR → Backend IR.
  *
  * Consumes resolved types and classifications.
  * No type lookups, no string parsing, no re-inference.
  */
-import { parseTsType, tyToLean } from "./types.js";
+import { parseTsType } from "./types.js";
+// ── Generic IR walkers ──────────────────────────────────────
+/**
+ * Map over all sub-expressions in an Expr. `f` is called on each node;
+ * if it returns non-null, that replaces the node (and recursion stops).
+ * If it returns null, the walker recurses into children.
+ */
+function mapExpr(e, f) {
+    const hit = f(e);
+    if (hit)
+        return hit;
+    const r = (x) => mapExpr(x, f);
+    switch (e.kind) {
+        case "var":
+        case "num":
+        case "bool":
+        case "str":
+        case "constructor":
+        case "emptyMap":
+        case "emptySet":
+        case "havoc": return e;
+        case "binop": return { ...e, left: r(e.left), right: r(e.right) };
+        case "unop": return { ...e, expr: r(e.expr) };
+        case "implies": return { ...e, premises: e.premises.map(r), conclusion: r(e.conclusion) };
+        case "app": return { ...e, args: e.args.map(r) };
+        case "field": return { ...e, obj: r(e.obj) };
+        case "toNat": return { ...e, expr: r(e.expr) };
+        case "index": return { ...e, arr: r(e.arr), idx: r(e.idx) };
+        case "record": return { ...e, spread: e.spread ? r(e.spread) : null, fields: e.fields.map(fi => ({ ...fi, value: r(fi.value) })) };
+        case "arrayLiteral": return { ...e, elems: e.elems.map(r) };
+        case "if": return { ...e, cond: r(e.cond), then: r(e.then), else: r(e.else) };
+        case "match": return { ...e, arms: e.arms.map(a => ({ ...a, body: r(a.body) })) };
+        case "forall": return { ...e, body: r(e.body) };
+        case "exists": return { ...e, body: r(e.body) };
+        case "let": return { ...e, value: r(e.value), body: r(e.body) };
+        case "methodCall": return { ...e, obj: r(e.obj), args: e.args.map(r) };
+        case "lambda": return e;
+    }
+}
+/** Map over all expressions in a statement tree. */
+function mapStmt(s, f) {
+    const r = (e) => mapExpr(e, f);
+    switch (s.kind) {
+        case "let": return { ...s, value: r(s.value) };
+        case "assign": return { ...s, value: r(s.value) };
+        case "bind": return { ...s, value: r(s.value) };
+        case "let-bind": return { ...s, value: r(s.value) };
+        case "return": return { ...s, value: r(s.value) };
+        case "break":
+        case "continue": return s;
+        case "if": return { ...s, cond: r(s.cond), then: s.then.map(t => mapStmt(t, f)), else: s.else.map(t => mapStmt(t, f)) };
+        case "match": return { ...s, arms: s.arms.map(a => ({ ...a, body: a.body.map(t => mapStmt(t, f)) })) };
+        case "while": return { ...s, cond: r(s.cond), invariants: s.invariants.map(r), body: s.body.map(t => mapStmt(t, f)) };
+        case "forin": return { ...s, bound: r(s.bound), invariants: s.invariants.map(r), body: s.body.map(t => mapStmt(t, f)) };
+        case "ghostLet": return { ...s, value: r(s.value) };
+        case "ghostAssign": return { ...s, value: r(s.value) };
+        case "assert": return { ...s, expr: r(s.expr) };
+    }
+}
+function mapStmts(stmts, f) {
+    return stmts.map(s => mapStmt(s, f));
+}
 export const LEAN_OPTIONS = {
     backend: "lean",
     monadic: true,
-    dotMethods: {
-        array: {
-            map: { pure: "map", monadic: "mapM" },
-            filter: { pure: "filter", monadic: "filterM" },
-            every: { pure: "all", monadic: "allM" },
-            some: { pure: "any", monadic: "anyM" },
-            includes: { pure: "contains" },
-            find: { pure: "find?" },
-            with: { pure: "set!" },
-        },
-    },
-    methodTable: {
-        string: {
-            indexOf: "JSString.indexOf",
-            slice: "JSString.slice",
-        },
-        array: {
-            push: "Array.push",
-        },
-    },
 };
 export const DAFNY_OPTIONS = {
     backend: "dafny",
     monadic: false,
-    dotMethods: {
-        array: {
-            map: { pure: "map" },
-            filter: { pure: "filter" },
-            every: { pure: "every" },
-            some: { pure: "some" },
-            includes: { pure: "includes" },
-            with: { pure: "with" },
-        },
-    },
-    methodTable: {
-        string: {
-            indexOf: "StringIndexOf",
-            slice: "StringSlice",
-        },
-        array: {
-            push: "SeqPush",
-        },
-    },
 };
 /** Active options — set before each transform call. */
 let _opts = LEAN_OPTIONS;
@@ -58,13 +80,10 @@ let _opts = LEAN_OPTIONS;
 function matchBinder(fieldName) {
     return `_${fieldName}`;
 }
+const _forofCounters = new Map();
 function isNat(ty) { return ty.kind === "nat"; }
 function isArray(ty) { return ty.kind === "array"; }
 function isUser(ty) { return ty.kind === "user"; }
-// ── Method lookup (uses active options) ─────────────────────
-function lookupDotMethod(recvTy, method) {
-    return _opts.dotMethods[recvTy.kind]?.[method];
-}
 /** Check if transformed lambda body contains monadic binds. */
 function isMonadicBody(stmts) {
     for (const s of stmts) {
@@ -84,32 +103,20 @@ function isMonadicBody(stmts) {
     }
     return false;
 }
-/** Lean modules that don't need explicit imports. */
-const BUILTIN_MODULES = new Set(["Array", "String", "List", "Nat", "Int"]);
-/** Map from Lean module prefix → import path. */
-const MODULE_IMPORTS = {
-    "JSString": "LemmaScript.JSString",
-};
-const usedImports = new Set();
-function lookupMethod(recvTy, method) {
-    const tyKey = recvTy.kind === "array" ? "array" : recvTy.kind;
-    const lean = _opts.methodTable[tyKey]?.[method];
-    if (lean) {
-        const mod = lean.split(".")[0];
-        if (!BUILTIN_MODULES.has(mod))
-            usedImports.add(mod);
-    }
-    return lean;
-}
 // ── Transform expressions ────────────────────────────────────
+/** Prop-valued operators (for specs/invariants). */
 const OP_MAP = {
     "===": "=", "!==": "≠", ">=": "≥", "<=": "≤", ">": ">", "<": "<",
     "&&": "∧", "||": "∨", "+": "+", "-": "-", "*": "*", "/": "/", "%": "%",
     "==": "=", "!=": "≠",
 };
+/** Bool-valued operators (for code-level conditions needing Decidable). */
+const BOOL_OP_MAP = {
+    ...OP_MAP, "===": "==", "!==": "!=",
+};
 function transformExpr(e) { return lowerExpr(e, null); }
 /**
- * Lower a typed expression to Lean IR.
+ * Lower a typed expression to Backend IR.
  *
  * When `binds` is non-null, embedded method calls are extracted into
  * `let ← ` binds (monadic lifting / selective ANF).  Lifting propagates
@@ -139,6 +146,20 @@ function lowerExpr(e, binds) {
         case "unop":
             if (e.op === "-" && e.expr.kind === "num")
                 return { kind: "num", value: -e.expr.value };
+            // String truthiness: !str → str == ""
+            if (e.op === "!" && e.expr.ty.kind === "string")
+                return { kind: "binop", op: "=", left: lowerExpr(e.expr, binds), right: { kind: "str", value: "" } };
+            // Optional truthiness: !opt → opt is None
+            if (e.op === "!" && e.expr.ty.kind === "optional") {
+                const bound = matchBinder("value");
+                return {
+                    kind: "match", scrutinee: lowerExpr(e.expr, binds),
+                    arms: [
+                        { pattern: `.some ${bound}`, body: { kind: "bool", value: false } },
+                        { pattern: ".none", body: { kind: "bool", value: true } },
+                    ],
+                };
+            }
             return { kind: "unop", op: e.op === "!" ? "¬" : e.op, expr: lowerExpr(e.expr, binds) };
         case "binop": {
             // Implication: flatten (A && B) ==> C → implies [A, B] C
@@ -166,6 +187,47 @@ function lowerExpr(e, binds) {
                     : { kind: "str", value: e.right.value };
                 return { kind: "binop", op: e.op === "===" ? "=" : "≠", left, right };
             }
+            // Optional comparison: optExpr op val → match optExpr { Some(v) => v op val, None => false/true }
+            if (["===", "!==", ">=", "<=", ">", "<"].includes(e.op) &&
+                (e.left.ty.kind === "optional") !== (e.right.ty.kind === "optional")) {
+                const [optSide, valSide] = e.left.ty.kind === "optional" ? [e.left, e.right] : [e.right, e.left];
+                const optExpr = lowerExpr(optSide, binds);
+                // x === undefined → None?, x !== undefined → Some?
+                if (valSide.kind === "var" && valSide.name === "undefined") {
+                    const isNone = e.op === "===";
+                    return {
+                        kind: "match", scrutinee: optExpr,
+                        arms: [
+                            { pattern: ".some _", body: { kind: "bool", value: !isNone } },
+                            { pattern: ".none", body: { kind: "bool", value: isNone } },
+                        ],
+                    };
+                }
+                const valExpr = lowerExpr(valSide, binds);
+                const cmpOp = BOOL_OP_MAP[e.op] ?? e.op;
+                const noneVal = e.op === "!==" ? true : false;
+                const bound = matchBinder("value");
+                return {
+                    kind: "match", scrutinee: optExpr,
+                    arms: [
+                        { pattern: `.some ${bound}`, body: { kind: "binop", op: cmpOp, left: { kind: "var", name: bound }, right: valExpr } },
+                        { pattern: ".none", body: { kind: "bool", value: noneVal } },
+                    ],
+                };
+            }
+            // || on optional → match Some/None with default
+            if (e.op === "||" && e.left.ty.kind === "optional") {
+                const optExpr = lowerExpr(e.left, binds);
+                const defaultExpr = lowerExpr(e.right, binds);
+                const bound = matchBinder("value");
+                return {
+                    kind: "match", scrutinee: optExpr,
+                    arms: [
+                        { pattern: `.some ${bound}`, body: { kind: "var", name: bound } },
+                        { pattern: ".none", body: defaultExpr },
+                    ],
+                };
+            }
             return {
                 kind: "binop",
                 op: OP_MAP[e.op] ?? e.op,
@@ -178,6 +240,8 @@ function lowerExpr(e, binds) {
                 return { kind: "field", obj: transformExpr(e.obj), field: "size" };
             if (e.field === "length" && e.obj.ty.kind === "string")
                 return { kind: "field", obj: transformExpr(e.obj), field: "length" };
+            if (e.field === "size" && (e.obj.ty.kind === "map" || e.obj.ty.kind === "set"))
+                return { kind: "field", obj: transformExpr(e.obj), field: "collectionSize" };
             return { kind: "field", obj: transformExpr(e.obj), field: e.field };
         case "index": {
             const idx = transformExpr(e.idx);
@@ -185,43 +249,47 @@ function lowerExpr(e, binds) {
             return { kind: "index", arr: transformExpr(e.obj), idx: wrappedIdx };
         }
         case "call": {
-            // Math.floor(a / b): Lean int div floors (erase), Dafny truncates (emit JSFloorDiv)
+            // Math.ceil(x): CeilReal on real args, identity on int
+            if (e.fn.kind === "field" && e.fn.field === "ceil" && e.fn.obj.kind === "var" && e.fn.obj.name === "Math" && e.args.length === 1) {
+                const arg = e.args[0];
+                if (arg.ty.kind === "real")
+                    return { kind: "app", fn: "CeilReal", args: [lowerExpr(arg, binds)] };
+                return lowerExpr(arg, binds);
+            }
+            // Math.floor(x): FloorReal on real args, JSFloorDiv for int division, identity on int
             if (e.fn.kind === "field" && e.fn.field === "floor" && e.fn.obj.kind === "var" && e.fn.obj.name === "Math" && e.args.length === 1) {
                 const arg = e.args[0];
+                if (arg.ty.kind === "real")
+                    return { kind: "app", fn: "FloorReal", args: [lowerExpr(arg, binds)] };
                 if (_opts.backend === "dafny" && arg.kind === "binop" && arg.op === "/")
                     return { kind: "app", fn: "JSFloorDiv", args: [lowerExpr(arg.left, binds), lowerExpr(arg.right, binds)] };
                 return lowerExpr(arg, binds);
             }
-            // Built-in method call: receiver.method(args)
+            // Method call: receiver.method(args) → methodCall node
             if (e.fn.kind === "field") {
-                // Remapped methods: leanFn receiver args
-                const lean = lookupMethod(e.fn.obj.ty, e.fn.field);
-                if (lean)
-                    return { kind: "app", fn: lean, args: [lowerExpr(e.fn.obj, binds), ...e.args.map(a => lowerExpr(a, binds))] };
-                // Dot-notation methods: receiver.leanName args
-                const dotEntry = lookupDotMethod(e.fn.obj.ty, e.fn.field);
-                if (dotEntry) {
-                    const recv = lowerExpr(e.fn.obj, binds);
-                    const args = e.args.map((a, i) => {
-                        const lowered = lowerExpr(a, binds);
-                        // set! index (first arg) needs .toNat when Int-typed
-                        if (dotEntry.pure === "set!" && i === 0 && !isNat(a.ty))
-                            return { kind: "toNat", expr: lowered };
-                        return lowered;
-                    });
-                    // Check if any lambda arg has monadic body → use monadic variant
-                    const needsMonadic = _opts.monadic && args.some(a => a.kind === "lambda" && isMonadicBody(a.body));
-                    const method = needsMonadic && dotEntry.monadic ? dotEntry.monadic : dotEntry.pure;
-                    const result = { kind: "dotCall", obj: recv, method, args };
-                    // Monadic HOF call is itself monadic — lift via binds like a method call
-                    if (_opts.monadic && needsMonadic && binds) {
-                        const name = `_t${_liftCounter++}`;
-                        binds.push({ kind: "let-bind", name, value: result });
-                        return { kind: "var", name };
-                    }
-                    return result;
+                const recv = lowerExpr(e.fn.obj, binds);
+                let method = e.fn.field;
+                const args = e.args.map((a, i) => {
+                    const lowered = lowerExpr(a, binds);
+                    // arr.with index (first arg) needs .toNat when Int-typed
+                    if (e.fn.kind === "field" && e.fn.field === "with" && e.fn.obj.ty.kind === "array" && i === 0 && !isNat(a.ty))
+                        return { kind: "toNat", expr: lowered };
+                    return lowered;
+                });
+                // Spec-context map get: result type is non-optional → direct access
+                if (method === "get" && e.fn.obj.ty.kind === "map" && e.ty.kind !== "optional") {
+                    method = "getDirect";
+                }
+                // Check if any lambda arg has monadic body
+                const needsMonadic = _opts.monadic && args.some(a => a.kind === "lambda" && isMonadicBody(a.body));
+                const result = { kind: "methodCall", obj: recv, objTy: e.fn.obj.ty, method, args, monadic: needsMonadic };
+                // Monadic HOF call is itself monadic — lift via binds like a method call
+                if (_opts.monadic && needsMonadic && binds) {
+                    const name = `_t${_liftCounter++}`;
+                    binds.push({ kind: "let-bind", name, value: result });
+                    return { kind: "var", name };
                 }
-                throw new Error(`Unsupported method call: .${e.fn.field}() on ${e.fn.obj.ty.kind}`);
+                return result;
             }
             if (e.fn.kind !== "var")
                 throw new Error(`Unsupported call expression: ${e.fn.kind}`);
@@ -231,15 +299,46 @@ function lowerExpr(e, binds) {
         case "record":
             return { kind: "record", spread: e.spread ? lowerExpr(e.spread, binds) : null, fields: e.fields.map(f => ({ name: f.name, value: lowerExpr(f.value, binds) })) };
         case "arrayLiteral":
+            if (e.ty.kind === "map" && e.elems.length === 0)
+                return { kind: "emptyMap" };
+            if (e.ty.kind === "set" && e.elems.length === 0)
+                return { kind: "emptySet" };
             return { kind: "arrayLiteral", elems: e.elems.map(el => lowerExpr(el, binds)) };
         case "lambda":
-            return { kind: "lambda", params: e.params.map(p => ({ name: p.name, type: tyToLean(p.ty) })), body: transformStmts(e.body, []) };
+            return { kind: "lambda", params: e.params.map(p => ({ name: p.name, type: p.ty })), body: transformStmts(e.body, []) };
         case "forall":
-            return { kind: "forall", var: e.var, type: tyToLean(e.varTy), body: transformExpr(e.body) };
+            return { kind: "forall", var: e.var, type: e.varTy, body: transformExpr(e.body) };
         case "exists":
-            return { kind: "exists", var: e.var, type: tyToLean(e.varTy), body: transformExpr(e.body) };
-        case "conditional":
-            return { kind: "if", cond: lowerExpr(e.cond, binds), then: lowerExpr(e.then, binds), else: lowerExpr(e.else, binds) };
+            return { kind: "exists", var: e.var, type: e.varTy, body: transformExpr(e.body) };
+        case "conditional": {
+            const cond = lowerExpr(e.cond, binds);
+            let thenExpr = lowerExpr(e.then, binds);
+            let elseExpr = lowerExpr(e.else, binds);
+            // Optional ternary: wrap non-undefined branch in Some, undefined branch in None
+            if (e.ty.kind === "optional") {
+                if (e.then.kind === "var" && e.then.name === "undefined") {
+                    thenExpr = { kind: "constructor", name: ".none" };
+                }
+                else {
+                    thenExpr = { kind: "app", fn: "Some", args: [thenExpr] };
+                }
+                if (e.else.kind === "var" && e.else.name === "undefined") {
+                    elseExpr = { kind: "constructor", name: ".none" };
+                }
+                else {
+                    elseExpr = { kind: "app", fn: "Some", args: [elseExpr] };
+                }
+            }
+            return { kind: "if", cond, then: thenExpr, else: elseExpr };
+        }
+        case "havoc":
+            // Dafny's * only works in var/assign positions — lift to own declaration
+            if (binds) {
+                const name = `_t${_liftCounter++}`;
+                binds.push({ kind: "let", name, type: e.ty, mutable: false, value: { kind: "havoc", type: e.ty } });
+                return { kind: "var", name };
+            }
+            return { kind: "havoc", type: e.ty };
     }
 }
 function flattenImpl(e) {
@@ -281,31 +380,17 @@ function ensuresToMatch(e, typeDecls) {
     return { kind: "match", scrutinee: obj.name, arms: [{ pattern, body: rhs }, { pattern: "_", body: { kind: "bool", value: true } }] };
 }
 function replaceFieldAccess(e, varName, fields) {
-    if (e.kind === "field" && e.obj.kind === "var" && e.obj.name === varName) {
-        const f = fields.find(f => f.name === e.field);
-        if (f)
-            return { kind: "var", name: matchBinder(f.name) };
-    }
-    const r = (x) => replaceFieldAccess(x, varName, fields);
-    switch (e.kind) {
-        case "binop": return { ...e, left: r(e.left), right: r(e.right) };
-        case "unop": return { ...e, expr: r(e.expr) };
-        case "implies": return { ...e, premises: e.premises.map(r), conclusion: r(e.conclusion) };
-        case "forall": return { ...e, body: r(e.body) };
-        case "exists": return { ...e, body: r(e.body) };
-        case "app": return { ...e, args: e.args.map(r) };
-        case "record": return { ...e, spread: e.spread ? r(e.spread) : null, fields: e.fields.map(f => ({ ...f, value: r(f.value) })) };
-        case "if": return { ...e, cond: r(e.cond), then: r(e.then), else: r(e.else) };
-        case "let":
-            // If this let shadows the matched variable, stop replacing in the body
-            if (e.name === varName)
-                return { ...e, value: r(e.value) };
-            return { ...e, value: r(e.value), body: r(e.body) };
-        case "index": return { ...e, arr: r(e.arr), idx: r(e.idx) };
-        case "field": return { ...e, obj: r(e.obj) };
-        case "match": return { ...e, arms: e.arms.map(a => ({ ...a, body: r(a.body) })) };
-        default: return e;
-    }
+    return mapExpr(e, x => {
+        if (x.kind === "field" && x.obj.kind === "var" && x.obj.name === varName) {
+            const f = fields.find(f => f.name === x.field);
+            if (f)
+                return { kind: "var", name: matchBinder(f.name) };
+        }
+        // If this let shadows the matched variable, stop replacing in the body
+        if (x.kind === "let" && x.name === varName)
+            return { ...x, value: replaceFieldAccess(x.value, varName, fields) };
+        return null;
+    });
 }
 // ── Transform statements ─────────────────────────────────────
 function transformStmts(stmts, typeDecls) {
@@ -321,20 +406,76 @@ function transformStmts(stmts, typeDecls) {
                 i += chain.consumed;
                 continue;
             }
+            // Detect optional check → match on Some/None
+            const opt = parseOptionalCheck(s.cond);
+            if (opt) {
+                const rest = stmts.slice(i + 1);
+                result.push(emitOptionalMatch(opt.varName, opt.negated, s, typeDecls, rest));
+                // If rest was consumed into the Some branch, skip remaining
+                const someBranch = opt.negated ? s.else : s.then;
+                if (someBranch.length === 0 && rest.length > 0) {
+                    return result;
+                }
+                i++;
+                continue;
+            }
         }
         // Transform for-of → for-in over range
         if (s.kind === "forof") {
-            const arrExpr = transformExpr(s.iterable);
-            const idxName = `_${s.varName}_idx`;
+            const varName = s.names[0];
+            const varTy = s.nameTypes[0] ?? { kind: "unknown" };
+            let iterExpr = transformExpr(s.iterable);
+            // Map iteration: for (const [k, v] of map) → iterate keys, look up values
+            if (s.names.length >= 2 && s.iterable.ty.kind === "map") {
+                const keyName = s.names[0], valueName = s.names[1];
+                const keyTy = s.nameTypes[0] ?? { kind: "unknown" };
+                const valueTy = s.nameTypes[1] ?? { kind: "unknown" };
+                const keysSeqName = `_${keyName}_keys`;
+                const convExpr = { kind: "app", fn: "SetToSeq", args: [{ kind: "field", obj: iterExpr, field: "keys" }] };
+                result.push({ kind: "let", name: keysSeqName, type: { kind: "array", elem: keyTy }, mutable: false, value: convExpr });
+                const keysVar = { kind: "var", name: keysSeqName };
+                const count = _forofCounters.get(keyName) ?? 0;
+                _forofCounters.set(keyName, count + 1);
+                const suffix = count === 0 ? "" : `${count + 1}`;
+                const idxName = `_${keyName}_idx${suffix}`;
+                const idx = { kind: "var", name: idxName };
+                const arrSize = { kind: "field", obj: keysVar, field: "size" };
+                const bodyStmts = transformStmts(s.body, typeDecls);
+                const letKey = { kind: "let", name: keyName, type: keyTy, mutable: false, value: { kind: "index", arr: keysVar, idx } };
+                const letVal = { kind: "let", name: valueName, type: valueTy, mutable: false,
+                    value: { kind: "methodCall", obj: iterExpr, objTy: s.iterable.ty, method: "getDirect", args: [{ kind: "var", name: keyName }], monadic: false } };
+                const boundInv = { kind: "binop", op: "≤", left: idx, right: arrSize };
+                result.push({
+                    kind: "forin", idx: idxName, bound: arrSize,
+                    invariants: [boundInv, ...s.invariants.map(transformExpr)],
+                    body: [letKey, letVal, ...bodyStmts],
+                });
+                i++;
+                continue;
+            }
+            // Sets aren't indexable — bind SetToSeq to a variable for iteration
+            if (s.iterable.ty.kind === "set") {
+                const seqName = `_${varName}_seq`;
+                const convExpr = { kind: "app", fn: "SetToSeq", args: [iterExpr] };
+                const elemTy = varTy.kind !== "unknown" ? varTy : { kind: "string" };
+                result.push({ kind: "let", name: seqName, type: { kind: "array", elem: elemTy }, mutable: false, value: convExpr });
+                iterExpr = { kind: "var", name: seqName };
+            }
+            const count = _forofCounters.get(varName) ?? 0;
+            _forofCounters.set(varName, count + 1);
+            const suffix = count === 0 ? "" : `${count + 1}`;
+            const idxName = `_${varName}_idx${suffix}`;
             const idx = { kind: "var", name: idxName };
-            const arrSize = { kind: "field", obj: arrExpr, field: "size" };
+            const arrSize = { kind: "field", obj: iterExpr, field: "size" };
             const bodyStmts = transformStmts(s.body, typeDecls);
-            const letElem = { kind: "let", name: s.varName, type: tyToLean(s.varTy), mutable: false, value: { kind: "index", arr: arrExpr, idx } };
+            const letElem = { kind: "let", name: varName, type: varTy, mutable: false, value: { kind: "index", arr: iterExpr, idx } };
+            // Auto-add bound invariant: idx ≤ bound (always true for range loops)
+            const boundInv = { kind: "binop", op: "≤", left: idx, right: arrSize };
             result.push({
                 kind: "forin",
                 idx: idxName,
                 bound: arrSize,
-                invariants: s.invariants.map(transformExpr),
+                invariants: [boundInv, ...s.invariants.map(transformExpr)],
                 body: [letElem, ...bodyStmts],
             });
             i++;
@@ -354,8 +495,71 @@ function liftMethodCalls(e) {
 function transformStmt(s, typeDecls) {
     switch (s.kind) {
         case "let": {
+            // Whole-init havoc: emit directly, no lifting
+            if (s.init.kind === "havoc") {
+                return [{ kind: "let", name: s.name, type: s.ty, mutable: s.mutable, value: { kind: "havoc", type: s.init.ty } }];
+            }
+            // arr.shift()! → let x = arr[0]; arr = arr[1..]
+            const init = s.init.kind === "call" ? s.init : undefined;
+            if (init && init.fn.kind === "field" && init.fn.field === "shift" && init.fn.obj.ty.kind === "array") {
+                const arrName = init.fn.obj.kind === "var" ? init.fn.obj.name : undefined;
+                if (arrName) {
+                    const arrVar = { kind: "var", name: arrName };
+                    const letHead = { kind: "let", name: s.name, type: s.ty, mutable: s.mutable,
+                        value: { kind: "index", arr: arrVar, idx: { kind: "num", value: 0 } } };
+                    const sliceTail = { kind: "assign", target: arrName,
+                        value: { kind: "methodCall", obj: arrVar, objTy: init.fn.obj.ty, method: "slice", args: [{ kind: "num", value: 1 }], monadic: false } };
+                    return [letHead, sliceTail];
+                }
+            }
+            // new Map(arr.map(n => [n.field, n])) → let m = map[]; for (n of arr) m[n.field] := n
+            if (init && init.fn.kind === "var" && init.fn.name === "__mapFromArray" &&
+                init.args.length === 1 && init.args[0].kind === "call" &&
+                init.args[0].fn.kind === "field" && init.args[0].fn.field === "map" &&
+                init.args[0].args.length === 1 && init.args[0].args[0].kind === "lambda") {
+                const arrExpr = init.args[0].fn.obj;
+                const lam = init.args[0].args[0];
+                const param = lam.params[0]?.name ?? "_";
+                const lamBody = Array.isArray(lam.body) ? lam.body : [{ kind: "return", value: lam.body }];
+                const retStmt = lamBody.find(b => b.kind === "return");
+                if (retStmt && retStmt.value.kind === "arrayLiteral" && retStmt.value.elems.length === 2) {
+                    const keyExpr = retStmt.value.elems[0];
+                    const valExpr = retStmt.value.elems[1];
+                    const arrIR = transformExpr(arrExpr);
+                    const arrTy = arrExpr.ty;
+                    const elemTy = arrTy.kind === "array" ? arrTy.elem : { kind: "unknown" };
+                    const idxName = `_${param}_idx`;
+                    const idx = { kind: "var", name: idxName };
+                    const arrSize = { kind: "field", obj: arrIR, field: "size" };
+                    const elemVar = { kind: "var", name: param };
+                    const keyIR = transformExpr(keyExpr);
+                    const valIR = transformExpr(valExpr);
+                    const mapSet = { kind: "methodCall", obj: { kind: "var", name: s.name }, objTy: s.ty, method: "set", args: [keyIR, valIR], monadic: false };
+                    // Auto-invariant: all processed elements' keys are in the map
+                    const kVar = { kind: "var", name: "ki" };
+                    const mapHasKey = {
+                        kind: "implies",
+                        premises: [
+                            { kind: "binop", op: "≥", left: kVar, right: { kind: "num", value: 0 } },
+                            { kind: "binop", op: "<", left: kVar, right: idx },
+                        ],
+                        conclusion: { kind: "methodCall", obj: { kind: "var", name: s.name }, objTy: s.ty, method: "has", args: [keyIR.kind === "field" ? { kind: "field", obj: { kind: "index", arr: arrIR, idx: kVar }, field: keyIR.field } : keyIR], monadic: false },
+                    };
+                    const autoInv = { kind: "forall", var: "ki", type: { kind: "int" }, body: mapHasKey };
+                    const stmts = [
+                        { kind: "let", name: s.name, type: s.ty, mutable: true, value: { kind: "emptyMap" } },
+                        { kind: "forin", idx: idxName, bound: arrSize,
+                            invariants: [{ kind: "binop", op: "≤", left: idx, right: arrSize }, autoInv],
+                            body: [
+                                { kind: "let", name: param, type: elemTy, mutable: false, value: { kind: "index", arr: arrIR, idx } },
+                                { kind: "assign", target: s.name, value: mapSet },
+                            ] },
+                    ];
+                    return stmts;
+                }
+            }
             const { binds, expr } = liftMethodCalls(s.init);
-            return [...binds, { kind: "let", name: s.name, type: tyToLean(s.ty), mutable: s.mutable, value: expr }];
+            return [...binds, { kind: "let", name: s.name, type: s.ty, mutable: s.mutable, value: expr }];
         }
         case "assign": {
             // Top-level method call → direct monadic bind, no lifting needed
@@ -371,6 +575,37 @@ function transformStmt(s, typeDecls) {
         case "break": return [{ kind: "break" }];
         case "continue": return [{ kind: "continue" }];
         case "expr": {
+            // Mutating collection call: m.set(k, v) → m := m.set(k, v)
+            // Same for s.add(x) on sets, arr.push(x)
+            if (s.expr.kind === "call" && s.expr.fn.kind === "field" &&
+                s.expr.fn.obj.kind === "var" &&
+                ((s.expr.fn.obj.ty.kind === "map" || s.expr.fn.obj.ty.kind === "set") &&
+                    (s.expr.fn.field === "set" || s.expr.fn.field === "add" || s.expr.fn.field === "delete")) ||
+                (s.expr.kind === "call" && s.expr.fn.kind === "field" &&
+                    s.expr.fn.obj.kind === "var" && s.expr.fn.obj.ty.kind === "array" &&
+                    s.expr.fn.field === "push")) {
+                const receiver = s.expr.fn.obj.name;
+                const { binds, expr } = liftMethodCalls(s.expr);
+                return [...binds, { kind: "assign", target: receiver, value: expr }];
+            }
+            // Optional chaining on map.get: m.get(k)?.push(v) → if k in m { m[k] := m[k] + [v] }
+            if (s.expr.kind === "call" && s.expr.fn.kind === "field" &&
+                s.expr.fn.obj.kind === "call" && s.expr.fn.obj.fn.kind === "field" &&
+                s.expr.fn.obj.fn.obj.ty.kind === "map" && s.expr.fn.obj.fn.field === "get" &&
+                s.expr.fn.field === "push") {
+                const mapExpr = s.expr.fn.obj.fn.obj;
+                const mapName = mapExpr.kind === "var" ? mapExpr.name : undefined;
+                const keyExpr = lowerExpr(s.expr.fn.obj.args[0], null);
+                const pushArg = lowerExpr(s.expr.args[0], null);
+                if (mapName) {
+                    const mapVar = { kind: "var", name: mapName };
+                    const directGet = { kind: "methodCall", obj: mapVar, objTy: mapExpr.ty, method: "getDirect", args: [keyExpr], monadic: false };
+                    const pushed = { kind: "methodCall", obj: directGet, objTy: mapExpr.ty.value, method: "push", args: [pushArg], monadic: false };
+                    const updated = { kind: "methodCall", obj: mapVar, objTy: mapExpr.ty, method: "set", args: [keyExpr, pushed], monadic: false };
+                    const hasCond = { kind: "methodCall", obj: mapVar, objTy: mapExpr.ty, method: "has", args: [keyExpr], monadic: false };
+                    return [{ kind: "if", cond: hasCond, then: [{ kind: "assign", target: mapName, value: updated }], else: [] }];
+                }
+            }
             const { binds, expr } = liftMethodCalls(s.expr);
             return [...binds, { kind: "assign", target: "_", value: expr }];
         }
@@ -388,10 +623,18 @@ function transformStmt(s, typeDecls) {
                     doneWith: s.doneWith ? transformExpr(s.doneWith) : null,
                     body: transformStmts(s.body, typeDecls),
                 }];
+        case "throw":
+            return [{ kind: "assert", expr: { kind: "bool", value: false } }];
         case "forof":
             throw new Error("forof should be transformed to forin (range loop) in transformStmts");
         case "switch":
             return [emitSwitchStmt(s, typeDecls)];
+        case "ghostLet":
+            return [{ kind: "ghostLet", name: s.name, type: s.ty, value: transformExpr(s.init) }];
+        case "ghostAssign":
+            return [{ kind: "ghostAssign", target: s.target, value: transformExpr(s.value) }];
+        case "assert":
+            return [{ kind: "assert", expr: transformExpr(s.expr) }];
     }
 }
 function detectDiscriminantChain(stmts) {
@@ -443,6 +686,41 @@ function parseDiscriminantCond(cond) {
         return null;
     return { varName: cond.left.obj.name, typeName: cond.left.obj.ty.name, variant: cond.right.value };
 }
+function emitOptionalMatch(varName, negated, s, typeDecls, restStmts) {
+    let someBranch = negated ? s.else : s.then;
+    const noneBranch = negated ? s.then : s.else;
+    // Early-return pattern: if (x === undefined) { return ... } — Some branch is empty,
+    // so include remaining statements as the Some body
+    if (someBranch.length === 0 && restStmts && restStmts.length > 0) {
+        someBranch = restStmts;
+    }
+    const bound = matchBinder(`${varName}_val`);
+    const someBody = transformStmts(someBranch, typeDecls);
+    const r = (e) => replaceVar(e, varName, { kind: "var", name: bound });
+    const someReplaced = someBody.map(stmt => mapStmtExprs(stmt, r));
+    const arms = [
+        { pattern: `.some ${bound}`, body: someReplaced },
+        { pattern: ".none", body: noneBranch.length > 0 ? transformStmts(noneBranch, typeDecls) : [] },
+    ];
+    return { kind: "match", scrutinee: varName, arms };
+}
+/** Apply an expression transform to all expressions in a statement (convenience wrapper). */
+function mapStmtExprs(s, r) {
+    return mapStmt(s, e => r(e));
+}
+/** Detect `v !== undefined` or `undefined !== v` where v has optional type. */
+function parseOptionalCheck(cond) {
+    if (cond.kind !== "binop" || (cond.op !== "!==" && cond.op !== "==="))
+        return null;
+    let varExpr = null;
+    if (cond.right.kind === "var" && cond.right.name === "undefined")
+        varExpr = cond.left;
+    if (cond.left.kind === "var" && cond.left.name === "undefined")
+        varExpr = cond.right;
+    if (!varExpr || varExpr.kind !== "var" || varExpr.ty.kind !== "optional")
+        return null;
+    return { varName: varExpr.name, negated: cond.op === "===" };
+}
 function emitMatchStmt(chain, typeDecls) {
     const decl = typeDecls.find(d => d.name === chain.typeName);
     const arms = chain.cases.map(c => {
@@ -476,36 +754,26 @@ function emitSwitchStmt(s, typeDecls) {
 function replaceFieldAccessInStmts(stmts, varName, fields) {
     if (fields.length === 0)
         return stmts;
+    const f = (e) => {
+        if (e.kind === "field" && e.obj.kind === "var" && e.obj.name === varName) {
+            const fi = fields.find(fi => fi.name === e.field);
+            if (fi)
+                return { kind: "var", name: matchBinder(fi.name) };
+        }
+        return null;
+    };
     const result = [];
     for (const s of stmts) {
         // If a let shadows the matched variable, stop replacing from here on
         if (s.kind === "let" && s.name === varName) {
-            const r = (e) => replaceFieldAccess(e, varName, fields);
-            result.push({ ...s, value: r(s.value) });
-            // Remaining statements see the shadowed name — no more replacement
+            result.push(s.value ? { ...s, value: mapExpr(s.value, f) } : s);
             result.push(...stmts.slice(result.length));
             break;
         }
-        result.push(replaceFieldAccessInStmt(s, varName, fields));
+        result.push(mapStmt(s, f));
     }
     return result;
 }
-function replaceFieldAccessInStmt(s, varName, fields) {
-    const r = (e) => replaceFieldAccess(e, varName, fields);
-    switch (s.kind) {
-        case "let": return { ...s, value: r(s.value) };
-        case "assign": return { ...s, value: r(s.value) };
-        case "bind": return { ...s, value: r(s.value) };
-        case "let-bind": return { ...s, value: r(s.value) };
-        case "return": return { ...s, value: r(s.value) };
-        case "break":
-        case "continue": return s;
-        case "if": return { ...s, cond: r(s.cond), then: replaceFieldAccessInStmts(s.then, varName, fields), else: replaceFieldAccessInStmts(s.else, varName, fields) };
-        case "match": return { ...s, arms: s.arms.map(a => ({ ...a, body: replaceFieldAccessInStmts(a.body, varName, fields) })) };
-        case "while": return { ...s, cond: r(s.cond), body: replaceFieldAccessInStmts(s.body, varName, fields) };
-        case "forin": return { ...s, invariants: s.invariants.map(r), body: replaceFieldAccessInStmts(s.body, varName, fields) };
-    }
-}
 // ── Pure function generation ─────────────────────────────────
 function transformPureBody(stmts, typeDecls) {
     // Detect discriminant if-chain
@@ -526,6 +794,29 @@ function transformPureBody(stmts, typeDecls) {
                 return { kind: "let", name: s.name, value: transformExpr(s.init), body: restExpr };
             }
             case "if": {
+                // Optional narrowing: if (x === undefined) → match x { None => ..., Some(x_val) => ... }
+                const optCheck = parseOptionalCheck(s.cond);
+                if (optCheck) {
+                    let someBranch = optCheck.negated ? s.else : s.then;
+                    const noneBranch = optCheck.negated ? s.then : (s.else.length > 0 ? s.else : rest);
+                    if (someBranch.length === 0)
+                        someBranch = rest;
+                    const bound = matchBinder(`${optCheck.varName}_val`);
+                    const someExpr = transformPureBody(someBranch, typeDecls);
+                    if (!someExpr)
+                        return null;
+                    const noneExpr = transformPureBody(noneBranch, typeDecls);
+                    if (!noneExpr)
+                        return null;
+                    const someReplaced = replaceVar(someExpr, optCheck.varName, { kind: "var", name: bound });
+                    return {
+                        kind: "match", scrutinee: optCheck.varName,
+                        arms: [
+                            { pattern: `.some ${bound}`, body: someReplaced },
+                            { pattern: ".none", body: noneExpr },
+                        ],
+                    };
+                }
                 const thenExpr = transformPureBody(s.then, typeDecls);
                 if (!thenExpr)
                     return null;
@@ -607,7 +898,7 @@ function transformTypeDecl(d) {
             kind: "inductive", name: d.name,
             constructors: d.variants.map(v => ({
                 name: v.name,
-                fields: v.fields.map(f => ({ name: f.name, type: tyToLean(parseTsType(f.tsType)) })),
+                fields: v.fields.map(f => ({ name: f.name, type: parseTsType(f.tsType) })),
             })),
             deriving: ["Repr", "Inhabited"],
         };
@@ -615,38 +906,59 @@ function transformTypeDecl(d) {
     else {
         return {
             kind: "structure", name: d.name,
-            fields: d.fields.map(f => ({ name: f.name, type: tyToLean(parseTsType(f.tsType)) })),
+            fields: d.fields.map(f => ({ name: f.name, type: parseTsType(f.tsType) })),
             deriving: ["Repr", "Inhabited", "DecidableEq"],
         };
     }
 }
 // ── Helpers ──────────────────────────────────────────────────
+/** Find parameter names that are reassigned anywhere in the body. */
+function findReassignedNames(stmts, names) {
+    const found = new Set();
+    function scan(stmts) {
+        for (const s of stmts) {
+            if (s.kind === "assign" && names.has(s.target))
+                found.add(s.target);
+            if (s.kind === "ghostAssign" && names.has(s.target))
+                found.add(s.target);
+            // 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)) {
+                found.add(s.expr.fn.obj.name);
+            }
+            if (s.kind === "if") {
+                scan(s.then);
+                scan(s.else);
+            }
+            if (s.kind === "while")
+                scan(s.body);
+            if (s.kind === "forof")
+                scan(s.body);
+            if (s.kind === "switch") {
+                for (const c of s.cases)
+                    scan(c.body);
+                scan(s.defaultBody);
+            }
+        }
+    }
+    scan(stmts);
+    return found;
+}
 /** Replace all occurrences of a variable name with a new expression. */
 function replaceVar(e, name, replacement) {
-    const r = (x) => replaceVar(x, name, replacement);
-    switch (e.kind) {
-        case "var": return e.name === name ? replacement : e;
-        case "num":
-        case "bool":
-        case "str":
-        case "constructor": return e;
-        case "binop": return { ...e, left: r(e.left), right: r(e.right) };
-        case "unop": return { ...e, expr: r(e.expr) };
-        case "implies": return { ...e, premises: e.premises.map(r), conclusion: r(e.conclusion) };
-        case "app": return { ...e, args: e.args.map(r) };
-        case "field": return { ...e, obj: r(e.obj) };
-        case "toNat": return { ...e, expr: r(e.expr) };
-        case "index": return { ...e, arr: r(e.arr), idx: r(e.idx) };
-        case "record": return { ...e, spread: e.spread ? r(e.spread) : null, fields: e.fields.map(f => ({ ...f, value: r(f.value) })) };
-        case "arrayLiteral": return { ...e, elems: e.elems.map(r) };
-        case "if": return { ...e, cond: r(e.cond), then: r(e.then), else: r(e.else) };
-        case "match": return { ...e, arms: e.arms.map(a => ({ ...a, body: r(a.body) })) };
-        case "forall": return { ...e, body: e.var === name ? e : { ...e, body: r(e.body) } };
-        case "exists": return { ...e, body: e.var === name ? e : { ...e, body: r(e.body) } };
-        case "let": return { ...e, value: r(e.value), body: e.name === name ? e : { ...e, body: r(e.body) } };
-        case "dotCall": return { ...e, obj: r(e.obj), args: e.args.map(r) };
-        case "lambda": return e; // don't descend into lambdas
-    }
+    return mapExpr(e, x => {
+        if (x.kind === "var" && x.name === name)
+            return replacement;
+        // Don't descend past bindings that shadow the name
+        if (x.kind === "forall" && x.var === name)
+            return x;
+        if (x.kind === "exists" && x.var === name)
+            return x;
+        if (x.kind === "let" && x.name === name)
+            return { ...x, value: replaceVar(x.value, name, replacement) };
+        return null;
+    });
 }
 // ── Top-level transform ──────────────────────────────────────
 /** Transform for Dafny backend — same logic, Dafny options. */
@@ -661,7 +973,15 @@ export function transformModuleDafny(mod) {
     }
 }
 export function transformModule(mod, specImport) {
+    _forofCounters.clear();
     const typeDecls = mod.typeDecls.map(transformTypeDecl);
+    // Module-level constants
+    const constDecls = (mod.constants ?? []).map(c => ({
+        kind: "const",
+        name: c.name,
+        type: c.ty,
+        value: transformExpr(c.value),
+    }));
     // Pure function mirrors
     const pureDefs = [];
     for (const fn of mod.functions) {
@@ -676,8 +996,8 @@ export function transformModule(mod, specImport) {
         pureDefs.push({
             kind: "def",
             name: fn.name,
-            params: fn.params.map(p => ({ name: p.name, type: tyToLean(p.ty) })),
-            returnType: tyToLean(fn.returnTy),
+            params: fn.params.map(p => ({ name: p.name, type: p.ty })),
+            returnType: fn.returnTy,
             requires: fn.requires.map(transformExpr),
             ensures,
             body,
@@ -686,9 +1006,6 @@ export function transformModule(mod, specImport) {
     const base = mod.file.split("/").pop()?.replace(/\.ts$/, "") ?? "module";
     // Types file
     const typesImports = ["LemmaScript"];
-    for (const m of usedImports)
-        typesImports.push(MODULE_IMPORTS[m] ?? m);
-    usedImports.clear();
     let typesFile = null;
     const pureNamespace = pureDefs.length > 0
         ? [{ kind: "namespace", name: "Pure", decls: pureDefs }]
@@ -713,25 +1030,54 @@ export function transformModule(mod, specImport) {
             else
                 ensures.push(transformExpr(e));
         }
-        const body = pureDefNames.has(fn.name)
+        _forofCounters.clear();
+        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 })) } }]
             : transformStmts(fn.body, mod.typeDecls);
+        // Shadow reassigned parameters with mutable locals
+        const paramNames = new Set(fn.params.map(p => p.name));
+        const reassigned = findReassignedNames(fn.body, paramNames);
+        if (reassigned.size > 0) {
+            const shadows = fn.params
+                .filter(p => reassigned.has(p.name))
+                .map(p => ({ kind: "let", name: p.name, type: p.ty, mutable: true, value: { kind: "var", name: p.name } }));
+            body = [...shadows, ...body];
+        }
         return {
             kind: "method",
             name: fn.name,
-            params: fn.params.map(p => ({ name: p.name, type: tyToLean(p.ty) })),
-            returnType: tyToLean(fn.returnTy),
+            params: fn.params.map(p => ({ name: p.name, type: p.ty })),
+            returnType: fn.returnTy,
             requires: fn.requires.map(transformExpr),
             ensures,
             body,
         };
     });
+    // Class declarations
+    const classDecls = (mod.classes ?? []).map(cls => {
+        const classMethods = cls.methods.map(fn => {
+            const ensures = fn.ensures.map(transformExpr);
+            _forofCounters.clear();
+            const body = transformStmts(fn.body, mod.typeDecls);
+            return {
+                kind: "method",
+                name: fn.name,
+                params: fn.params.map(p => ({ name: p.name, type: p.ty })),
+                returnType: fn.returnTy,
+                requires: fn.requires.map(transformExpr),
+                ensures,
+                body,
+            };
+        });
+        return {
+            kind: "class",
+            name: cls.name,
+            fields: cls.fields.map(f => ({ name: f.name, type: f.ty })),
+            methods: classMethods,
+        };
+    });
     const defImport = specImport ?? (typesFile ? `«${base}.types»` : null);
     const defBaseImports = defImport ? [defImport] : ["LemmaScript"];
-    if (!typesFile)
-        for (const m of usedImports)
-            defBaseImports.push(MODULE_IMPORTS[m] ?? m);
-    usedImports.clear();
     const defFile = {
         comment: "  Generated by lsc from " + (mod.file.split("/").pop() ?? "") + "\n  Do not edit — re-run `lsc gen` to regenerate.",
         imports: defBaseImports,
@@ -739,7 +1085,7 @@ export function transformModule(mod, specImport) {
             { key: "loom.semantics.termination", value: '"total"' },
             { key: "loom.semantics.choice", value: '"demonic"' },
         ],
-        decls: methods,
+        decls: [...constDecls, ...methods, ...classDecls],
     };
     return { typesFile, defFile };
 }