lemmascript 0.1.0 → 0.3.0

package/tools/dist/transform.js

@@ -1,70 +1,142 @@
 /**
- * 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": {
+            const scr = typeof e.scrutinee === "string" ? e.scrutinee : r(e.scrutinee);
+            return { ...e, scrutinee: scr, 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));
+}
+/** Map over all sub-expressions in a TExpr (typed IR). */
+function mapTExpr(e, f) {
+    const hit = f(e);
+    if (hit)
+        return hit;
+    const r = (x) => mapTExpr(x, f);
+    switch (e.kind) {
+        case "var":
+        case "num":
+        case "str":
+        case "bool":
+        case "result":
+        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 "call": return { ...e, fn: r(e.fn), args: e.args.map(r) };
+        case "index": return { ...e, obj: r(e.obj), idx: r(e.idx) };
+        case "field": return { ...e, obj: r(e.obj) };
+        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 "conditional": return { ...e, cond: r(e.cond), then: r(e.then), else: r(e.else) };
+        case "forall": return { ...e, body: r(e.body) };
+        case "exists": return { ...e, body: r(e.body) };
+        case "lambda": return e;
+    }
+}
+/** Map over all expressions in a TStmt tree (typed IR). */
+function mapTStmt(s, f) {
+    const r = (e) => mapTExpr(e, f);
+    switch (s.kind) {
+        case "let": return { ...s, init: r(s.init) };
+        case "assign": return { ...s, value: r(s.value) };
+        case "return": return { ...s, value: r(s.value) };
+        case "break":
+        case "continue":
+        case "throw": return s;
+        case "expr": return { ...s, expr: r(s.expr) };
+        case "if": return { ...s, cond: r(s.cond), then: s.then.map(t => mapTStmt(t, f)), else: s.else.map(t => mapTStmt(t, f)) };
+        case "while": return { ...s, cond: r(s.cond), invariants: s.invariants.map(r), body: s.body.map(t => mapTStmt(t, f)) };
+        case "switch": return { ...s, expr: r(s.expr), cases: s.cases.map(c => ({ ...c, body: c.body.map(t => mapTStmt(t, f)) })), defaultBody: s.defaultBody.map(t => mapTStmt(t, f)) };
+        case "forof": return { ...s, iterable: r(s.iterable), invariants: s.invariants.map(r), body: s.body.map(t => mapTStmt(t, f)) };
+        case "ghostLet": return { ...s, init: r(s.init) };
+        case "ghostAssign": return { ...s, value: r(s.value) };
+        case "assert": return { ...s, expr: r(s.expr) };
+    }
+}
 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;
-/** Prefix match-bound field names to avoid capturing user variables. */
-function matchBinder(fieldName) {
-    return `_${fieldName}`;
+let _opts = DAFNY_OPTIONS;
+/** Type declarations — set once per module transform for discriminated union handling. */
+let _typeDecls = [];
+/** Prefix match-bound field names to avoid capturing user variables.
+ *  When prefix is given (the scrutinee name), include it to avoid
+ *  collisions in nested matches on different variables. */
+function matchBinder(fieldName, prefix) {
+    return prefix ? `_${prefix}_${fieldName}` : `_${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 +156,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 +199,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 +240,89 @@ 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 } },
+                    ],
+                };
+            }
+            // || undefined on optional → identity (no-op: x || undefined = x)
+            if (e.op === "||" && e.left.ty.kind === "optional" &&
+                e.right.kind === "var" && e.right.name === "undefined") {
+                return lowerExpr(e.left, binds);
+            }
+            // || 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 },
+                    ],
+                };
+            }
+            // || on map index → if key in map then map[key] else default
+            if (e.op === "||" && e.left.kind === "index" && e.left.obj.ty.kind === "map") {
+                const map = lowerExpr(e.left.obj, binds);
+                const key = lowerExpr(e.left.idx, binds);
+                const right = lowerExpr(e.right, binds);
+                return {
+                    kind: "if",
+                    cond: { kind: "binop", op: "in", left: key, right: map },
+                    then: { kind: "index", arr: map, idx: key }, else: right,
+                };
+            }
+            // || on non-optional string/array/user → if non-empty then x else default
+            if (e.op === "||" && (e.left.ty.kind === "string" || e.left.ty.kind === "array" ||
+                (e.left.ty.kind === "user" && e.right.ty.kind === "string"))) {
+                const left = lowerExpr(e.left, binds);
+                const right = lowerExpr(e.right, binds);
+                return {
+                    kind: "if",
+                    cond: { kind: "binop", op: ">", left: { kind: "field", obj: left, field: "size" }, right: { kind: "num", value: 0 } },
+                    then: left, else: right,
+                };
+            }
+            // int + string → NatToString(int) + string (string concatenation)
+            if (e.op === "+" && _opts.backend === "dafny") {
+                const isIntL = e.left.ty.kind === "int" || e.left.ty.kind === "nat";
+                const isIntR = e.right.ty.kind === "int" || e.right.ty.kind === "nat";
+                if (isIntL && e.right.ty.kind === "string") {
+                    return { kind: "binop", op: "+",
+                        left: { kind: "app", fn: "NatToString", args: [lowerExpr(e.left, binds)] },
+                        right: lowerExpr(e.right, binds) };
+                }
+                if (e.left.ty.kind === "string" && isIntR) {
+                    return { kind: "binop", op: "+",
+                        left: lowerExpr(e.left, binds),
+                        right: { kind: "app", fn: "NatToString", args: [lowerExpr(e.right, binds)] } };
+                }
+            }
             return {
                 kind: "binop",
                 op: OP_MAP[e.op] ?? e.op,
@@ -178,6 +335,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,61 +344,198 @@ 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.abs/min/max → preamble functions
+            if (e.fn.kind === "field" && e.fn.obj.kind === "var" && e.fn.obj.name === "Math") {
+                if (e.fn.field === "abs" && e.args.length === 1)
+                    return { kind: "app", fn: "MathAbs", args: [lowerExpr(e.args[0], binds)] };
+                if (e.fn.field === "min" && e.args.length === 2)
+                    return { kind: "app", fn: "MathMin", args: e.args.map(a => lowerExpr(a, binds)) };
+                if (e.fn.field === "max" && e.args.length === 2)
+                    return { kind: "app", fn: "MathMax", args: e.args.map(a => lowerExpr(a, binds)) };
+            }
+            // 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}`);
             const prefix = e.callKind === "spec-pure" && _opts.backend === "lean" ? "Pure." : "";
             return { kind: "app", fn: prefix + e.fn.name, args: e.args.map(a => lowerExpr(a, 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 "record": {
+            // Discriminated union: { kind: 'NoOp' } → constructor NoOp
+            if (e.ty.kind === "user" && !e.spread) {
+                const tyName = e.ty.name;
+                // Match base type name (strip generic args: "Result<Model, Err>" → "Result")
+                const baseName = tyName.includes("<") ? tyName.slice(0, tyName.indexOf("<")) : tyName;
+                const decl = _typeDecls.find(d => d.name === baseName && (d.kind === "discriminated-union" || d.kind === "string-union"));
+                if (decl && decl.discriminant) {
+                    const discField = e.fields.find(f => f.name === decl.discriminant);
+                    if (discField && (discField.value.kind === "str" || discField.value.kind === "bool")) {
+                        const variantName = String(discField.value.kind === "str" ? discField.value.value : discField.value.value);
+                        const variant = decl.variants?.find(v => v.name === variantName);
+                        if (variant) {
+                            const nonDiscFields = e.fields.filter(f => f.name !== decl.discriminant);
+                            if (nonDiscFields.length === 0) {
+                                return { kind: "constructor", name: variantName, type: tyName };
+                            }
+                            // Constructor with args: match variant field order
+                            const args = variant.fields.map(vf => {
+                                const ef = nonDiscFields.find(f => f.name === vf.name);
+                                return ef ? lowerExpr(ef.value, binds) : { kind: "var", name: "None" };
+                            });
+                            return { kind: "app", fn: variantName, args };
+                        }
+                    }
+                }
+            }
+            // For spread records, wrap non-optional values in Some for optional fields
+            if (e.spread) {
+                const spreadTy = e.spread.ty.kind === "optional" ? e.spread.ty.inner : e.spread.ty;
+                const structName = spreadTy.kind === "user" ? spreadTy.name : undefined;
+                const structDecl = structName ? _typeDecls.find(d => d.name === structName && d.kind === "record") : undefined;
+                const loweredFields = e.fields.map(f => {
+                    let value = lowerExpr(f.value, binds);
+                    if (structDecl?.fields) {
+                        const fieldDecl = structDecl.fields.find(sf => sf.name === f.name);
+                        if (fieldDecl) {
+                            const fieldTy = parseTsType(fieldDecl.tsType);
+                            const isUndef = f.value.kind === "var" && f.value.name === "undefined";
+                            if (fieldTy.kind === "optional" && f.value.ty.kind !== "optional" && !isUndef) {
+                                value = { kind: "app", fn: "Some", args: [value] };
+                            }
+                        }
+                    }
+                    return { name: f.name, value };
+                });
+                return { kind: "record", spread: lowerExpr(e.spread, binds), fields: loweredFields };
+            }
+            return { kind: "record", spread: 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" };
+            // Set with initial elements: new Set([a, b]) → {a, b}
+            if (e.ty.kind === "set")
+                return { kind: "app", fn: "SetLiteral", args: e.elems.map(el => lowerExpr(el, binds)) };
             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);
+            // Explicit !== undefined with narrowedExpr → match Some/None on the optional expression
+            if (e.narrowedVar && e.narrowedExpr) {
+                const scrutinee = lowerExpr(e.narrowedExpr, binds);
+                const bound = matchBinder(e.narrowedVar);
+                if (bound !== e.narrowedVar) {
+                    thenExpr = replaceVar(thenExpr, e.narrowedVar, { kind: "var", name: bound });
+                }
+                const wrapSomeNone = (expr, raw) => (raw.kind === "var" && raw.name === "undefined")
+                    ? { kind: "constructor", name: ".none" }
+                    : { kind: "app", fn: "Some", args: [expr] };
+                thenExpr = wrapSomeNone(thenExpr, e.then);
+                elseExpr = wrapSomeNone(elseExpr, e.else);
+                return {
+                    kind: "match", scrutinee,
+                    arms: [
+                        { pattern: `.some ${bound}`, body: thenExpr },
+                        { pattern: ".none", body: elseExpr },
+                    ],
+                };
+            }
+            // Optional cond with narrowedVar → match Some/None (truthiness)
+            if (e.narrowedVar && e.cond.ty.kind === "optional") {
+                const bound = matchBinder(e.narrowedVar);
+                // Replace the synthetic/narrowed var with the match-bound name
+                if (bound !== e.narrowedVar) {
+                    thenExpr = replaceVar(thenExpr, e.narrowedVar, { kind: "var", name: bound });
+                }
+                // The match produces an Optional: wrap branches in Some/None.
+                // Either branch being undefined signals None; otherwise wrap in Some.
+                const wrapSomeNone = (expr, raw) => (raw.kind === "var" && raw.name === "undefined")
+                    ? { kind: "constructor", name: ".none" }
+                    : { kind: "app", fn: "Some", args: [expr] };
+                thenExpr = wrapSomeNone(thenExpr, e.then);
+                elseExpr = wrapSomeNone(elseExpr, e.else);
+                return {
+                    kind: "match", scrutinee: cond,
+                    arms: [
+                        { pattern: `.some ${bound}`, body: thenExpr },
+                        { pattern: ".none", body: elseExpr },
+                    ],
+                };
+            }
+            // Non-optional: regular if with optional wrapping
+            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) {
@@ -275,37 +571,23 @@ function ensuresToMatch(e, typeDecls) {
     if (!variant)
         return null;
     const fields = variant.fields;
-    const pattern = fields.length > 0 ? `.${variantName} ${fields.map(f => matchBinder(f.name)).join(" ")}` : `.${variantName}`;
+    const pattern = fields.length > 0 ? `.${variantName} ${fields.map(f => matchBinder(f.name, obj.name)).join(" ")}` : `.${variantName}`;
     let rhs = transformExpr(e.right);
     rhs = replaceFieldAccess(rhs, obj.name, fields);
     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, varName) };
+        }
+        // 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 +603,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 +692,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,10 +772,52 @@ 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 }];
         }
         case "if": {
+            // Restructure && with optional check: extract the leftmost optional check
+            // from a && chain and nest the rest inside. Handles left-associative chains:
+            // if ((x !== undefined && b) && c) → if (x !== undefined) { if (b && c) { ... } }
+            if (s.cond.kind === "binop" && s.cond.op === "&&" && s.else.length === 0) {
+                const extracted = extractLeftmostOptional(s.cond);
+                if (extracted) {
+                    const innerIf = { kind: "if", cond: extracted.rest, then: s.then, else: [] };
+                    const outerIf = { kind: "if", cond: extracted.optCond, then: [innerIf], else: [] };
+                    return transformStmts([outerIf], typeDecls);
+                }
+            }
             // Lift from condition only (Lean rule: don't lift from branches)
             const { binds, expr: cond } = liftMethodCalls(s.cond);
             return [...binds, { kind: "if", cond, then: transformStmts(s.then, typeDecls), else: transformStmts(s.else, typeDecls) }];
@@ -388,10 +831,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,14 +894,65 @@ 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));
+}
+/** Extract the leftmost optional check from a && chain, returning the check and the rest.
+ *  (x !== undefined && b) && c → { optCond: x !== undefined, rest: b && c } */
+function extractLeftmostOptional(cond) {
+    if (cond.kind !== "binop" || cond.op !== "&&")
+        return null;
+    const check = parseOptionalCheck(cond.left);
+    if (check && !check.negated)
+        return { optCond: cond.left, rest: cond.right };
+    if (cond.left.kind === "binop" && cond.left.op === "&&") {
+        const inner = extractLeftmostOptional(cond.left);
+        if (inner)
+            return { optCond: inner.optCond, rest: { ...cond, left: inner.rest } };
+    }
+    return null;
+}
+/** 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 => {
         const variant = decl?.variants?.find(v => v.name === c.variant);
         const fields = variant?.fields ?? [];
-        const pattern = fields.length > 0 ? `.${c.variant} ${fields.map(f => matchBinder(f.name)).join(" ")}` : `.${c.variant}`;
-        let body = transformStmts(c.body, typeDecls);
-        body = replaceFieldAccessInStmts(body, chain.varName, fields);
+        const pattern = fields.length > 0 ? `.${c.variant} ${fields.map(f => matchBinder(f.name, chain.varName)).join(" ")}` : `.${c.variant}`;
+        // Replace field accesses in TStmt BEFORE transforming, so optional narrowing sees simple vars
+        const replaced = replaceFieldAccessInTStmts(c.body, chain.varName, fields);
+        const body = transformStmts(replaced, typeDecls);
         return { pattern, body };
     });
     if (chain.fallthrough.length > 0)
@@ -464,48 +966,56 @@ function emitSwitchStmt(s, typeDecls) {
     const arms = s.cases.map(c => {
         const variant = decl?.variants?.find(v => v.name === c.label);
         const fields = variant?.fields ?? [];
-        const pattern = fields.length > 0 ? `.${c.label} ${fields.map(f => matchBinder(f.name)).join(" ")}` : `.${c.label}`;
-        let body = transformStmts(c.body, typeDecls);
-        body = replaceFieldAccessInStmts(body, varName, fields);
+        const pattern = fields.length > 0 ? `.${c.label} ${fields.map(f => matchBinder(f.name, varName)).join(" ")}` : `.${c.label}`;
+        // Replace field accesses in TStmt BEFORE transforming, so optional narrowing sees simple vars
+        const replaced = replaceFieldAccessInTStmts(c.body, varName, fields);
+        const body = transformStmts(replaced, typeDecls);
         return { pattern, body };
     });
     if (s.defaultBody.length > 0)
         arms.push({ pattern: "_", body: transformStmts(s.defaultBody, typeDecls) });
     return { kind: "match", scrutinee: varName, arms };
 }
+/** Replace obj.field → binder var in typed IR (before transform).
+ *  Uses the variant's declared field type since the resolve phase may not
+ *  resolve field types on discriminated unions correctly. */
+function replaceFieldAccessInTStmts(stmts, varName, fields) {
+    if (fields.length === 0)
+        return stmts;
+    return stmts.map(s => mapTStmt(s, e => {
+        if (e.kind === "field" && e.obj.kind === "var" && e.obj.name === varName) {
+            const fi = fields.find(fi => fi.name === e.field);
+            if (fi) {
+                const ty = e.ty.kind !== "unknown" ? e.ty : parseTsType(fi.tsType);
+                return { kind: "var", name: matchBinder(fi.name, varName), ty };
+            }
+        }
+        return null;
+    }));
+}
 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, varName) };
+        }
+        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 +1036,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;
@@ -545,11 +1078,12 @@ function transformPureSwitch(s, typeDecls) {
     const decl = typeDecls.find(d => d.name === (s.expr.ty.kind === "user" ? s.expr.ty.name : ""));
     if (!decl)
         return null;
+    const varName = s.expr.kind === "var" ? s.expr.name : undefined;
     const arms = [];
     for (const c of s.cases) {
         const variant = decl.variants?.find(v => v.name === c.label);
         const fields = variant?.fields ?? [];
-        const pattern = fields.length > 0 ? `.${c.label} ${fields.map(f => matchBinder(f.name)).join(" ")}` : `.${c.label}`;
+        const pattern = fields.length > 0 ? `.${c.label} ${fields.map(f => matchBinder(f.name, varName)).join(" ")}` : `.${c.label}`;
         let body = transformPureBody(c.body, typeDecls);
         if (!body)
             return null;
@@ -573,7 +1107,7 @@ function transformPureMatch(chain, typeDecls) {
     for (const c of chain.cases) {
         const variant = decl?.variants?.find(v => v.name === c.variant);
         const fields = variant?.fields ?? [];
-        const pattern = fields.length > 0 ? `.${c.variant} ${fields.map(f => matchBinder(f.name)).join(" ")}` : `.${c.variant}`;
+        const pattern = fields.length > 0 ? `.${c.variant} ${fields.map(f => matchBinder(f.name, chain.varName)).join(" ")}` : `.${c.variant}`;
         let body = transformPureBody(c.body, typeDecls);
         if (!body)
             return null;
@@ -605,50 +1139,89 @@ function transformTypeDecl(d) {
     else if (d.kind === "discriminated-union") {
         return {
             kind: "inductive", name: d.name,
+            typeParams: d.typeParams,
             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"],
         };
     }
+    else if (d.kind === "alias") {
+        return {
+            kind: "type-alias", name: d.name,
+            target: parseTsType(d.aliasOf),
+        };
+    }
     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 Lean backend — same logic, Lean options. */
+export function transformModuleLean(mod, specImport) {
+    const prev = _opts;
+    _opts = LEAN_OPTIONS;
+    try {
+        return transformModule(mod, specImport);
+    }
+    finally {
+        _opts = prev;
+    }
+}
 /** Transform for Dafny backend — same logic, Dafny options. */
 export function transformModuleDafny(mod) {
     const prev = _opts;
@@ -661,7 +1234,16 @@ export function transformModuleDafny(mod) {
     }
 }
 export function transformModule(mod, specImport) {
+    _forofCounters.clear();
+    _typeDecls = mod.typeDecls;
     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 +1258,9 @@ 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),
+            typeParams: fn.typeParams,
+            params: fn.params.map(p => ({ name: p.name, type: p.ty })),
+            returnType: fn.returnTy,
             requires: fn.requires.map(transformExpr),
             ensures,
             body,
@@ -686,9 +1269,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 +1293,56 @@ 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),
+            typeParams: fn.typeParams,
+            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,
+                typeParams: fn.typeParams,
+                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 +1350,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 };
 }