lemmascript 0.3.0 → 0.3.2

package/tools/dist/transform.js

@@ -66,9 +66,6 @@ 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);
@@ -133,6 +130,12 @@ let _typeDecls = [];
 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"; }
 function isArray(ty) { return ty.kind === "array"; }
@@ -177,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
@@ -340,6 +350,9 @@ function lowerExpr(e, binds) {
             return { kind: "field", obj: transformExpr(e.obj), field: e.field };
         case "index": {
             const idx = transformExpr(e.idx);
+            if (e.obj.ty.kind === "map") {
+                return { kind: "methodCall", obj: transformExpr(e.obj), objTy: e.obj.ty, method: "get", args: [idx], monadic: false };
+            }
             const wrappedIdx = isArray(e.obj.ty) && !isNat(e.idx.ty) ? { kind: "toNat", expr: idx } : idx;
             return { kind: "index", arr: transformExpr(e.obj), idx: wrappedIdx };
         }
@@ -427,23 +440,41 @@ function lowerExpr(e, binds) {
                     }
                 }
             }
-            // For spread records, wrap non-optional values in Some for optional fields
+            // 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 => {
-                    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] };
+                    // 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 };
@@ -466,7 +497,9 @@ 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);
             // Explicit !== undefined with narrowedExpr → match Some/None on the optional expression
@@ -476,11 +509,11 @@ function lowerExpr(e, binds) {
                 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);
+                // 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);
+                }
                 return {
                     kind: "match", scrutinee,
                     arms: [
@@ -497,12 +530,8 @@ function lowerExpr(e, binds) {
                     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);
+                thenExpr = wrapOptionalBranch(thenExpr, e.then);
+                elseExpr = wrapOptionalBranch(elseExpr, e.else);
                 return {
                     kind: "match", scrutinee: cond,
                     arms: [
@@ -513,18 +542,8 @@ function lowerExpr(e, binds) {
             }
             // 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] };
-                }
+                thenExpr = wrapOptionalBranch(thenExpr, e.then);
+                elseExpr = wrapOptionalBranch(elseExpr, e.else);
             }
             return { kind: "if", cond, then: thenExpr, else: elseExpr };
         }
@@ -571,7 +590,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, obj.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 } }] };
@@ -604,13 +623,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++;
@@ -622,6 +640,31 @@ function transformStmts(stmts, typeDecls) {
             const varName = s.names[0];
             const varTy = s.nameTypes[0] ?? { kind: "unknown" };
             let iterExpr = transformExpr(s.iterable);
+            // Map key-only iteration: for (const k in record) → iterate keys only
+            if (s.names.length === 1 && s.iterable.ty.kind === "map") {
+                const keyName = s.names[0];
+                const keyTy = s.nameTypes[0] ?? s.iterable.ty.key ?? { 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 boundInv = { kind: "binop", op: "≤", left: idx, right: arrSize };
+                result.push({
+                    kind: "forin", idx: idxName, bound: arrSize,
+                    invariants: [boundInv, ...s.invariants.map(transformExpr)],
+                    body: [letKey, ...bodyStmts],
+                });
+                i++;
+                continue;
+            }
             // 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];
@@ -894,7 +937,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,
@@ -903,19 +946,69 @@ 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));
 }
+// ── 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) {
@@ -931,7 +1024,9 @@ function extractLeftmostOptional(cond) {
     }
     return null;
 }
-/** Detect `v !== undefined` or `undefined !== v` where v has optional type. */
+/** 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;
@@ -940,21 +1035,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, 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 };
-    });
+        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 };
@@ -962,59 +1086,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, 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 };
-    });
+    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 };
 }
-/** 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)
+/** 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;
     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 };
+        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;
     }));
 }
-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) {
-            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) {
@@ -1037,24 +1141,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 },
                         ],
                     };
@@ -1075,22 +1174,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 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, varName)).join(" ")}` : `.${c.label}`;
-        let body = transformPureBody(c.body, typeDecls);
-        if (!body)
+    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)
@@ -1102,22 +1200,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, chain.varName)).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);
@@ -1142,7 +1239,7 @@ function transformTypeDecl(d) {
             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"],
         };
@@ -1150,13 +1247,13 @@ function transformTypeDecl(d) {
     else if (d.kind === "alias") {
         return {
             kind: "type-alias", name: d.name,
-            target: parseTsType(d.aliasOf),
+            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"],
         };
     }
@@ -1235,6 +1332,7 @@ 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