lemmascript 0.3.2 → 0.4.0

package/tools/dist/transform.js

@@ -21,10 +21,10 @@ function mapExpr(e, f) {
         case "num":
         case "bool":
         case "str":
-        case "constructor":
         case "emptyMap":
         case "emptySet":
         case "havoc": return e;
+        case "constructor": return e.args ? { ...e, args: e.args.map(r) } : 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) };
@@ -58,7 +58,10 @@ function mapStmt(s, f) {
         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 "match": {
+            const scr = typeof s.scrutinee === "string" ? s.scrutinee : r(s.scrutinee);
+            return { ...s, scrutinee: scr, 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) };
@@ -87,6 +90,15 @@ function mapTExpr(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 "conditional": return { ...e, cond: r(e.cond), then: r(e.then), else: r(e.else) };
+        case "optChain": return { ...e, obj: r(e.obj),
+            chain: e.chain.map(s => s.kind === "call" ? { ...s, args: s.args.map(r) }
+                : s.kind === "index" ? { ...s, idx: r(s.idx) }
+                    : s) };
+        case "nullish": return { ...e, left: r(e.left), right: r(e.right) };
+        case "someMatch": return { ...e, scrutinee: r(e.scrutinee), someBody: r(e.someBody), noneBody: r(e.noneBody) };
+        case "tagMatch": return { ...e, scrutinee: r(e.scrutinee),
+            cases: e.cases.map(c => ({ ...c, body: r(c.body) })),
+            fallthrough: e.fallthrough ? r(e.fallthrough) : null };
         case "forall": return { ...e, body: r(e.body) };
         case "exists": return { ...e, body: r(e.body) };
         case "lambda": return e;
@@ -110,6 +122,10 @@ function mapTStmt(s, 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) };
+        case "someMatch": return { ...s, scrutinee: r(s.scrutinee), someBody: s.someBody.map(t => mapTStmt(t, f)), noneBody: s.noneBody.map(t => mapTStmt(t, f)) };
+        case "tagMatch": return { ...s, scrutinee: r(s.scrutinee),
+            cases: s.cases.map(c => ({ ...c, body: c.body.map(t => mapTStmt(t, f)) })),
+            fallthrough: s.fallthrough.map(t => mapTStmt(t, f)) };
     }
 }
 export const LEAN_OPTIONS = {
@@ -183,18 +199,26 @@ function transformExpr(e) { return lowerExpr(e, null); }
 /** 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] };
+    // Set type: "Option" so Lean emits `Option.some`/`Option.none` (qualified).
+    // The dotted form `.some`/`.none` would be ambiguous in expression positions
+    // like the scrutinee of an outer match. Dafny treats `Option.Some` and bare
+    // `Some` equivalently — the qualification is harmless there.
+    if (raw.kind === "var" && raw.name === "undefined")
+        return { kind: "constructor", name: "none", type: "Option" };
+    if (raw.ty.kind === "optional")
+        return expr; // already Option<T>, don't double-wrap
+    return { kind: "constructor", name: "some", type: "Option", 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
-    if (binds && e.kind === "call" && e.callKind === "method") {
+    // Monadic lifting: extract embedded method calls to let-binds.
+    // `callKind: "method"` means a global var-fn call (classifyCall returns
+    // "method" only for `fn.kind === "var"`). Receiver method calls have
+    // callKind "unknown" and fall through to the regular case below where
+    // they become `methodCall`.
+    if (binds && e.kind === "call" && e.callKind === "method" && e.fn.kind === "var") {
         const name = `_t${_liftCounter++}`;
-        const fn = e.fn.kind === "var" ? e.fn.name : `${lowerExpr(e.fn, binds)}`;
         const args = e.args.map(a => lowerExpr(a, binds));
-        binds.push({ kind: "let-bind", name, value: { kind: "app", fn, args } });
+        binds.push({ kind: "let-bind", name, value: { kind: "app", fn: e.fn.name, args } });
         return { kind: "var", name };
     }
     switch (e.kind) {
@@ -347,11 +371,27 @@ function lowerExpr(e, binds) {
                 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" };
+            // Boolean discriminant bare access: `result.ok` where Result has variants
+            // {ok: true, ...} | {ok: false, ...}. String discriminants are always used via
+            // comparison (x.kind === 'Foo' → x.Foo?), but boolean discriminants are used
+            // as bare truthiness checks. Emit as the Dafny discriminator predicate for the
+            // 'true' variant: result.ok → result.true_?
+            if (e.isDiscriminant && e.obj.ty.kind === "user") {
+                const baseName = e.obj.ty.name.includes("<") ? e.obj.ty.name.slice(0, e.obj.ty.name.indexOf("<")) : e.obj.ty.name;
+                const decl = _typeDecls.find(d => d.name === baseName && d.kind === "discriminated-union");
+                if (decl?.variants?.some(v => v.name === "true")) {
+                    return { kind: "field", obj: transformExpr(e.obj), field: "true_?" };
+                }
+            }
             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 };
+                // Mirrors the .get() → .getDirect switch at line ~453: when resolve has
+                // narrowed the index type to non-optional (via `k in m` atoms in scope),
+                // emit direct access; otherwise keep the Option-producing `get`.
+                const method = e.ty.kind !== "optional" ? "getDirect" : "get";
+                return { kind: "methodCall", obj: transformExpr(e.obj), objTy: e.obj.ty, method, 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 };
@@ -393,10 +433,23 @@ function lowerExpr(e, binds) {
                         return { kind: "toNat", expr: lowered };
                     return lowered;
                 });
+                // arr.concat(otherArr): array argument → real concatenation, not push
+                if (method === "concat" && e.fn.obj.ty.kind === "array" && e.args.length === 1 && e.args[0].ty.kind === "array") {
+                    return { kind: "binop", op: "arrayConcat", left: recv, right: args[0] };
+                }
                 // 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";
                 }
+                // map.set(k, v): if v is an Optional-wrapped map get, unwrap to getDirect
+                // (the desugared spread { ...m, [k]: m2[k] } becomes m.set(k, m2.get(k)),
+                // but the value should be direct access, not Optional)
+                if (method === "set" && e.fn.obj.ty.kind === "map" && args.length === 2) {
+                    const val = args[1];
+                    if (val.kind === "methodCall" && val.method === "get" && val.objTy.kind === "map") {
+                        args[1] = { ...val, 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 };
@@ -479,6 +532,10 @@ function lowerExpr(e, binds) {
                 });
                 return { kind: "record", spread: lowerExpr(e.spread, binds), fields: loweredFields };
             }
+            // Empty record with map type → empty map
+            if (e.fields.length === 0 && !e.spread && e.ty.kind === "map") {
+                return { kind: "emptyMap" };
+            }
             return { kind: "record", spread: null, fields: e.fields.map(f => ({ name: f.name, value: lowerExpr(f.value, binds) })) };
         }
         case "arrayLiteral":
@@ -497,56 +554,21 @@ function lowerExpr(e, binds) {
         case "exists":
             return { kind: "exists", var: e.var, type: e.varTy, body: transformExpr(e.body) };
         case "conditional": {
-            // 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);
+            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 });
-                }
-                // 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: [
-                        { 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 };
         }
+        case "optChain":
+            // Narrow should have rewritten optChain to someMatch.
+            throw new Error(`optChain reached transform — narrow should have rewritten it`);
+        case "nullish":
+            // Narrow should have rewritten nullish to someMatch.
+            throw new Error(`nullish reached transform — narrow should have rewritten it`);
         case "havoc":
             // Dafny's * only works in var/assign positions — lift to own declaration
             if (binds) {
@@ -555,6 +577,41 @@ function lowerExpr(e, binds) {
                 return { kind: "var", name };
             }
             return { kind: "havoc", type: e.ty };
+        case "someMatch": {
+            let someBody;
+            let scrutinee;
+            const path = asTAccessPath(e.scrutinee);
+            if (path) {
+                // Pure access path (var or any depth of obj.f.g.h) — substitute the
+                // path with the binder pre-lowering.
+                const replaced = replacePathInTExpr(e.someBody, path, e.binder, e.binderTy);
+                someBody = lowerExpr(replaced, binds);
+                scrutinee = path.fields.length === 0 ? path.rootVar : lowerExpr(e.scrutinee, binds);
+            }
+            else {
+                // Complex scrutinee — narrow pre-bound the someBody to use the binder directly,
+                // so no substitution needed. Used by optChain rewrites.
+                someBody = lowerExpr(e.someBody, binds);
+                scrutinee = lowerExpr(e.scrutinee, binds);
+            }
+            let noneBody = lowerExpr(e.noneBody, binds);
+            if (e.ty.kind === "optional") {
+                someBody = wrapOptionalBranch(someBody, e.someBody);
+                noneBody = wrapOptionalBranch(noneBody, e.noneBody);
+            }
+            return {
+                kind: "match", scrutinee,
+                arms: [
+                    { pattern: `.some ${e.binder}`, body: someBody },
+                    { pattern: ".none", body: noneBody },
+                ],
+            };
+        }
+        case "tagMatch":
+            // Spec/expr-position tagMatch — narrow shouldn't produce these (it only
+            // emits stmt-form tagMatch from if-chain detection on stmts). If reached,
+            // bug in narrow.
+            throw new Error(`tagMatch reached lowerExpr — narrow only emits stmt-form tagMatch`);
     }
 }
 function flattenImpl(e) {
@@ -614,27 +671,6 @@ function transformStmts(stmts, typeDecls) {
     let i = 0;
     while (i < stmts.length) {
         const s = stmts[i];
-        // Detect discriminant if-chain → match
-        if (s.kind === "if") {
-            const chain = detectDiscriminantChain(stmts.slice(i));
-            if (chain) {
-                result.push(emitMatchStmt(chain.chain, typeDecls));
-                i += chain.consumed;
-                continue;
-            }
-            // Detect optional check → match on Some/None
-            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 origSome = optMatch.check.negated ? s.else : s.then;
-                if (origSome.length === 0 && i + 1 < stmts.length) {
-                    return result;
-                }
-                i++;
-                continue;
-            }
-        }
         // Transform for-of → for-in over range
         if (s.kind === "forof") {
             const varName = s.names[0];
@@ -644,13 +680,13 @@ function transformStmts(stmts, typeDecls) {
             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 keysSeqName = `_${keyName}_keys${suffix}`;
+                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 idxName = `_${keyName}_idx${suffix}`;
                 const idx = { kind: "var", name: idxName };
                 const arrSize = { kind: "field", obj: keysVar, field: "size" };
@@ -670,13 +706,13 @@ function transformStmts(stmts, typeDecls) {
                 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 keysSeqName = `_${keyName}_keys${suffix}`;
+                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 idxName = `_${keyName}_idx${suffix}`;
                 const idx = { kind: "var", name: idxName };
                 const arrSize = { kind: "field", obj: keysVar, field: "size" };
@@ -828,15 +864,19 @@ function transformStmt(s, typeDecls) {
                 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;
+            // Optional chaining on map.get at statement level: m.get(k)?.push(v)
+            // → if k in m { m[k] := m[k] + [v] } (actual mutation, not value-discard).
+            // Narrow rewrote this to a someMatch — destructure to find the underlying
+            // m.get(k) scrutinee and the .push(v) body call.
+            if (s.expr.kind === "someMatch" &&
+                s.expr.scrutinee.kind === "call" && s.expr.scrutinee.fn.kind === "field" &&
+                s.expr.scrutinee.fn.field === "get" && s.expr.scrutinee.fn.obj.ty.kind === "map" &&
+                s.expr.someBody.kind === "call" && s.expr.someBody.fn.kind === "field" &&
+                s.expr.someBody.fn.field === "push") {
+                const mapExpr = s.expr.scrutinee.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);
+                const keyExpr = lowerExpr(s.expr.scrutinee.args[0], null);
+                const pushArg = lowerExpr(s.expr.someBody.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 };
@@ -850,18 +890,7 @@ 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)
+            // 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) }];
         }
@@ -886,178 +915,34 @@ function transformStmt(s, typeDecls) {
             return [{ kind: "ghostAssign", target: s.target, value: transformExpr(s.value) }];
         case "assert":
             return [{ kind: "assert", expr: transformExpr(s.expr) }];
-    }
-}
-function detectDiscriminantChain(stmts) {
-    if (stmts.length === 0 || stmts[0].kind !== "if")
-        return null;
-    const first = parseDiscriminantCond(stmts[0].cond);
-    if (!first)
-        return null;
-    const cases = [];
-    // Follow else branches within one if-else-if tree
-    function collectElse(s) {
-        const p = parseDiscriminantCond(s.cond);
-        if (!p || p.varName !== first.varName)
-            return [s];
-        cases.push({ variant: p.variant, body: s.then });
-        if (s.else.length === 0)
-            return [];
-        if (s.else.length === 1 && s.else[0].kind === "if")
-            return collectElse(s.else[0]);
-        return s.else;
-    }
-    // Walk consecutive top-level ifs on the same discriminant
-    let consumed = 0;
-    for (let i = 0; i < stmts.length; i++) {
-        const s = stmts[i];
-        if (s.kind !== "if")
-            break;
-        const p = parseDiscriminantCond(s.cond);
-        if (!p || p.varName !== first.varName)
-            break;
-        cases.push({ variant: p.variant, body: s.then });
-        consumed = i + 1;
-        if (s.else.length > 0) {
-            const ft = (s.else.length === 1 && s.else[0].kind === "if") ? collectElse(s.else[0]) : s.else;
-            return cases.length > 0 ? { chain: { ...first, cases, fallthrough: ft }, consumed } : null;
+        case "someMatch": {
+            const path = asTAccessPath(s.scrutinee);
+            if (path) {
+                const replaced = replacePathInTStmts(s.someBody, path, s.binder, s.binderTy);
+                const someBody = transformStmts(replaced, typeDecls);
+                const noneBody = transformStmts(s.noneBody, typeDecls);
+                const scrutinee = path.fields.length === 0 ? path.rootVar : transformExpr(s.scrutinee);
+                return [{
+                        kind: "match", scrutinee,
+                        arms: [
+                            { pattern: `.some ${s.binder}`, body: someBody },
+                            { pattern: ".none", body: noneBody },
+                        ],
+                    }];
+            }
+            throw new Error(`someMatch stmt scrutinee must be a pure access path, got ${s.scrutinee.kind}`);
+        }
+        case "tagMatch": {
+            const varName = s.scrutinee.kind === "var" ? s.scrutinee.name : "?";
+            const chain = { varName, typeName: s.typeName, cases: s.cases, fallthrough: s.fallthrough };
+            return [emitMatchStmt(chain, typeDecls)];
         }
     }
-    if (cases.length === 0)
-        return null;
-    return { chain: { ...first, cases, fallthrough: stmts.slice(consumed) }, consumed: stmts.length };
-}
-function parseDiscriminantCond(cond) {
-    // Pattern: x.discriminant === "variant"
-    if (cond.kind !== "binop" || cond.op !== "===" || cond.right.kind !== "str")
-        return null;
-    if (cond.left.kind !== "field" || !cond.left.isDiscriminant)
-        return null;
-    if (cond.left.obj.kind !== "var" || cond.left.obj.ty.kind !== "user")
-        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, 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,
-    // so include remaining statements as the Some body
-    if (someBranch.length === 0 && restStmts && restStmts.length > 0) {
-        someBranch = restStmts;
-    }
-    const bound = matchBinder(`${varName}_val`);
-    // 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) {
-    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;
-    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)
-        return null;
-    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;
-}
-/** 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.
@@ -1079,10 +964,32 @@ function buildMatchArms(cases, varName, typeName, typeDecls, transformBody) {
 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) });
+    if (chain.fallthrough.length > 0) {
+        const remaining = remainingVariant(chain, typeDecls);
+        if (remaining) {
+            // Exactly one variant left — destructure so the fallthrough body can
+            // access variant-specific fields (Lean requires this; Dafny tolerates `_`).
+            const pattern = buildMatchPattern(remaining.name, remaining.fields, chain.varName);
+            const body = transformStmts(replaceFieldAccessInTStmts(chain.fallthrough, chain.varName, remaining.fields), typeDecls);
+            arms.push({ pattern, body });
+        }
+        else {
+            arms.push({ pattern: "_", body: transformStmts(chain.fallthrough, typeDecls) });
+        }
+    }
     return { kind: "match", scrutinee: chain.varName, arms };
 }
+/** If the chain has matched all variants but one, return that remaining variant. */
+function remainingVariant(chain, typeDecls) {
+    const decl = typeDecls.find(d => d.name === chain.typeName);
+    if (!decl?.variants)
+        return null;
+    const matched = new Set(chain.cases.map(c => c.variant));
+    const remaining = decl.variants.filter(v => !matched.has(v.name));
+    if (remaining.length !== 1)
+        return null;
+    return remaining[0];
+}
 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;
@@ -1092,10 +999,7 @@ function emitSwitchStmt(s, typeDecls) {
         arms.push({ pattern: "_", body: transformStmts(s.defaultBody, typeDecls) });
     return { kind: "match", scrutinee: varName, arms };
 }
-/** 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.
+/** Replace obj.field → replacement var in typed IR.
  *  Uses the TExpr's resolved type when available, falling back to `fallbackTy`. */
 function replaceFieldsInTStmts(stmts, objName, replacements) {
     if (replacements.length === 0)
@@ -1120,13 +1024,71 @@ function replaceFieldAccessInTStmts(stmts, varName, fields) {
         fallbackTy: f.type ?? parseTsType(f.tsType),
     })));
 }
+/** Replace obj.field → replacement var in typed IR expressions (before lowering).
+ *  Mirrors replaceFieldsInTStmts but operates on a single TExpr tree. */
+function replaceFieldInTExpr(expr, objName, replacements) {
+    if (replacements.length === 0)
+        return expr;
+    return mapTExpr(expr, 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;
+    });
+}
+function asTAccessPath(e) {
+    if (e.kind === "var")
+        return { rootVar: e.name, fields: [] };
+    if (e.kind === "field") {
+        const inner = asTAccessPath(e.obj);
+        if (!inner)
+            return null;
+        return { rootVar: inner.rootVar, fields: [...inner.fields, e.field] };
+    }
+    return null;
+}
+/** Does TExpr `e` match the given access path exactly? */
+function matchesAccessPath(e, path) {
+    const collected = [];
+    let cur = e;
+    while (cur.kind === "field") {
+        collected.unshift(cur.field);
+        cur = cur.obj;
+    }
+    if (cur.kind !== "var" || cur.name !== path.rootVar)
+        return false;
+    if (collected.length !== path.fields.length)
+        return false;
+    return collected.every((f, i) => f === path.fields[i]);
+}
+/** Replace every TExpr matching `path` with `var(binder, binderTy)`. */
+function replacePathInTExpr(expr, path, binder, binderTy) {
+    return mapTExpr(expr, e => matchesAccessPath(e, path)
+        ? { kind: "var", name: binder, ty: binderTy } : null);
+}
+function replacePathInTStmts(stmts, path, binder, binderTy) {
+    return stmts.map(s => mapTStmt(s, e => matchesAccessPath(e, path)
+        ? { kind: "var", name: binder, ty: binderTy } : null));
+}
+/** Unwrap optional type on match-bound variables in TExpr.
+ *  After replaceFieldInTExpr, the replaced variable carries the original optional
+ *  type from the field declaration. The match binding unwraps it to the inner type. */
+function fixBoundType(expr, boundName) {
+    return mapTExpr(expr, e => e.kind === "var" && e.name === boundName && e.ty.kind === "optional"
+        ? { ...e, ty: e.ty.inner } : null);
+}
 // ── Pure function generation ─────────────────────────────────
 function transformPureBody(stmts, typeDecls) {
-    // Detect discriminant if-chain
-    if (stmts.length > 0 && stmts[0].kind === "if") {
-        const chain = detectDiscriminantChain(stmts);
-        if (chain)
-            return transformPureMatch(chain.chain, typeDecls);
+    // tagMatch (from narrow's discriminant detection) is the leading stmt and consumes the rest.
+    if (stmts.length > 0 && stmts[0].kind === "tagMatch") {
+        const t = stmts[0];
+        const varName = t.scrutinee.kind === "var" ? t.scrutinee.name : "?";
+        const chain = { varName, typeName: t.typeName, cases: t.cases, fallthrough: t.fallthrough };
+        return transformPureMatch(chain, typeDecls);
     }
     for (let i = 0; i < stmts.length; i++) {
         const s = stmts[i];
@@ -1140,34 +1102,39 @@ 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 optMatch = prepareOptionalMatch(s, rest);
-                if (optMatch) {
-                    const someExpr = transformPureBody(optMatch.someBranch, typeDecls);
+                // Append rest to both branches so nested ifs that fall through
+                // can reach the continuation (e.g. early return inside then-branch).
+                const thenExpr = transformPureBody([...s.then, ...rest], typeDecls);
+                if (!thenExpr)
+                    return null;
+                const elseStmts = s.else.length > 0 ? [...s.else, ...rest] : rest;
+                const elseExpr = transformPureBody(elseStmts, typeDecls);
+                if (!elseExpr)
+                    return null;
+                return { kind: "if", cond: transformExpr(s.cond), then: thenExpr, else: elseExpr };
+            }
+            case "switch": return transformPureSwitch(s, typeDecls);
+            case "someMatch": {
+                const path = asTAccessPath(s.scrutinee);
+                if (path) {
+                    const replaced = replacePathInTStmts(s.someBody, path, s.binder, s.binderTy);
+                    const someExpr = transformPureBody([...replaced, ...rest], typeDecls);
                     if (!someExpr)
                         return null;
-                    const noneExpr = transformPureBody(optMatch.noneBranch, typeDecls);
+                    const noneExpr = transformPureBody([...s.noneBody, ...rest], typeDecls);
                     if (!noneExpr)
                         return null;
-                    const someReplaced = replaceVar(someExpr, optMatch.check.varName, { kind: "var", name: optMatch.bound });
+                    const scrutinee = path.fields.length === 0 ? path.rootVar : transformExpr(s.scrutinee);
                     return {
-                        kind: "match", scrutinee: optMatch.check.varName,
+                        kind: "match", scrutinee,
                         arms: [
-                            { pattern: `.some ${optMatch.bound}`, body: someReplaced },
+                            { pattern: `.some ${s.binder}`, body: someExpr },
                             { pattern: ".none", body: noneExpr },
                         ],
                     };
                 }
-                const thenExpr = transformPureBody(s.then, typeDecls);
-                if (!thenExpr)
-                    return null;
-                const elseBranch = s.else.length > 0 ? s.else : rest;
-                const elseExpr = transformPureBody(elseBranch, typeDecls);
-                if (!elseExpr)
-                    return null;
-                return { kind: "if", cond: transformExpr(s.cond), then: thenExpr, else: elseExpr };
+                throw new Error(`someMatch pure-body scrutinee must be a pure access path, got ${s.scrutinee.kind}`);
             }
-            case "switch": return transformPureSwitch(s, typeDecls);
             default: return null;
         }
     }
@@ -1217,10 +1184,22 @@ function transformPureMatch(chain, typeDecls) {
     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);
-        if (!body)
-            return null;
-        arms.push({ pattern: "_", body });
+        const remaining = remainingVariant(chain, typeDecls);
+        if (remaining) {
+            // Exactly one variant left — destructure for variant-specific field access.
+            let body = transformPureBody(chain.fallthrough, typeDecls);
+            if (!body)
+                return null;
+            if (remaining.fields.length > 0)
+                body = replaceFieldAccess(body, chain.varName, remaining.fields);
+            arms.push({ pattern: buildMatchPattern(remaining.name, remaining.fields, chain.varName), body });
+        }
+        else {
+            const body = transformPureBody(chain.fallthrough, typeDecls);
+            if (!body)
+                return null;
+            arms.push({ pattern: "_", body });
+        }
     }
     return { kind: "match", scrutinee: chain.varName, arms };
 }
@@ -1293,17 +1272,38 @@ function findReassignedNames(stmts, names) {
     return found;
 }
 /** Replace all occurrences of a variable name with a new expression. */
-function replaceVar(e, name, replacement) {
+/**
+ * Replace all occurrences of variable `name` with `replacement`.
+ * If `narrowing` is true, the replacement is an unwrapped Optional value
+ * (e.g., replacing `x: Option<T>` with `x_val: T`). In that case, when the
+ * variable appears directly as a record spread field value, it's wrapped in
+ * Some() to preserve the field's Optional type.
+ */
+function replaceVar(e, name, replacement, narrowing) {
+    const rec = (expr) => replaceVar(expr, name, replacement, narrowing);
     return mapExpr(e, x => {
         if (x.kind === "var" && x.name === name)
             return replacement;
+        // Record spread: wrap direct variable uses in field values with Some when narrowing
+        if (narrowing && x.kind === "record" && x.spread) {
+            return {
+                ...x,
+                spread: rec(x.spread),
+                fields: x.fields.map(f => {
+                    if (f.value.kind === "var" && f.value.name === name) {
+                        return { ...f, value: { kind: "app", fn: "Some", args: [replacement] } };
+                    }
+                    return { ...f, value: rec(f.value) };
+                }),
+            };
+        }
         // 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 { ...x, value: replaceVar(x.value, name, replacement, narrowing) };
         return null;
     });
 }
@@ -1344,25 +1344,43 @@ export function transformModule(mod, specImport) {
     }));
     // Pure function mirrors
     const pureDefs = [];
+    const defByMethods = [];
     for (const fn of mod.functions) {
         if (!fn.isPure)
             continue;
         const body = transformPureBody(fn.body, mod.typeDecls);
-        if (!body)
-            continue;
-        // For ensures, replace \result (→ "res") with the function call
-        const fnCall = { kind: "app", fn: fn.name, args: fn.params.map(p => ({ kind: "var", name: p.name })) };
-        const ensures = fn.ensures.map(e => replaceVar(transformExpr(e), "res", fnCall));
-        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),
-            ensures,
-            body,
-        });
+        if (body) {
+            // For ensures, replace \result (→ "res") with the function call
+            const fnCall = { kind: "app", fn: fn.name, args: fn.params.map(p => ({ kind: "var", name: p.name })) };
+            const ensures = fn.ensures.map(e => replaceVar(transformExpr(e), "res", fnCall));
+            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),
+                ensures,
+                decreases: fn.decreases ? transformExpr(fn.decreases) : null,
+                body,
+            });
+        }
+        else if (fn.forcePure) {
+            // //@ pure but body can't be auto-converted — emit function by method
+            _forofCounters.clear();
+            const methodBody = transformStmts(fn.body, mod.typeDecls);
+            defByMethods.push({
+                kind: "def-by-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: fn.ensures.map(transformExpr),
+                decreases: fn.decreases ? transformExpr(fn.decreases) : null,
+                methodBody,
+            });
+        }
     }
     const base = mod.file.split("/").pop()?.replace(/\.ts$/, "") ?? "module";
     // Types file
@@ -1381,7 +1399,8 @@ export function transformModule(mod, specImport) {
     }
     // Def file: Velvet methods
     // Pure functions get a thin wrapper that calls Pure.fnName
-    const pureDefNames = new Set(pureDefs.map(d => d.name));
+    // def-by-method functions also skip their method wrappers
+    const pureDefNames = new Set([...pureDefs.map(d => d.name), ...defByMethods.map(d => d.name)]);
     const methods = mod.functions.map(fn => {
         const ensures = [];
         for (const e of fn.ensures) {
@@ -1448,7 +1467,7 @@ export function transformModule(mod, specImport) {
             { key: "loom.semantics.termination", value: '"total"' },
             { key: "loom.semantics.choice", value: '"demonic"' },
         ],
-        decls: [...constDecls, ...methods, ...classDecls],
+        decls: [...constDecls, ...defByMethods, ...methods, ...classDecls],
     };
     return { typesFile, defFile };
 }