lemmascript 0.2.0 → 0.3.1

package/tools/dist/transform.js

@@ -35,7 +35,10 @@ function mapExpr(e, f) {
         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 "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) };
@@ -63,8 +66,51 @@ function mapStmt(s, f) {
         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",
@@ -75,10 +121,20 @@ export const DAFNY_OPTIONS = {
     monadic: false,
 };
 /** 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}`;
+}
+/** Build a match arm pattern like `.VariantName _v_field1 _v_field2` from variant info. */
+function buildMatchPattern(variantName, fields, scopePrefix) {
+    if (fields.length === 0)
+        return `.${variantName}`;
+    return `.${variantName} ${fields.map(f => matchBinder(f.name, scopePrefix)).join(" ")}`;
 }
 const _forofCounters = new Map();
 function isNat(ty) { return ty.kind === "nat"; }
@@ -124,6 +180,13 @@ function transformExpr(e) { return lowerExpr(e, null); }
  * a method call can appear inline in TS.  It does NOT propagate into
  * field, index, record, forall, or exists sub-expressions.
  */
+/** Wrap an expression in Some/None for optional-typed conditionals.
+ *  If the raw TExpr is `undefined`, emit `.none`; otherwise wrap in `Some`. */
+function wrapOptionalBranch(expr, raw) {
+    return (raw.kind === "var" && raw.name === "undefined")
+        ? { kind: "constructor", name: ".none" }
+        : { kind: "app", fn: "Some", args: [expr] };
+}
 function lowerExpr(e, binds) {
     // Monadic lifting: extract embedded method calls to let-binds
     // Pass binds through to args so nested method calls are also lifted
@@ -215,6 +278,11 @@ function lowerExpr(e, binds) {
                     ],
                 };
             }
+            // || 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);
@@ -228,6 +296,43 @@ function lowerExpr(e, binds) {
                     ],
                 };
             }
+            // || 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,
@@ -249,6 +354,15 @@ function lowerExpr(e, binds) {
             return { kind: "index", arr: transformExpr(e.obj), idx: wrappedIdx };
         }
         case "call": {
+            // 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];
@@ -296,13 +410,82 @@ function lowerExpr(e, binds) {
             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, propagate declared field types and wrap optionals
+            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;
+                // Also check discriminated-union variants for field types
+                const unionDecl = structName ? _typeDecls.find(d => d.name === structName && d.kind === "discriminated-union") : undefined;
+                const loweredFields = e.fields.map(f => {
+                    // Propagate declared field type onto value if it has unknown type
+                    let fieldValue = f.value;
+                    const fieldDecl = structDecl?.fields?.find(sf => sf.name === f.name);
+                    let declaredTy;
+                    if (fieldDecl) {
+                        declaredTy = fieldDecl.type;
+                    }
+                    else if (unionDecl?.variants) {
+                        for (const v of unionDecl.variants) {
+                            const vf = v.fields.find(vf => vf.name === f.name);
+                            if (vf) {
+                                declaredTy = vf.type;
+                                break;
+                            }
+                        }
+                    }
+                    if (declaredTy && fieldValue.ty.kind === "unknown") {
+                        fieldValue = { ...fieldValue, ty: declaredTy };
+                    }
+                    let value = lowerExpr(fieldValue, binds);
+                    // Wrap non-optional values in Some for optional fields
+                    if (declaredTy?.kind === "optional") {
+                        const isUndef = f.value.kind === "var" && f.value.name === "undefined";
+                        if (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: p.ty })), body: transformStmts(e.body, []) };
@@ -311,23 +494,53 @@ function lowerExpr(e, binds) {
         case "exists":
             return { kind: "exists", var: e.var, type: e.varTy, body: transformExpr(e.body) };
         case "conditional": {
-            const cond = lowerExpr(e.cond, binds);
+            // When narrowedExpr is set, the match replaces the condition — don't lift from it
+            const condBinds = (e.narrowedVar && e.narrowedExpr) ? null : binds;
+            const cond = lowerExpr(e.cond, condBinds);
             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] };
+            // 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 });
                 }
-                if (e.else.kind === "var" && e.else.name === "undefined") {
-                    elseExpr = { kind: "constructor", name: ".none" };
+                // Wrap in Some/None only when result is optional (one branch is undefined)
+                if (e.ty.kind === "optional") {
+                    thenExpr = wrapOptionalBranch(thenExpr, e.then);
+                    elseExpr = wrapOptionalBranch(elseExpr, e.else);
                 }
-                else {
-                    elseExpr = { kind: "app", fn: "Some", args: [elseExpr] };
+                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.
+                thenExpr = wrapOptionalBranch(thenExpr, e.then);
+                elseExpr = wrapOptionalBranch(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") {
+                thenExpr = wrapOptionalBranch(thenExpr, e.then);
+                elseExpr = wrapOptionalBranch(elseExpr, e.else);
             }
             return { kind: "if", cond, then: thenExpr, else: elseExpr };
         }
@@ -374,7 +587,7 @@ 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 = buildMatchPattern(variantName, fields, obj.name);
     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 } }] };
@@ -384,7 +597,7 @@ function replaceFieldAccess(e, varName, fields) {
         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) };
+                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)
@@ -407,13 +620,12 @@ function transformStmts(stmts, typeDecls) {
                 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));
+            const optMatch = prepareOptionalMatch(s, stmts.slice(i + 1));
+            if (optMatch) {
+                result.push(emitOptionalMatch(optMatch.check.varName, optMatch.check.negated, s, typeDecls, stmts.slice(i + 1), optMatch.check.fieldExpr));
                 // 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) {
+                const origSome = optMatch.check.negated ? s.else : s.then;
+                if (origSome.length === 0 && i + 1 < stmts.length) {
                     return result;
                 }
                 i++;
@@ -610,6 +822,17 @@ function transformStmt(s, typeDecls) {
             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) }];
@@ -686,7 +909,7 @@ 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) {
+function emitOptionalMatch(varName, negated, s, typeDecls, restStmts, fieldExpr) {
     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,
@@ -695,20 +918,87 @@ function emitOptionalMatch(varName, negated, s, typeDecls, restStmts) {
         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 };
+    // Replace the narrowed variable/field in the Some branch body.
+    // Field chains: replace in TStmt before transform (so downstream narrowing sees simple vars).
+    // Simple vars: replace in IR after transform (the original mechanism).
+    let someBody;
+    if (fieldExpr && fieldExpr.kind === "field" && fieldExpr.obj.kind === "var") {
+        const innerTy = fieldExpr.ty.kind === "optional" ? fieldExpr.ty.inner : fieldExpr.ty;
+        const replaced = replaceFieldsInTStmts(someBranch, fieldExpr.obj.name, [
+            { fieldName: fieldExpr.field, newName: bound, fallbackTy: innerTy },
+        ]);
+        someBody = transformStmts(replaced, typeDecls);
+    }
+    else {
+        const transformed = transformStmts(someBranch, typeDecls);
+        someBody = transformed.map(stmt => mapStmtExprs(stmt, e => replaceVar(e, varName, { kind: "var", name: bound })));
+    }
+    return {
+        kind: "match", scrutinee: varName,
+        arms: [
+            { pattern: `.some ${bound}`, body: someBody },
+            { pattern: ".none", body: noneBranch.length > 0 ? transformStmts(noneBranch, typeDecls) : [] },
+        ],
+    };
 }
 /** 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. */
+// ── Optional narrowing helpers ──────────────────────────────
+//
+// Optional narrowing converts TS `if (x === undefined)` patterns to Dafny
+// `match x { Some(val) => ..., None => ... }`.
+//
+// The resolve phase (resolve.ts) handles:
+//   - Flow narrowing: after `if (x === undefined) return`, x is non-optional
+//   - && narrowing: in `x !== undefined && f(x)`, f(x) sees x as non-optional
+//   - Conditional narrowing: in `x !== undefined ? x.field : default`, sets
+//     narrowedVar/narrowedExpr on TExpr for the transform phase
+//
+// The transform phase (here) handles:
+//   - Statement-level: `transformStmts` detects optional checks → `emitOptionalMatch`
+//   - Expression-level: `lowerExpr` conditional reads narrowedVar/narrowedExpr → match
+//   - && restructuring: `extractLeftmostOptional` splits `&&` chains into nested ifs
+//     so `emitOptionalMatch` can detect the inner optional check
+//
+// Both phases detect `v !== undefined` patterns. The resolve phase uses
+// `detectOptionalCheck` (on RawExpr), the transform uses `parseOptionalCheck` (on TExpr).
+// These are separate because they operate on different IR types, but both handle
+// simple variables and field access chains.
+/** Shared logic for optional match in both imperative and pure function paths.
+ *  Detects optional check, selects branches, handles early-return consumption.
+ *  Returns null if the condition is not an optional check. */
+function prepareOptionalMatch(s, restStmts) {
+    const check = parseOptionalCheck(s.cond);
+    if (!check)
+        return null;
+    let someBranch = check.negated ? s.else : s.then;
+    const noneBranch = check.negated ? s.then : (s.else.length > 0 ? s.else : restStmts);
+    // Early-return pattern: Some branch is empty → consume rest of block
+    if (someBranch.length === 0 && restStmts.length > 0)
+        someBranch = restStmts;
+    const bound = matchBinder(`${check.varName}_val`);
+    return { check, someBranch, noneBranch, bound };
+}
+/** 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.
+ *  Also handles field access chains like `obj.field !== undefined`.
+ *  When `fieldExpr` is returned, callers must use field-aware replacement. */
 function parseOptionalCheck(cond) {
     if (cond.kind !== "binop" || (cond.op !== "!==" && cond.op !== "==="))
         return null;
@@ -717,20 +1007,50 @@ function parseOptionalCheck(cond) {
         varExpr = cond.left;
     if (cond.left.kind === "var" && cond.left.name === "undefined")
         varExpr = cond.right;
-    if (!varExpr || varExpr.kind !== "var" || varExpr.ty.kind !== "optional")
+    if (!varExpr)
         return null;
-    return { varName: varExpr.name, negated: cond.op === "===" };
+    if (varExpr.kind === "var" && varExpr.ty.kind === "optional") {
+        return { varName: varExpr.name, negated: cond.op === "===" };
+    }
+    if (varExpr.kind === "field" && varExpr.ty.kind === "optional") {
+        // Serialize field chain as a dotted name for use as match scrutinee
+        const chain = serializeFieldChain(varExpr);
+        if (chain)
+            return { varName: chain, negated: cond.op === "===", fieldExpr: varExpr };
+    }
+    return null;
 }
-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);
+/** Serialize a field access chain to a dotted variable path, or null if not a simple chain. */
+function serializeFieldChain(e) {
+    if (e.kind === "var")
+        return e.name;
+    if (e.kind === "field") {
+        const parent = serializeFieldChain(e.obj);
+        return parent ? `${parent}.${e.field}` : null;
+    }
+    return null;
+}
+/** Build match arms from variant cases — shared by imperative and pure paths.
+ *  Looks up variant fields from typeDecls, builds patterns via buildMatchPattern,
+ *  and delegates body transformation to the caller-provided function.
+ *  Returns null if any body transformation returns null (pure path abort). */
+function buildMatchArms(cases, varName, typeName, typeDecls, transformBody) {
+    const decl = typeName ? typeDecls.find(d => d.name === typeName) : undefined;
+    const arms = [];
+    for (const c of cases) {
+        const variant = decl?.variants?.find(v => v.name === c.name);
         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);
-        return { pattern, body };
-    });
+        const pattern = buildMatchPattern(c.name, fields, varName);
+        const body = transformBody(c.body, varName, fields);
+        if (body === null)
+            return null;
+        arms.push({ pattern, body });
+    }
+    return arms;
+}
+function emitMatchStmt(chain, typeDecls) {
+    const cases = chain.cases.map(c => ({ name: c.variant, body: c.body }));
+    const arms = buildMatchArms(cases, chain.varName, chain.typeName, typeDecls, (body, vn, fields) => transformStmts(replaceFieldAccessInTStmts(body, vn, fields), typeDecls));
     if (chain.fallthrough.length > 0)
         arms.push({ pattern: "_", body: transformStmts(chain.fallthrough, typeDecls) });
     return { kind: "match", scrutinee: chain.varName, arms };
@@ -738,41 +1058,39 @@ function emitMatchStmt(chain, typeDecls) {
 function emitSwitchStmt(s, typeDecls) {
     const varName = s.expr.kind === "var" ? s.expr.name : "?";
     const typeName = s.expr.ty.kind === "user" ? s.expr.ty.name : undefined;
-    const decl = typeName ? typeDecls.find(d => d.name === typeName) : undefined;
-    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);
-        return { pattern, body };
-    });
+    const cases = s.cases.map(c => ({ name: c.label, body: c.body }));
+    const arms = buildMatchArms(cases, varName, typeName, typeDecls, (body, vn, fields) => transformStmts(replaceFieldAccessInTStmts(body, vn, fields), typeDecls));
     if (s.defaultBody.length > 0)
         arms.push({ pattern: "_", body: transformStmts(s.defaultBody, typeDecls) });
     return { kind: "match", scrutinee: varName, arms };
 }
-function replaceFieldAccessInStmts(stmts, varName, fields) {
-    if (fields.length === 0)
+/** Replace obj.field → replacement var in typed IR (before transform).
+ *  Used by discriminant match/switch and optional match to rewrite field accesses
+ *  into simple variables before the transform phase, so downstream narrowing
+ *  (parseOptionalCheck, extractLeftmostOptional) sees simple variable references.
+ *  Uses the TExpr's resolved type when available, falling back to `fallbackTy`. */
+function replaceFieldsInTStmts(stmts, objName, replacements) {
+    if (replacements.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 stmts.map(s => mapTStmt(s, e => {
+        if (e.kind === "field" && e.obj.kind === "var" && e.obj.name === objName) {
+            const r = replacements.find(r => r.fieldName === e.field);
+            if (r) {
+                const ty = e.ty.kind !== "unknown" ? e.ty : r.fallbackTy;
+                return { kind: "var", name: r.newName, ty };
+            }
         }
         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) {
-            result.push(s.value ? { ...s, value: mapExpr(s.value, f) } : s);
-            result.push(...stmts.slice(result.length));
-            break;
-        }
-        result.push(mapStmt(s, f));
-    }
-    return result;
+    }));
+}
+/** Replace all variant fields of obj → match binder vars in typed IR.
+ *  Thin wrapper around replaceFieldsInTStmts for discriminant match/switch. */
+function replaceFieldAccessInTStmts(stmts, varName, fields) {
+    return replaceFieldsInTStmts(stmts, varName, fields.map(f => ({
+        fieldName: f.name,
+        newName: matchBinder(f.name, varName),
+        fallbackTy: f.type ?? parseTsType(f.tsType),
+    })));
 }
 // ── Pure function generation ─────────────────────────────────
 function transformPureBody(stmts, typeDecls) {
@@ -795,24 +1113,19 @@ function transformPureBody(stmts, typeDecls) {
             }
             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);
+                const optMatch = prepareOptionalMatch(s, rest);
+                if (optMatch) {
+                    const someExpr = transformPureBody(optMatch.someBranch, typeDecls);
                     if (!someExpr)
                         return null;
-                    const noneExpr = transformPureBody(noneBranch, typeDecls);
+                    const noneExpr = transformPureBody(optMatch.noneBranch, typeDecls);
                     if (!noneExpr)
                         return null;
-                    const someReplaced = replaceVar(someExpr, optCheck.varName, { kind: "var", name: bound });
+                    const someReplaced = replaceVar(someExpr, optMatch.check.varName, { kind: "var", name: optMatch.bound });
                     return {
-                        kind: "match", scrutinee: optCheck.varName,
+                        kind: "match", scrutinee: optMatch.check.varName,
                         arms: [
-                            { pattern: `.some ${bound}`, body: someReplaced },
+                            { pattern: `.some ${optMatch.bound}`, body: someReplaced },
                             { pattern: ".none", body: noneExpr },
                         ],
                     };
@@ -833,21 +1146,21 @@ function transformPureBody(stmts, typeDecls) {
     return null;
 }
 function transformPureSwitch(s, typeDecls) {
-    const decl = typeDecls.find(d => d.name === (s.expr.ty.kind === "user" ? s.expr.ty.name : ""));
-    if (!decl)
+    const typeName = s.expr.ty.kind === "user" ? s.expr.ty.name : "";
+    if (!typeDecls.find(d => d.name === typeName))
         return null;
-    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}`;
-        let body = transformPureBody(c.body, typeDecls);
-        if (!body)
+    const varName = s.expr.kind === "var" ? s.expr.name : undefined;
+    const cases = s.cases.map(c => ({ name: c.label, body: c.body }));
+    const arms = buildMatchArms(cases, varName, typeName, typeDecls, (body, vn, fields) => {
+        let result = transformPureBody(body, typeDecls);
+        if (!result)
             return null;
-        if (fields.length > 0 && s.expr.kind === "var")
-            body = replaceFieldAccess(body, s.expr.name, fields);
-        arms.push({ pattern, body });
-    }
+        if (fields.length > 0 && vn)
+            result = replaceFieldAccess(result, vn, fields);
+        return result;
+    });
+    if (!arms)
+        return null;
     if (s.defaultBody.length > 0) {
         const body = transformPureBody(s.defaultBody, typeDecls);
         if (!body)
@@ -859,22 +1172,21 @@ function transformPureSwitch(s, typeDecls) {
     return { kind: "match", scrutinee: s.expr.name, arms };
 }
 function transformPureMatch(chain, typeDecls) {
-    const decl = typeDecls.find(d => d.name === chain.typeName);
-    const arms = [];
-    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}`;
-        let body = transformPureBody(c.body, typeDecls);
-        if (!body)
+    const cases = chain.cases.map(c => ({ name: c.variant, body: c.body }));
+    const arms = buildMatchArms(cases, chain.varName, chain.typeName, typeDecls, (body, vn, fields) => {
+        let result = transformPureBody(body, typeDecls);
+        if (!result)
             return null;
-        if (fields.length > 0)
-            body = replaceFieldAccess(body, chain.varName, fields);
-        arms.push({ pattern, body });
-    }
+        if (fields.length > 0 && vn)
+            result = replaceFieldAccess(result, vn, fields);
+        return result;
+    });
+    if (!arms)
+        return null;
     // Idiomatic TS often has an unreachable fallthrough after exhaustive if-chains on
     // discriminated unions. Skip the catch-all arm when all variants are matched,
     // since Lean errors on redundant match arms.
+    const decl = typeDecls.find(d => d.name === chain.typeName);
     const allCovered = decl?.variants && chain.cases.length >= decl.variants.length;
     if (chain.fallthrough.length > 0 && !allCovered) {
         const body = transformPureBody(chain.fallthrough, typeDecls);
@@ -896,17 +1208,24 @@ 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: parseTsType(f.tsType) })),
+                fields: v.fields.map(f => ({ name: f.name, type: f.type })),
             })),
             deriving: ["Repr", "Inhabited"],
         };
     }
+    else if (d.kind === "alias") {
+        return {
+            kind: "type-alias", name: d.name,
+            target: d.aliasOfTy,
+        };
+    }
     else {
         return {
             kind: "structure", name: d.name,
-            fields: d.fields.map(f => ({ name: f.name, type: parseTsType(f.tsType) })),
+            fields: d.fields.map(f => ({ name: f.name, type: f.type })),
             deriving: ["Repr", "Inhabited", "DecidableEq"],
         };
     }
@@ -961,6 +1280,17 @@ function replaceVar(e, name, replacement) {
     });
 }
 // ── 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;
@@ -974,6 +1304,8 @@ export function transformModuleDafny(mod) {
 }
 export function transformModule(mod, specImport) {
     _forofCounters.clear();
+    _liftCounter = 0;
+    _typeDecls = mod.typeDecls;
     const typeDecls = mod.typeDecls.map(transformTypeDecl);
     // Module-level constants
     const constDecls = (mod.constants ?? []).map(c => ({
@@ -996,6 +1328,7 @@ export function transformModule(mod, specImport) {
         pureDefs.push({
             kind: "def",
             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),
@@ -1046,6 +1379,7 @@ export function transformModule(mod, specImport) {
         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),
@@ -1062,6 +1396,7 @@ export function transformModule(mod, specImport) {
             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),