lemmascript 0.3.1 → 0.3.3

package/tools/dist/transform.js

@@ -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) };
@@ -347,9 +350,24 @@ 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 };
+            }
             const wrappedIdx = isArray(e.obj.ty) && !isNat(e.idx.ty) ? { kind: "toNat", expr: idx } : idx;
             return { kind: "index", arr: transformExpr(e.obj), idx: wrappedIdx };
         }
@@ -394,6 +412,15 @@ function lowerExpr(e, binds) {
                 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 };
@@ -476,6 +503,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":
@@ -494,19 +525,16 @@ 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);
-            let thenExpr = lowerExpr(e.then, binds);
-            let elseExpr = lowerExpr(e.else, binds);
-            // Explicit !== undefined with narrowedExpr → match Some/None on the optional expression
+            // Phase 0: Complex expression check (call results, etc.) — resolve substituted
+            // and set narrowedVar + narrowedExpr because transform can't detect these.
             if (e.narrowedVar && e.narrowedExpr) {
                 const scrutinee = lowerExpr(e.narrowedExpr, binds);
                 const bound = matchBinder(e.narrowedVar);
+                let thenExpr = lowerExpr(e.then, binds);
+                let elseExpr = lowerExpr(e.else, binds);
                 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);
@@ -519,14 +547,16 @@ function lowerExpr(e, binds) {
                     ],
                 };
             }
-            // Optional cond with narrowedVar → match Some/None (truthiness)
+            // Phase 1: Truthiness — cond itself is optional (e.g. opt ? X : Y)
+            // Uses narrowedVar set by resolve's Phase 1 (unchanged).
             if (e.narrowedVar && e.cond.ty.kind === "optional") {
+                const cond = lowerExpr(e.cond, binds);
+                let thenExpr = lowerExpr(e.then, binds);
+                let elseExpr = lowerExpr(e.else, binds);
                 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 {
@@ -537,7 +567,26 @@ function lowerExpr(e, binds) {
                     ],
                 };
             }
-            // Non-optional: regular if with optional wrapping
+            // Phase 2: Explicit check — x !== undefined ? A : B
+            // Transform detects the pattern itself (no narrowedVar/narrowedExpr from resolve).
+            const check = parseOptionalCheck(e.cond);
+            if (check && !check.negated) {
+                return lowerOptionalConditional(e, check, null, binds);
+            }
+            // Phase 3: && with optional check — x !== undefined && guard(x) ? A : B
+            if (e.cond.kind === "binop" && e.cond.op === "&&") {
+                const extracted = extractLeftmostOptional(e.cond);
+                if (extracted) {
+                    const innerCheck = parseOptionalCheck(extracted.optCond);
+                    if (innerCheck && !innerCheck.negated) {
+                        return lowerOptionalConditional(e, innerCheck, extracted.rest, binds);
+                    }
+                }
+            }
+            // Phase 4: Regular conditional (no optional narrowing)
+            const cond = lowerExpr(e.cond, binds);
+            let thenExpr = lowerExpr(e.then, binds);
+            let elseExpr = lowerExpr(e.else, binds);
             if (e.ty.kind === "optional") {
                 thenExpr = wrapOptionalBranch(thenExpr, e.then);
                 elseExpr = wrapOptionalBranch(elseExpr, e.else);
@@ -637,18 +686,43 @@ 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 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" };
+                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];
                 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" };
@@ -770,6 +844,21 @@ function transformStmt(s, typeDecls) {
                     return stmts;
                 }
             }
+            // Desugar: const x = optCheck && guard ? A : B
+            // → var x := B; if (optCheck && guard) { x := A; }
+            // This avoids putting method calls (from guard) inside a match expression,
+            // which Dafny doesn't allow.  The if-case in transformStmt handles the &&
+            // via extractLeftmostOptional → emitOptionalMatch (statement-level match).
+            if (s.init.kind === "conditional" && s.init.cond.kind === "binop" && s.init.cond.op === "&&") {
+                const extracted = extractLeftmostOptional(s.init.cond);
+                if (extracted) {
+                    const desugared = [
+                        { kind: "let", name: s.name, ty: s.ty, mutable: true, init: s.init.else },
+                        { kind: "if", cond: s.init.cond, then: [{ kind: "assign", target: s.name, value: s.init.then }], else: [] },
+                    ];
+                    return desugared.flatMap(ds => transformStmt(ds, typeDecls));
+                }
+            }
             const { binds, expr } = liftMethodCalls(s.init);
             return [...binds, { kind: "let", name: s.name, type: s.ty, mutable: s.mutable, value: expr }];
         }
@@ -917,7 +1006,8 @@ function emitOptionalMatch(varName, negated, s, typeDecls, restStmts, fieldExpr)
     if (someBranch.length === 0 && restStmts && restStmts.length > 0) {
         someBranch = restStmts;
     }
-    const bound = matchBinder(`${varName}_val`);
+    const sanitized = varName.replace(/\./g, "_");
+    const bound = matchBinder(`${sanitized}_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).
@@ -931,10 +1021,14 @@ function emitOptionalMatch(varName, negated, s, typeDecls, restStmts, fieldExpr)
     }
     else {
         const transformed = transformStmts(someBranch, typeDecls);
-        someBody = transformed.map(stmt => mapStmtExprs(stmt, e => replaceVar(e, varName, { kind: "var", name: bound })));
+        someBody = transformed.map(stmt => mapStmtExprs(stmt, e => replaceVar(e, varName, { kind: "var", name: bound }, true)));
     }
+    // For field chains, use an Expr scrutinee so outer match replaceVar can
+    // substitute the object variable (e.g. task → i_task_val in task.deletedFromList).
+    // String scrutinees are opaque to replaceVar/mapExpr.
+    const scrutinee = fieldExpr ? transformExpr(fieldExpr) : varName;
     return {
-        kind: "match", scrutinee: varName,
+        kind: "match", scrutinee,
         arms: [
             { pattern: `.some ${bound}`, body: someBody },
             { pattern: ".none", body: noneBranch.length > 0 ? transformStmts(noneBranch, typeDecls) : [] },
@@ -947,25 +1041,21 @@ function mapStmtExprs(s, r) {
 }
 // ── Optional narrowing helpers ──────────────────────────────
 //
-// Optional narrowing converts TS `if (x === undefined)` patterns to Dafny
-// `match x { Some(val) => ..., None => ... }`.
+// Optional narrowing converts TS `if (x === undefined)` / `x !== undefined ? a : b`
+// patterns to `match x { Some(val) => ..., None => ... }`.
 //
-// The resolve phase (resolve.ts) handles:
+// The resolve phase (resolve.ts) handles TYPE narrowing only:
 //   - 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
+//   - Conditional type narrowing: extends env (simple vars) or narrowedFields context
+//     (field chains) so the then-branch resolves with the unwrapped type
 //
-// The transform phase (here) handles:
+// The transform phase (here) handles ALL structural narrowing:
 //   - 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.
+//   - Expression-level: `lowerExpr` detects optional checks via `parseOptionalCheck`
+//     and `extractLeftmostOptional` → `lowerOptionalConditional`
+//   - && restructuring: `extractLeftmostOptional` splits `&&` chains, generating
+//     match with guard: `match x { Some(val) => if guard then A else B, None => B }`
 /** 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. */
@@ -978,7 +1068,8 @@ function prepareOptionalMatch(s, 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`);
+    const sanitized = check.varName.replace(/\./g, "_");
+    const bound = matchBinder(`${sanitized}_val`);
     return { check, someBranch, noneBranch, bound };
 }
 /** Extract the leftmost optional check from a && chain, returning the check and the rest.
@@ -1092,6 +1183,96 @@ 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;
+    });
+}
+/** 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);
+}
+/** Lower a conditional with an optional check (Phase 2: explicit, Phase 3: && with guard).
+ *  Generates: match scrutinee { Some(val) => [if guard then] A [else B], None => B }
+ *  For field chains, replaces field accesses in TExpr before lowering.
+ *  For simple vars, replaces variable names in Expr after lowering. */
+function lowerOptionalConditional(e, check, guard, binds) {
+    const sanitized = check.varName.replace(/\./g, "_");
+    const bound = matchBinder(`${sanitized}_val`);
+    const isFieldChain = check.fieldExpr && check.fieldExpr.kind === "field" &&
+        check.fieldExpr.obj.kind === "var";
+    // Determine which branch is Some (unwrapped) vs None
+    let thenTExpr = e.then;
+    let guardTExpr = guard;
+    // For field chains: replace field access with bound var in TExpr before lowering
+    if (isFieldChain) {
+        const fe = check.fieldExpr;
+        const innerTy = fe.ty.kind === "optional" ? fe.ty.inner : fe.ty;
+        const replacements = [{ fieldName: fe.field, newName: bound, fallbackTy: innerTy }];
+        thenTExpr = fixBoundType(replaceFieldInTExpr(thenTExpr, fe.obj.name, replacements), bound);
+        if (guardTExpr) {
+            guardTExpr = fixBoundType(replaceFieldInTExpr(guardTExpr, fe.obj.name, replacements), bound);
+        }
+    }
+    let thenExpr = lowerExpr(thenTExpr, binds);
+    let elseExpr = lowerExpr(e.else, binds);
+    // For simple vars: replace after lowering
+    if (!isFieldChain) {
+        thenExpr = replaceVar(thenExpr, check.varName, { kind: "var", name: bound }, true);
+    }
+    // Optional wrapping: check if either branch is undefined (produces optional result)
+    const isOptionalResult = (e.then.kind === "var" && e.then.name === "undefined") ||
+        (e.else.kind === "var" && e.else.name === "undefined");
+    if (isOptionalResult) {
+        thenExpr = wrapOptionalBranch(thenExpr, e.then);
+        elseExpr = wrapOptionalBranch(elseExpr, e.else);
+    }
+    // Build Some arm body — add guard for && patterns
+    // Note: if the guard has method calls, the impure path desugars the let to a
+    // statement-level if+match in transformStmt, so this expression-level path only
+    // runs for pure guards. Use null for binds to avoid escaping the match scope.
+    let someBody;
+    if (guard) {
+        let guardExpr = lowerExpr(guardTExpr, null);
+        if (!isFieldChain) {
+            guardExpr = replaceVar(guardExpr, check.varName, { kind: "var", name: bound }, true);
+        }
+        // Guard-else gets the same expression as None arm
+        let guardElse = lowerExpr(e.else, null);
+        if (isOptionalResult) {
+            guardElse = wrapOptionalBranch(guardElse, e.else);
+        }
+        someBody = { kind: "if", cond: guardExpr, then: thenExpr, else: guardElse };
+    }
+    else {
+        someBody = thenExpr;
+    }
+    // Build scrutinee
+    const scrutinee = isFieldChain
+        ? lowerExpr(check.fieldExpr, binds)
+        : check.varName;
+    return {
+        kind: "match", scrutinee,
+        arms: [
+            { pattern: `.some ${bound}`, body: someBody },
+            { pattern: ".none", body: elseExpr },
+        ],
+    };
+}
 // ── Pure function generation ─────────────────────────────────
 function transformPureBody(stmts, typeDecls) {
     // Detect discriminant if-chain
@@ -1112,16 +1293,42 @@ function transformPureBody(stmts, typeDecls) {
                 return { kind: "let", name: s.name, value: transformExpr(s.init), body: restExpr };
             }
             case "if": {
+                // Restructure && with optional check: split into nested ifs so
+                // prepareOptionalMatch can detect the optional check and bind the unwrapped value
+                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 transformPureBody([outerIf, ...rest], typeDecls);
+                    }
+                }
                 // Optional narrowing: if (x === undefined) → match x { None => ..., Some(x_val) => ... }
                 const optMatch = prepareOptionalMatch(s, rest);
                 if (optMatch) {
-                    const someExpr = transformPureBody(optMatch.someBranch, typeDecls);
+                    // For field chains (a.dueDate !== undefined), replace in TStmt before transform
+                    let someBranch = [...optMatch.someBranch, ...rest];
+                    const fe = optMatch.check.fieldExpr;
+                    if (fe && fe.kind === "field" && fe.obj.kind === "var") {
+                        const innerTy = fe.ty.kind === "optional" ? fe.ty.inner : fe.ty;
+                        someBranch = replaceFieldsInTStmts(someBranch, fe.obj.name, [
+                            { fieldName: fe.field, newName: optMatch.bound, fallbackTy: innerTy },
+                        ]);
+                        // The replacement keeps the original optional type, but the match binding
+                        // unwraps it. Fix the type so downstream record coercion re-wraps with Some().
+                        someBranch = someBranch.map(s => mapTStmt(s, e => e.kind === "var" && e.name === optMatch.bound && e.ty.kind === "optional"
+                            ? { ...e, ty: e.ty.inner } : null));
+                    }
+                    const someExpr = transformPureBody(someBranch, typeDecls);
                     if (!someExpr)
                         return null;
                     const noneExpr = transformPureBody(optMatch.noneBranch, typeDecls);
                     if (!noneExpr)
                         return null;
-                    const someReplaced = replaceVar(someExpr, optMatch.check.varName, { kind: "var", name: optMatch.bound });
+                    // For simple vars, replace in Expr after transform
+                    const someReplaced = optMatch.check.fieldExpr
+                        ? someExpr
+                        : replaceVar(someExpr, optMatch.check.varName, { kind: "var", name: optMatch.bound }, true);
                     return {
                         kind: "match", scrutinee: optMatch.check.varName,
                         arms: [
@@ -1130,11 +1337,13 @@ function transformPureBody(stmts, typeDecls) {
                         ],
                     };
                 }
-                const thenExpr = transformPureBody(s.then, 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 elseBranch = s.else.length > 0 ? s.else : rest;
-                const elseExpr = transformPureBody(elseBranch, typeDecls);
+                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 };
@@ -1265,17 +1474,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;
     });
 }
@@ -1316,25 +1546,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
@@ -1353,7 +1601,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) {
@@ -1420,7 +1669,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 };
 }