lemmascript 0.3.2 → 0.4.0

package/tools/dist/resolve.js

@@ -6,72 +6,6 @@
  */
 import { parseTsType } from "./types.js";
 import { parseExpr } from "./specparser.js";
-// ── Raw expression substitution ─────────────────────────────
-let _synVarCounter = 0;
-/**
- * Structural equality for raw field-access chains (var and field nodes only).
- * Exact within a single expression scope — raw IR has no bindings that could
- * cause name collisions (those are introduced by resolve, which runs after).
- */
-function rawExprEquals(a, b) {
-    if (a.kind === "var" && b.kind === "var")
-        return a.name === b.name;
-    if (a.kind === "field" && b.kind === "field")
-        return a.field === b.field && rawExprEquals(a.obj, b.obj);
-    if (a.kind === "call" && b.kind === "call")
-        return rawExprEquals(a.fn, b.fn) && a.args.length === b.args.length && a.args.every((arg, i) => rawExprEquals(arg, b.args[i]));
-    if (a.kind === "index" && b.kind === "index")
-        return rawExprEquals(a.obj, b.obj) && rawExprEquals(a.idx, b.idx);
-    return false;
-}
-/** Return the root variable name of a field-access chain, or null. */
-function rawChainRoot(e) {
-    if (e.kind === "var")
-        return e.name;
-    if (e.kind === "field")
-        return rawChainRoot(e.obj);
-    return null;
-}
-/**
- * Replace all occurrences of `target` in `expr` with `replacement`.
- * Only matches field-access chains (see rawExprEquals). Stops at lambda
- * boundaries that shadow the chain's root variable.
- */
-function substituteRawExpr(expr, target, replacement) {
-    if (rawExprEquals(expr, target))
-        return replacement;
-    const root = rawChainRoot(target);
-    const sub = (e) => substituteRawExpr(e, target, replacement);
-    switch (expr.kind) {
-        case "var":
-        case "num":
-        case "str":
-        case "bool":
-        case "result":
-        case "havoc":
-        case "emptyCollection":
-            return expr;
-        case "binop": return { ...expr, left: sub(expr.left), right: sub(expr.right) };
-        case "unop": return { ...expr, expr: sub(expr.expr) };
-        case "call": return { ...expr, fn: sub(expr.fn), args: expr.args.map(sub) };
-        case "field": return { ...expr, obj: sub(expr.obj) };
-        case "index": return { ...expr, obj: sub(expr.obj), idx: sub(expr.idx) };
-        case "record":
-            return { ...expr, spread: expr.spread ? sub(expr.spread) : null,
-                fields: expr.fields.map(f => ({ ...f, value: sub(f.value) })) };
-        case "arrayLiteral": return { ...expr, elems: expr.elems.map(sub) };
-        case "conditional": return { ...expr, cond: sub(expr.cond), then: sub(expr.then), else: sub(expr.else) };
-        case "nonNull": return { ...expr, expr: sub(expr.expr) };
-        case "forall":
-        case "exists":
-            return { ...expr, body: sub(expr.body) };
-        case "lambda":
-            // Don't cross lambda boundaries that shadow the chain's root variable
-            if (root && expr.params.some(p => p.name === root))
-                return expr;
-            return { ...expr, body: Array.isArray(expr.body) ? expr.body : sub(expr.body) };
-    }
-}
 function lookup(env, name) {
     if (!env)
         return undefined;
@@ -80,6 +14,21 @@ function lookup(env, name) {
 function extend(env, name, ty) {
     return { name, ty, parent: env };
 }
+function asRawAccessPath(e) {
+    if (e.kind === "var")
+        return { rootVar: e.name, fields: [] };
+    if (e.kind === "field") {
+        const inner = asRawAccessPath(e.obj);
+        if (!inner)
+            return null;
+        return { rootVar: inner.rootVar, fields: [...inner.fields, e.field] };
+    }
+    return null;
+}
+function accessPathsEqual(a, b) {
+    return a.rootVar === b.rootVar && a.fields.length === b.fields.length &&
+        a.fields.every((f, i) => f === b.fields[i]);
+}
 function withEnv(ctx, env) {
     return { ...ctx, env };
 }
@@ -106,15 +55,25 @@ function wrapSome(value, optionalTy) {
         args: [value], ty: optionalTy, callKind: "pure",
     };
 }
-/** Detect `v !== undefined` or `undefined !== v` where v: optional<T>.
- *  Handles simple variables, field access chains, and arbitrary expressions.
- *  When `fieldExpr` is returned, callers must use `substituteRawExpr` to narrow.
- *
- *  Does NOT recurse into `&&` — callers that need to detect optional checks
- *  inside `&&` conditions should check `cond.left` explicitly. */
+/** Detect optional checks: `v !== undefined` (positive narrows then-branch),
+ *  `v === undefined` (negative narrows else-branch), or `!v` (equivalent to
+ *  `=== undefined`).
+ *  Returns:
+ *    - simple var: `varName` set, `fieldExpr` unset
+ *    - complex (field chain or call): `fieldExpr` set, `varName` empty
+ *    - inThen: true for `!==` (truthy), false for `===` and `!v` (falsy).
+ *  Does NOT recurse into `&&`. */
 function detectOptionalCheck(cond, ctx) {
-    if (cond.kind !== "binop" || (cond.op !== "!==" && cond.op !== "==="))
-        return null;
+    // `!v` where v is optional — same shape as `v === undefined` (inThen: false).
+    if (cond.kind === "unop" && cond.op === "!") {
+        const inner = classifyOptExpr(cond.expr, ctx);
+        return inner ? { ...inner, inThen: false } : null;
+    }
+    if (cond.kind !== "binop" || (cond.op !== "!==" && cond.op !== "===")) {
+        // Bare optional truthiness: `if (v)` where v: T | undefined — same as `v !== undefined`.
+        const inner = classifyOptExpr(cond, ctx);
+        return inner ? { ...inner, inThen: true } : null;
+    }
     // Identify the expression being checked against undefined
     let optExpr = null;
     if (cond.right.kind === "var" && cond.right.name === "undefined")
@@ -123,24 +82,112 @@ function detectOptionalCheck(cond, ctx) {
         optExpr = cond.right;
     if (!optExpr)
         return null;
-    // Simple variable — env lookup, no substitution needed
-    if (optExpr.kind === "var") {
-        const ty = lookup(ctx.env, optExpr.name);
+    const inner = classifyOptExpr(optExpr, ctx);
+    return inner ? { ...inner, inThen: cond.op === "!==" } : null;
+}
+/** Classify an expression as a simple var or field-chain optional, returning
+ *  the shape needed by detectOptionalCheck (sans inThen). */
+function classifyOptExpr(e, ctx) {
+    if (e.kind === "var") {
+        const ty = lookup(ctx.env, e.name);
         if (!ty || ty.kind !== "optional")
             return null;
-        return { varName: optExpr.name, innerTy: ty.inner, inThen: cond.op === "!==" };
+        return { varName: e.name, innerTy: ty.inner };
     }
-    // Field access chain or arbitrary expression — resolve to check type, needs substitution
-    const resolved = resolveExpr(optExpr, ctx);
-    if (resolved.ty.kind === "optional") {
-        const synVar = optExpr.kind === "field" ? `_narr${_synVarCounter++}` : `_opt${_synVarCounter++}`;
-        return {
-            varName: synVar, innerTy: resolved.ty.inner, inThen: cond.op === "!==",
-            fieldExpr: optExpr, narrowedExpr: resolved,
-        };
+    const resolved = resolveExpr(e, ctx);
+    if (resolved.ty.kind !== "optional")
+        return null;
+    return { varName: "", innerTy: resolved.ty.inner, fieldExpr: e };
+}
+/** Collect all optional narrowings from an early-return condition.
+ *  Handles single checks (x === undefined) and compound || chains
+ *  (x === undefined || y === undefined). */
+function collectEarlyReturnNarrowings(cond, ctx) {
+    if (cond.kind === "binop" && cond.op === "||") {
+        return [...collectEarlyReturnNarrowings(cond.left, ctx), ...collectEarlyReturnNarrowings(cond.right, ctx)];
+    }
+    const narrowed = detectOptionalCheck(cond, ctx);
+    if (narrowed && !narrowed.inThen && !narrowed.fieldExpr) {
+        return [{ varName: narrowed.varName, innerTy: narrowed.innerTy }];
+    }
+    return [];
+}
+/** TExpr → AccessPath. Counterpart to `asRawAccessPath` for resolved trees.
+ *  Used by `extractInAtoms` when pulling atoms out of typed spec expressions. */
+function asTExprAccessPath(e) {
+    if (e.kind === "var")
+        return { rootVar: e.name, fields: [] };
+    if (e.kind === "field") {
+        const inner = asTExprAccessPath(e.obj);
+        if (!inner)
+            return null;
+        return { rootVar: inner.rootVar, fields: [...inner.fields, e.field] };
     }
     return null;
 }
+/** Walk `e` collecting top-level `k in m` atoms where both sides are pure
+ *  access paths and the right side is map-typed. Descends through `&&` only.
+ *  Does NOT descend into `==>`, `||`, negation, `forall`, or `exists` — in
+ *  those positions an atom is only conditionally known (or a premise, not a
+ *  conclusion), so treating it as always-true in the enclosing scope would
+ *  be unsound. */
+function extractInAtoms(e) {
+    if (e.kind === "binop" && e.op === "in" && e.right.ty.kind === "map") {
+        const obj = asTExprAccessPath(e.right);
+        const idx = asTExprAccessPath(e.left);
+        if (obj && idx)
+            return [{ obj, idx }];
+        return [];
+    }
+    if (e.kind === "binop" && e.op === "&&") {
+        return [...extractInAtoms(e.left), ...extractInAtoms(e.right)];
+    }
+    return [];
+}
+/** Extract `k in m` atoms that hold when `e` is *false*. Currently only
+ *  strips an outer `!` and hands the inner to `extractInAtoms`; that covers
+ *  `if (!(k in m)) ...` for the else-branch and early-return patterns.
+ *  De Morgan over `||` / nested `!(a && b)` not handled yet. */
+function extractInAtomsNegated(e) {
+    if (e.kind === "unop" && e.op === "!")
+        return extractInAtoms(e.expr);
+    return [];
+}
+/** Extend a Ctx with `k in m` atoms. Deduplicates against existing atoms. */
+function withInAtoms(ctx, atoms) {
+    if (atoms.length === 0)
+        return ctx;
+    const existing = ctx.narrowedIndices;
+    const added = [];
+    for (const a of atoms) {
+        if (!existing.some(e => accessPathsEqual(e.obj, a.obj) && accessPathsEqual(e.idx, a.idx))) {
+            added.push(a);
+        }
+    }
+    if (added.length === 0)
+        return ctx;
+    return { ...ctx, narrowedIndices: [...existing, ...added] };
+}
+/** Walk an `&&` chain of `e !== undefined` checks, returning a Ctx with all
+ *  narrowings applied. Earlier checks are in scope for later checks (so the
+ *  right side of `&&` sees the left side's narrowings). */
+function collectAndChainNarrowings(cond, ctx) {
+    if (cond.kind === "binop" && cond.op === "&&") {
+        const leftCtx = collectAndChainNarrowings(cond.left, ctx);
+        return collectAndChainNarrowings(cond.right, leftCtx);
+    }
+    const n = detectOptionalCheck(cond, ctx);
+    if (!n || !n.inThen)
+        return ctx;
+    if (!n.fieldExpr) {
+        return withEnv(ctx, extend(ctx.env, n.varName, n.innerTy));
+    }
+    const path = asRawAccessPath(n.fieldExpr);
+    if (path) {
+        return { ...ctx, narrowedPaths: [...ctx.narrowedPaths, { path, narrowedTy: n.innerTy }] };
+    }
+    return ctx;
+}
 /** TS reference types that become value types in Dafny/Lean — const bindings need mutable var. */
 function isRefMutableInTS(ty) {
     return ty.kind === "array" || ty.kind === "map" || ty.kind === "set";
@@ -317,6 +364,34 @@ function inferMethodReturnTy(fn, args, ctx) {
     }
     return { kind: "unknown" };
 }
+/** Look up the type of `field` on `objTy`. Returns `unknown` if not found. */
+function lookupFieldTy(objTy, field, ctx) {
+    if (field === "length" && (objTy.kind === "array" || objTy.kind === "string")) {
+        return { ty: { kind: "nat" }, isDiscriminant: false };
+    }
+    if (field === "size" && (objTy.kind === "map" || objTy.kind === "set")) {
+        return { ty: { kind: "nat" }, isDiscriminant: false };
+    }
+    if (objTy.kind === "user") {
+        const baseTyName = objTy.name.includes("<") ? objTy.name.slice(0, objTy.name.indexOf("<")) : objTy.name;
+        const isDiscriminant = getDiscriminant(ctx, baseTyName) === field;
+        const decl = findDecl(ctx, baseTyName);
+        if (decl?.kind === "record") {
+            const f = decl.fields?.find(f => f.name === field);
+            if (f)
+                return { ty: f.type, isDiscriminant };
+        }
+        if (decl?.kind === "discriminated-union" && decl.variants) {
+            for (const variant of decl.variants) {
+                const f = variant.fields.find(f => f.name === field);
+                if (f)
+                    return { ty: f.type, isDiscriminant };
+            }
+        }
+        return { ty: { kind: "unknown" }, isDiscriminant };
+    }
+    return { ty: { kind: "unknown" }, isDiscriminant: false };
+}
 // ── Resolve expressions ──────────────────────────────────────
 function resolveExpr(e, ctx) {
     switch (e.kind) {
@@ -344,16 +419,12 @@ function resolveExpr(e, ctx) {
         }
         case "binop": {
             let left = resolveExpr(e.left, ctx);
-            // && narrowing: if left is "x !== undefined", narrow x for right side.
-            // Field chains are excluded — resolve can't substitute in sub-expressions;
-            // transform's emitOptionalMatch handles field chains in statement contexts.
+            // && and ==> narrowing: left-side optional checks narrow the right side.
+            // (For ==>, the premise is assumed in the conclusion — same principle.)
             let rightCtx = ctx;
             let rawRight = e.right;
-            if (e.op === "&&") {
-                const narrowed = detectOptionalCheck(e.left, ctx);
-                if (narrowed && narrowed.inThen && !narrowed.fieldExpr) {
-                    rightCtx = withEnv(ctx, extend(ctx.env, narrowed.varName, narrowed.innerTy));
-                }
+            if (e.op === "&&" || e.op === "==>") {
+                rightCtx = collectAndChainNarrowings(e.left, ctx);
             }
             let right = resolveExpr(rawRight, rightCtx);
             if (e.op === "===" || e.op === "!==") {
@@ -389,49 +460,122 @@ function resolveExpr(e, ctx) {
                 fn.obj.ty.elem.kind === "user") {
                 argCtx = { ...ctx, returnTy: fn.obj.ty.elem };
             }
-            const args = coerceCallArgs(rawArgs.map(a => resolveExpr(a, argCtx)), fn, ctx);
-            const ty = inferMethodReturnTy(fn, args, ctx);
+            // Propagate parameter types to arguments for record literal resolution
+            // (enables inline discriminated union construction in function arguments)
+            const paramTypes = fn.kind === "var" && ctx.fnParams.has(fn.name) ? ctx.fnParams.get(fn.name) : null;
+            const args = coerceCallArgs(rawArgs.map((a, i) => {
+                let aCtx = argCtx;
+                if (paramTypes && i < paramTypes.length && paramTypes[i].kind === "user") {
+                    aCtx = { ...aCtx, returnTy: paramTypes[i] };
+                }
+                return resolveExpr(a, aCtx);
+            }), fn, ctx);
+            let ty = inferMethodReturnTy(fn, args, ctx);
+            // For same-file function calls, use the known return type
+            if (ty.kind === "unknown" && fn.kind === "var" && ctx.fnReturns.has(fn.name)) {
+                ty = ctx.fnReturns.get(fn.name);
+            }
             return { kind: "call", fn, args, ty, callKind: classifyCall(e.fn, ctx) };
         }
         case "index": {
             const obj = resolveExpr(e.obj, ctx);
             const idx = resolveExpr(e.idx, ctx);
-            const idxTy = obj.ty.kind === "array" ? obj.ty.elem
-                : obj.ty.kind === "map" ? { kind: "optional", inner: obj.ty.value }
-                    : { kind: "unknown" };
+            // Map bracket access: default to Option<V>. But if the enclosing scope has a
+            // known `k in m` atom matching (obj, idx) — from requires, assert, an enclosing
+            // `if (k in m)`, or a loop invariant — narrow to V. Parallels how `narrowedPaths`
+            // narrows `obj.field.field` under optional-undefined checks.
+            let idxTy;
+            if (obj.ty.kind === "array") {
+                idxTy = obj.ty.elem;
+            }
+            else if (obj.ty.kind === "map") {
+                const objPath = asTExprAccessPath(obj);
+                const idxPath = asTExprAccessPath(idx);
+                const narrowed = objPath && idxPath && ctx.narrowedIndices.some(n => accessPathsEqual(n.obj, objPath) && accessPathsEqual(n.idx, idxPath));
+                idxTy = narrowed ? obj.ty.value : { kind: "optional", inner: obj.ty.value };
+            }
+            else {
+                idxTy = { kind: "unknown" };
+            }
             return { kind: "index", obj, idx, ty: idxTy };
         }
         case "field": {
             const obj = resolveExpr(e.obj, ctx);
             let isDiscriminant = false;
             let ty = { kind: "unknown" };
-            if (e.field === "length" && (obj.ty.kind === "array" || obj.ty.kind === "string")) {
-                ty = { kind: "nat" };
+            // Check narrowed path context (from conditional optional checks).
+            // Applies when the current field-access forms a pure access path AND
+            // that path is in the narrowedPaths list.
+            if (ctx.narrowedPaths.length > 0) {
+                const myPath = asRawAccessPath(e);
+                if (myPath) {
+                    const np = ctx.narrowedPaths.find(n => accessPathsEqual(n.path, myPath));
+                    if (np)
+                        ty = np.narrowedTy;
+                }
             }
-            else if (e.field === "size" && (obj.ty.kind === "map" || obj.ty.kind === "set")) {
-                ty = { kind: "nat" };
+            if (ty.kind === "unknown") {
+                const lookup = lookupFieldTy(obj.ty, e.field, ctx);
+                ty = lookup.ty;
+                isDiscriminant = lookup.isDiscriminant;
             }
-            else if (obj.ty.kind === "user") {
-                if (getDiscriminant(ctx, obj.ty.name) === e.field)
-                    isDiscriminant = true;
-                const decl = findDecl(ctx, obj.ty.name);
-                if (decl?.kind === "record") {
-                    const f = decl.fields?.find(f => f.name === e.field);
-                    if (f)
-                        ty = f.type;
+            return { kind: "field", obj, field: e.field, ty, isDiscriminant };
+        }
+        case "nullish": {
+            // left ?? right — result type is left's inner (when left is optional)
+            // or just left's type, unified with right's type.
+            const left = resolveExpr(e.left, ctx);
+            const right = resolveExpr(e.right, ctx);
+            const ty = left.ty.kind === "optional" ? left.ty.inner : left.ty;
+            return { kind: "nullish", left, right, ty };
+        }
+        case "optChain": {
+            // obj?.<chain> — obj has type Option<T>; we walk the chain stepping
+            // through types from T. The final result is Option<finalStepTy>
+            // (collapsed: if finalStepTy is already optional, we don't double-wrap).
+            // Narrow rewrites this to a someMatch with the chain applied to the binder.
+            const obj = resolveExpr(e.obj, ctx);
+            let stepInTy = obj.ty.kind === "optional" ? obj.ty.inner : obj.ty;
+            const chain = [];
+            for (const step of e.chain) {
+                if (step.kind === "field") {
+                    const fieldTy = lookupFieldTy(stepInTy, step.name, ctx).ty;
+                    chain.push({ kind: "field", name: step.name, ty: fieldTy });
+                    stepInTy = fieldTy;
+                }
+                else if (step.kind === "index") {
+                    const idx = resolveExpr(step.idx, ctx);
+                    const idxTy = stepInTy.kind === "array" ? stepInTy.elem
+                        : stepInTy.kind === "map" ? { kind: "optional", inner: stepInTy.value }
+                            : { kind: "unknown" };
+                    chain.push({ kind: "index", idx, ty: idxTy });
+                    stepInTy = idxTy;
                 }
-                // Also resolve fields from discriminated-union variants
-                if (ty.kind === "unknown" && decl?.kind === "discriminated-union" && decl.variants) {
-                    for (const variant of decl.variants) {
-                        const f = variant.fields.find(f => f.name === e.field);
-                        if (f) {
-                            ty = f.type;
-                            break;
-                        }
+                else {
+                    // call: prev step yielded a callable (typically a method via field).
+                    // Build a fake fn TExpr from prev steps to reuse inferMethodReturnTy.
+                    const args = step.args.map(a => resolveExpr(a, ctx));
+                    const lastField = chain.length > 0 && chain[chain.length - 1].kind === "field"
+                        ? chain[chain.length - 1] : null;
+                    let callTy = { kind: "unknown" };
+                    let callKind = "unknown";
+                    if (lastField) {
+                        // Build a synthetic field TExpr with the prior step's input type as obj
+                        // so inferMethodReturnTy can dispatch on the receiver type.
+                        const priorInTy = chain.length >= 2 ? chain[chain.length - 2].ty
+                            : (obj.ty.kind === "optional" ? obj.ty.inner : obj.ty);
+                        const fakeObj = { kind: "var", name: "_chain_recv", ty: priorInTy };
+                        const fakeFn = { kind: "field", obj: fakeObj, field: lastField.name, ty: lastField.ty };
+                        callTy = inferMethodReturnTy(fakeFn, args, ctx);
+                        callKind = "method";
                     }
+                    chain.push({ kind: "call", args, ty: callTy, callKind });
+                    stepInTy = callTy;
                 }
             }
-            return { kind: "field", obj, field: e.field, ty, isDiscriminant };
+            const finalTy = stepInTy;
+            const ty = finalTy.kind === "optional" ? finalTy : { kind: "optional", inner: finalTy };
+            return { kind: "optChain", obj, chain, ty };
         }
         case "record": {
             const spread = e.spread ? resolveExpr(e.spread, ctx) : null;
@@ -442,8 +586,13 @@ function resolveExpr(e, ctx) {
             // Clear returnTy for field values — it applies to THIS record, not nested ones
             const fieldCtx = recordTy ? { ...ctx, returnTy: { kind: "unknown" } } : ctx;
             const fields = e.fields.map(f => {
-                let value = resolveExpr(f.value, fieldCtx);
                 const fieldDecl = decl?.fields?.find(df => df.name === f.name);
+                // Propagate declared field type into context so nested records resolve
+                // their union variant correctly (e.g., { kind: 'Idle' } → EffectMode.Idle)
+                const valueCtx = (fieldDecl?.type?.kind === "user")
+                    ? { ...fieldCtx, returnTy: fieldDecl.type }
+                    : fieldCtx;
+                let value = resolveExpr(f.value, valueCtx);
                 if (fieldDecl) {
                     const declTy = fieldDecl.type;
                     value = coerceStr(value, declTy);
@@ -498,56 +647,43 @@ function resolveExpr(e, ctx) {
         }
         case "conditional": {
             const cond = resolveExpr(e.cond, ctx);
-            let narrowedVar;
-            let narrowedExprResolved;
-            let thenCtx = ctx;
-            let rawThen = e.then;
-            // Phase 1: Optional truthiness — cond itself is optional (e.g. opt ? X : Y)
-            if (cond.ty.kind === "optional") {
-                const innerTy = cond.ty.inner;
-                if (e.cond.kind === "var") {
-                    narrowedVar = e.cond.name;
-                    thenCtx = withEnv(ctx, extend(ctx.env, e.cond.name, innerTy));
-                }
-                else {
-                    narrowedVar = `_opt${_synVarCounter++}`;
-                    rawThen = substituteRawExpr(e.then, e.cond, { kind: "var", name: narrowedVar });
-                    thenCtx = withEnv(ctx, extend(ctx.env, narrowedVar, innerTy));
-                }
+            // Type narrowing for the then/else branches. Following TS, we narrow
+            // simple vars and any pure access path (`a.b.c.d`) — but not expressions
+            // with method calls or index ops (bind-first required).
+            // For &&-chains, all positive checks narrow the then-branch; earlier
+            // checks are in scope when resolving later ones.
+            let thenCtx = collectAndChainNarrowings(e.cond, ctx);
+            let elseCtx = ctx;
+            // Truthiness — cond itself is optional (`opt ? a : b`), only for simple vars.
+            if (cond.ty.kind === "optional" && e.cond.kind === "var") {
+                thenCtx = withEnv(thenCtx, extend(thenCtx.env, e.cond.name, cond.ty.inner));
             }
-            // Phase 2: Explicit check — v !== undefined (simple vars, field chains,
-            // complex expressions all handled uniformly by detectOptionalCheck)
-            if (!narrowedVar) {
-                const narrowed = detectOptionalCheck(e.cond, ctx)
-                    // Phase 3: && with optional check — (v !== undefined && ...) ? ... : ...
-                    ?? (e.cond.kind === "binop" && e.cond.op === "&&" ? detectOptionalCheck(e.cond.left, ctx) : null);
-                if (narrowed && narrowed.inThen) {
-                    narrowedVar = narrowed.varName;
-                    thenCtx = withEnv(ctx, extend(ctx.env, narrowed.varName, narrowed.innerTy));
-                    if (narrowed.fieldExpr) {
-                        narrowedExprResolved = narrowed.narrowedExpr ?? resolveExpr(narrowed.fieldExpr, ctx);
-                        rawThen = substituteRawExpr(e.then, narrowed.fieldExpr, { kind: "var", name: narrowed.varName });
-                    }
+            // Single === undefined check narrows the else-branch.
+            const single = detectOptionalCheck(e.cond, ctx);
+            if (single && !single.inThen && !single.fieldExpr) {
+                elseCtx = withEnv(elseCtx, extend(elseCtx.env, single.varName, single.innerTy));
+            }
+            if (!single && e.cond.kind === "binop" && e.cond.op === "||") {
+                for (const n of collectEarlyReturnNarrowings(e.cond, ctx)) {
+                    elseCtx = withEnv(elseCtx, extend(elseCtx.env, n.varName, n.innerTy));
                 }
             }
-            let then_ = resolveExpr(rawThen, thenCtx);
-            let else_ = resolveExpr(e.else, ctx);
+            // Map-index narrowing: `k in m` in the cond narrows the then-branch; `!(k in m)`
+            // narrows the else-branch. Parallel to the if-statement hook in resolveStmt.
+            thenCtx = withInAtoms(thenCtx, extractInAtoms(cond));
+            elseCtx = withInAtoms(elseCtx, extractInAtomsNegated(cond));
+            let then_ = resolveExpr(e.then, thenCtx);
+            let else_ = resolveExpr(e.else, elseCtx);
             then_ = coerceStr(then_, else_.ty);
             else_ = coerceStr(else_, then_.ty);
             let ty = then_.ty.kind !== "unknown" ? then_.ty : else_.ty;
-            // When one branch is undefined, result is optional
             if (then_.ty.kind === "void" && else_.ty.kind !== "void" && else_.ty.kind !== "unknown") {
                 ty = { kind: "optional", inner: else_.ty };
             }
             else if (else_.ty.kind === "void" && then_.ty.kind !== "void" && then_.ty.kind !== "unknown") {
                 ty = { kind: "optional", inner: then_.ty };
             }
-            // When narrowedExpr is set AND a branch is void, the match produces Optional
-            const hasVoidBranch = then_.ty.kind === "void" || else_.ty.kind === "void";
-            if (narrowedExprResolved && hasVoidBranch && ty.kind !== "optional") {
-                ty = { kind: "optional", inner: ty };
-            }
-            return { kind: "conditional", cond, then: then_, else: else_, ty, narrowedVar, narrowedExpr: narrowedExprResolved };
+            return { kind: "conditional", cond, then: then_, else: else_, ty };
         }
         case "emptyCollection": {
             const ty = parseTsType(e.tsType);
@@ -580,17 +716,34 @@ function splitConj(e) {
 function resolveBlock(stmts, ctx) {
     const result = [];
     let env = ctx.env;
+    let narrowedIndices = ctx.narrowedIndices;
     for (const s of stmts) {
-        const [typed, nextEnv] = resolveStmt(s, withEnv(ctx, env));
+        const currentCtx = { ...ctx, env, narrowedIndices };
+        const [typed, nextEnv] = resolveStmt(s, currentCtx);
         result.push(typed);
         env = nextEnv;
         // Flow narrowing: if (x === undefined) { return } narrows x for rest of block.
+        // Also handles compound: if (x === undefined || y === undefined) { return }
         // Field chains are excluded — resolve can't substitute in statement lists;
         // transform's emitOptionalMatch handles field chains in statement contexts.
         if (s.kind === "if" && s.then.length > 0 && s.then[s.then.length - 1].kind === "return" && s.else.length === 0) {
-            const narrowed = detectOptionalCheck(s.cond, withEnv(ctx, env));
-            if (narrowed && !narrowed.inThen && !narrowed.fieldExpr) {
-                env = extend(env, narrowed.varName, narrowed.innerTy);
+            const narrowings = collectEarlyReturnNarrowings(s.cond, withEnv(ctx, env));
+            for (const n of narrowings) {
+                env = extend(env, n.varName, n.innerTy);
+            }
+            // Map-index narrowing: `if (!(k in m)) return;` means `k in m` holds in rest.
+            if (typed.kind === "if") {
+                const addedAtoms = extractInAtomsNegated(typed.cond);
+                if (addedAtoms.length > 0) {
+                    narrowedIndices = withInAtoms({ ...ctx, narrowedIndices }, addedAtoms).narrowedIndices;
+                }
+            }
+        }
+        // Assert narrowing: `//@ assert k in m` adds atoms for the rest of the block.
+        if (typed.kind === "assert") {
+            const addedAtoms = extractInAtoms(typed.expr);
+            if (addedAtoms.length > 0) {
+                narrowedIndices = withInAtoms({ ...ctx, narrowedIndices }, addedAtoms).narrowedIndices;
             }
         }
     }
@@ -600,7 +753,10 @@ function resolveStmt(s, ctx) {
     switch (s.kind) {
         case "let": {
             const declTy = resolveTsType(s.tsType, ctx.overrides, s.name);
-            const init = coerceStr(resolveExpr(s.init, ctx), declTy);
+            // Propagate declared type as returnTy so nested record expressions
+            // resolve union variants correctly (e.g., EffectState → mode: EffectMode → { kind: 'Idle' })
+            const initCtx = declTy.kind === "user" ? { ...ctx, returnTy: declTy } : ctx;
+            const init = coerceStr(resolveExpr(s.init, initCtx), declTy);
             // Map indexing: TS says T, but access can fail → use Optional<T> from init
             const ty = (declTy.kind !== "optional" && init.ty.kind === "optional") ? init.ty : declTy;
             // const collections are mutable in value-semantics world (TS mutates in place, Dafny/Lean reassign)
@@ -629,30 +785,36 @@ function resolveStmt(s, ctx) {
             return [{ kind: "expr", expr: resolveExpr(s.expr, ctx) }, ctx.env];
         case "if": {
             // Narrow optional<T> → T when checking !== undefined or undefined !==.
-            // Also checks left side of && conditions: if (x !== undefined && ...) { ... }
-            // Field chains are excluded — resolve can't substitute in statement bodies;
-            // transform's emitOptionalMatch handles field chains in statement contexts.
-            let thenCtx = ctx, elseCtx = ctx;
-            const narrowed = detectOptionalCheck(s.cond, ctx)
-                ?? (s.cond.kind === "binop" && s.cond.op === "&&" ? detectOptionalCheck(s.cond.left, ctx) : null);
-            if (narrowed && !narrowed.fieldExpr) {
-                const env = extend(ctx.env, narrowed.varName, narrowed.innerTy);
-                if (narrowed.inThen)
-                    thenCtx = withEnv(ctx, env);
-                else
-                    elseCtx = withEnv(ctx, env);
+            // For &&-chains, all positive optional checks narrow the then-branch;
+            // earlier checks are in scope when resolving later ones.
+            // Single-check === undefined narrows the else-branch.
+            let thenCtx = collectAndChainNarrowings(s.cond, ctx);
+            let elseCtx = ctx;
+            const single = detectOptionalCheck(s.cond, ctx);
+            if (single && !single.inThen && !single.fieldExpr) {
+                elseCtx = withEnv(ctx, extend(ctx.env, single.varName, single.innerTy));
             }
-            return [{ kind: "if", cond: resolveExpr(s.cond, ctx), then: resolveBlock(s.then, thenCtx), else: resolveBlock(s.else, elseCtx) }, ctx.env];
+            // Narrow map index access across `k in m` / `!(k in m)` in the cond:
+            // positive atoms (from `k in m` or &&-chains containing it) → then-branch;
+            // negated atoms (from `!(k in m)`) → else-branch.
+            const resolvedCond = resolveExpr(s.cond, ctx);
+            thenCtx = withInAtoms(thenCtx, extractInAtoms(resolvedCond));
+            elseCtx = withInAtoms(elseCtx, extractInAtomsNegated(resolvedCond));
+            return [{ kind: "if", cond: resolvedCond, then: resolveBlock(s.then, thenCtx), else: resolveBlock(s.else, elseCtx) }, ctx.env];
         }
         case "while": {
             const whileSpecCtx = { ...ctx, inSpec: true };
+            const resolvedInvariants = resolveSpecs(s.invariants, whileSpecCtx);
+            // Invariants hold at the top of the body, so any `k in m` atoms among them
+            // narrow map index access in the body.
+            const bodyCtx = withInAtoms(ctx, resolvedInvariants.flatMap(extractInAtoms));
             return [{
                     kind: "while",
                     cond: resolveExpr(s.cond, ctx),
-                    invariants: resolveSpecs(s.invariants, whileSpecCtx),
+                    invariants: resolvedInvariants,
                     decreases: s.decreases ? resolveSpec(s.decreases, whileSpecCtx) : null,
                     doneWith: s.doneWith ? resolveSpec(s.doneWith, whileSpecCtx) : null,
-                    body: resolveBlock(s.body, ctx),
+                    body: resolveBlock(s.body, bodyCtx),
                 }, ctx.env];
         }
         case "forof": {
@@ -833,6 +995,8 @@ function collectCallsStmts(stmts, fns, out) {
 }
 function computePureFns(functions) {
     const allFnNames = new Set(functions.map(fn => fn.name));
+    // //@ pure functions are always considered pure — never taint callers
+    const forcePure = new Set(functions.filter(fn => fn.pure).map(fn => fn.name));
     // Build call graph: fn → set of same-file functions it calls
     const callGraph = new Map();
     for (const fn of functions) {
@@ -840,8 +1004,8 @@ function computePureFns(functions) {
         collectCallsStmts(fn.body, allFnNames, calls);
         callGraph.set(fn.name, calls);
     }
-    // Seed: syntactically non-pure functions
-    const nonPure = new Set(functions.filter(fn => !isSyntacticallyPure(fn.body)).map(fn => fn.name));
+    // Seed: syntactically non-pure functions (skip //@ pure)
+    const nonPure = new Set(functions.filter(fn => !forcePure.has(fn.name) && !isSyntacticallyPure(fn.body)).map(fn => fn.name));
     // Build reverse graph: fn → set of functions that call it
     const callers = new Map();
     for (const name of allFnNames)
@@ -850,12 +1014,12 @@ function computePureFns(functions) {
         for (const callee of callees)
             callers.get(callee).add(caller);
     }
-    // Propagate impurity through reverse call graph
+    // Propagate impurity through reverse call graph (skip //@ pure)
     const worklist = [...nonPure];
     while (worklist.length > 0) {
         const fn = worklist.pop();
         for (const caller of callers.get(fn) ?? []) {
-            if (!nonPure.has(caller)) {
+            if (!nonPure.has(caller) && !forcePure.has(caller)) {
                 nonPure.add(caller);
                 worklist.push(caller);
             }
@@ -888,7 +1052,7 @@ function containsReturn(stmts) {
     return false;
 }
 // ── Resolve function / module ────────────────────────────────
-function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), opts) {
+function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), opts) {
     const overrides = new Map(fn.typeAnnotations.map(a => [a.name, a.type]));
     const params = fn.params.map(p => ({ name: p.name, ty: resolveTsType(p.tsType, overrides, p.name) }));
     const returnTy = resolveTsType(fn.returnType, overrides, "\\result");
@@ -897,7 +1061,7 @@ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), opts) {
         env = extend(env, opts.thisBinding.name, opts.thisBinding.ty);
     for (const p of params)
         env = extend(env, p.name, p.ty);
-    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, inSpec: false, inLambda: false };
+    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
     const requiresCtx = { ...baseCtx, inSpec: true };
     const ensuresCtx = { ...baseCtx, allowResult: true, inSpec: true };
     // Apply type parameter constraints from //@ type T (==) annotations
@@ -905,21 +1069,31 @@ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), opts) {
         const constraint = overrides.get(tp);
         return constraint ? `${tp}${constraint}` : tp;
     });
+    // Resolve requires first so we can extract any `k in m` atoms and seed them
+    // into the body context and the ensures context — they hold for the whole
+    // body (pure fns) and for post-state references in the ensures (map params
+    // aren't mutated through their binding in the Dafny translation).
+    const resolvedRequires = resolveSpecs(fn.requires, requiresCtx);
+    const requiresAtoms = resolvedRequires.flatMap(extractInAtoms);
+    const bodyCtx = withInAtoms(baseCtx, requiresAtoms);
+    const ensuresCtxNarrowed = withInAtoms(ensuresCtx, requiresAtoms);
     return {
         name: fn.name, typeParams, params, returnTy,
-        requires: resolveSpecs(fn.requires, requiresCtx),
-        ensures: resolveSpecs(fn.ensures, ensuresCtx),
+        requires: resolvedRequires,
+        ensures: resolveSpecs(fn.ensures, ensuresCtxNarrowed),
+        decreases: fn.decreases ? resolveSpec(fn.decreases, requiresCtx) : null,
         isPure: opts?.forcePure !== undefined ? opts.forcePure : pureFns.has(fn.name),
-        body: resolveBlock(fn.body, baseCtx),
+        forcePure: fn.pure,
+        body: resolveBlock(fn.body, bodyCtx),
     };
 }
-function resolveClass(cls, typeDecls, pureFns, fnParams = new Map()) {
+function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map()) {
     const fields = cls.fields.map(f => ({ name: f.name, ty: parseTsType(f.tsType) }));
     // Create a synthetic record type for 'this' so field access resolves
     const thisType = { kind: "user", name: cls.name };
     const thisDecl = { name: cls.name, kind: "record", fields: cls.fields.map(f => ({ name: f.name, tsType: f.tsType, type: parseTsType(f.tsType) })) };
     const allTypeDecls = [...typeDecls, thisDecl];
-    const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, {
+    const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, fnReturns, {
         thisBinding: { name: "this", ty: thisType },
         forcePure: false, // class methods are never pure (they access this)
     }));
@@ -943,13 +1117,15 @@ function precomputeFieldTypes(typeDecls) {
 export function resolveModule(raw) {
     precomputeFieldTypes(raw.typeDecls);
     const pureFns = computePureFns(raw.functions);
-    // Pre-compute function parameter types for optional coercion
+    // Pre-compute function parameter and return types
     const fnParams = new Map();
+    const fnReturns = new Map();
     for (const fn of raw.functions) {
         const overrides = new Map(fn.typeAnnotations.map(a => [a.name, a.type]));
         fnParams.set(fn.name, fn.params.map(p => resolveTsType(p.tsType, overrides, p.name)));
+        fnReturns.set(fn.name, resolveTsType(fn.returnType, overrides, "\\result"));
     }
-    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, inSpec: false, inLambda: false };
+    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
     const constants = (raw.constants ?? []).map(c => ({
         name: c.name,
         ty: parseTsType(c.tsType),
@@ -959,7 +1135,7 @@ export function resolveModule(raw) {
         file: raw.file,
         typeDecls: raw.typeDecls,
         constants,
-        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams)),
-        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams)),
+        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams, fnReturns)),
+        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams, fnReturns)),
     };
 }