lemmascript 0.3.3 → 0.5.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,30 @@ function lookup(env, name) {
 function extend(env, name, ty) {
     return { name, ty, parent: env };
 }
+function envKeys(env) {
+    const out = [];
+    let e = env;
+    while (e) {
+        out.push(e.name);
+        e = e.parent;
+    }
+    return out;
+}
+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,17 +64,62 @@ 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 for a simple field chain (obj.field), callers
- *  can narrow via `narrowedFields` context. For complex expressions (calls, etc.),
- *  callers should fall back to `substituteRawExpr`.
- *
- *  Does NOT recurse into `&&` — callers that need to detect optional checks
- *  inside `&&` conditions should check `cond.left` explicitly. */
+/** Find the synth array-union TypeDecl named `name` (discriminant `__isArray__`). */
+function findSynthArrayUnion(name, typeDecls) {
+    const decl = typeDecls.find(d => d.name === name);
+    if (decl?.kind === "discriminated-union" && decl.discriminant === "__isArray__")
+        return decl;
+    return null;
+}
+/** Coerce `value` to `targetTy` at an assignment-position. Mirrors TS subtyping
+ *  for the two upcast shapes LS synthesizes:
+ *    - `T` into `optional<T>` slot  → wrap with `Some(...)`
+ *    - `T[]` into a synth `T[] | U` slot → wrap with `ArrayBranch(...)`
+ *    - `U`   into a synth `T[] | U` slot → wrap with `NonArrayBranch(...)`
+ *  Returns `value` unchanged if no coercion applies (types already match,
+ *  source is unknown, or no rule matches). */
+function coerceToTargetTy(value, targetTy, typeDecls) {
+    if (value.ty.kind === "unknown" || value.ty.kind === "void")
+        return value;
+    if (targetTy.kind === "optional" && value.ty.kind !== "optional") {
+        return wrapSome(value, targetTy);
+    }
+    if (targetTy.kind === "user") {
+        const synth = findSynthArrayUnion(targetTy.name, typeDecls);
+        if (synth && synth.variants && synth.variants.length === 2) {
+            const arrVariant = synth.variants.find(v => v.name === "ArrayBranch");
+            const nonVariant = synth.variants.find(v => v.name === "NonArrayBranch");
+            if (value.ty.kind === "array" && arrVariant) {
+                return { kind: "call", fn: { kind: "var", name: "ArrayBranch", ty: targetTy },
+                    args: [value], ty: targetTy, callKind: "pure" };
+            }
+            if (value.ty.kind !== "array" && nonVariant) {
+                return { kind: "call", fn: { kind: "var", name: "NonArrayBranch", ty: targetTy },
+                    args: [value], ty: targetTy, callKind: "pure" };
+            }
+        }
+    }
+    return value;
+}
+/** 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")
@@ -125,23 +128,28 @@ 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,
-        };
+    if (e.kind === "result") {
+        const ty = lookup(ctx.env, "\\result");
+        if (!ty || ty.kind !== "optional")
+            return null;
+        return { varName: "\\result", innerTy: ty.inner };
     }
-    return null;
+    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
@@ -156,16 +164,155 @@ function collectEarlyReturnNarrowings(cond, ctx) {
     }
     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";
 }
 function findDecl(ctx, name) {
-    return ctx.typeDecls.find(d => d.name === name);
+    const direct = ctx.typeDecls.find(d => d.name === name);
+    if (direct)
+        return direct;
+    // Dotted names (e.g. `Agent.Info`, `Permission.Ruleset`): fall back to the
+    // last segment, so `//@ declare-type Info { ... }` matches a reference to
+    // `Agent.Info` without forcing the user to repeat the namespace.
+    const dotIdx = name.lastIndexOf(".");
+    if (dotIdx >= 0)
+        return ctx.typeDecls.find(d => d.name === name.slice(dotIdx + 1));
+    return undefined;
+}
+/** Expand alias-kind typeDecls when the alias target is structural (array,
+ *  map, set, optional, or another user type). Primitive-typed aliases like
+ *  `type TaskId = number` stay as `user("TaskId")` so the generated Dafny
+ *  preserves the alias name. Recursive through compound types; cycle-safe. */
+function expandAlias(ty, typeDecls, seen = new Set()) {
+    if (ty.kind === "user") {
+        if (seen.has(ty.name))
+            return ty;
+        let decl = typeDecls.find(d => d.name === ty.name);
+        if (!decl && ty.name.includes(".")) {
+            const tail = ty.name.slice(ty.name.lastIndexOf(".") + 1);
+            decl = typeDecls.find(d => d.name === tail);
+        }
+        if (decl?.kind === "alias" && decl.aliasOfTy) {
+            const target = decl.aliasOfTy;
+            if (target.kind === "array" || target.kind === "map" || target.kind === "set" || target.kind === "optional" || target.kind === "user") {
+                return expandAlias(target, typeDecls, new Set([...seen, ty.name]));
+            }
+        }
+        return ty;
+    }
+    if (ty.kind === "optional")
+        return { kind: "optional", inner: expandAlias(ty.inner, typeDecls, seen) };
+    if (ty.kind === "array")
+        return { kind: "array", elem: expandAlias(ty.elem, typeDecls, seen) };
+    if (ty.kind === "set")
+        return { kind: "set", elem: expandAlias(ty.elem, typeDecls, seen) };
+    if (ty.kind === "map")
+        return { kind: "map", key: expandAlias(ty.key, typeDecls, seen), value: expandAlias(ty.value, typeDecls, seen) };
+    return ty;
 }
 function getDiscriminant(ctx, typeName) {
     return findDecl(ctx, typeName)?.discriminant;
 }
+/** A type ts-morph handed us that LemmaScript hasn't modeled: contains
+ *  `unknown` (TS `any`), or a `user` type whose name isn't a known declaration
+ *  (an opaque expanded union like `"AssistantMsg | ToolMsg"` that ts-morph
+ *  produced by expanding an alias LS shadows via declare-type). Used by
+ *  `case "let"` to decide when LS's own `init.ty` is the better source of
+ *  structure. */
+function isUnmodeledTy(ty, typeDecls) {
+    if (ty.kind === "unknown")
+        return true;
+    if (ty.kind === "optional")
+        return isUnmodeledTy(ty.inner, typeDecls);
+    if (ty.kind === "array")
+        return isUnmodeledTy(ty.elem, typeDecls);
+    if (ty.kind === "set")
+        return isUnmodeledTy(ty.elem, typeDecls);
+    if (ty.kind === "map")
+        return isUnmodeledTy(ty.key, typeDecls) || isUnmodeledTy(ty.value, typeDecls);
+    if (ty.kind === "user") {
+        const base = ty.name.includes("<") ? ty.name.slice(0, ty.name.indexOf("<")) : ty.name;
+        return !typeDecls.some(d => d.name === base);
+    }
+    return false;
+}
 /** Infer quantifier variable type from usage in body.
  *  If the variable is used as a map/set key (e.g. map.has(k), map.get(k)),
  *  return the collection's key type. Otherwise return null (default to int). */
@@ -232,8 +379,16 @@ function inferQuantVarType(varName, body, ctx) {
 function classifyCall(fn, ctx) {
     if (fn.kind === "field" && fn.obj.kind === "var" && fn.obj.name === "Math")
         return "pure";
+    if (fn.kind === "field" && fn.obj.kind === "var" && fn.obj.name === "Array" && fn.field === "isArray")
+        return "pure";
     if (fn.kind === "var" && (ctx.inSpec || ctx.inLambda) && ctx.pureFns.has(fn.name))
         return "spec-pure";
+    // Bare-name `//@ extern` declarations are emitted as `function {:axiom}` —
+    // pure from the verifier's perspective. Classify them as pure so callers
+    // don't get lifted to statement-level binds (which would force lambdas to
+    // become multi-statement, illegal in Dafny).
+    if (fn.kind === "var" && ctx.externs.has(fn.name))
+        return "pure";
     if (fn.kind === "var" && ctx.inSpec) {
         // Not a known pure function — could be external (Lean-defined spec helper).
         // Pass through as "pure" and let Lean catch any errors.
@@ -244,24 +399,56 @@ function classifyCall(fn, ctx) {
     return "unknown";
 }
 // ── Call resolution helpers ─────────────────────────────────
-/** Infer lambda param types from array method context (map, filter, etc.).
- *  Returns updated rawArgs with inferred tsType on the first lambda param. */
-function inferLambdaParamTypes(fn, rawArgs) {
+/** Infer lambda param types from array method context (map, filter, etc.)
+ *  AND from function-typed parameters of named callees (e.g., a `Comparator =
+ *  (a, b) => bool` parameter propagates `string, string` to the lambda's
+ *  inline params). Returns updated rawArgs with inferred tsType. */
+function tyToTsStr(ty) {
+    if (ty.kind === "user")
+        return ty.name;
+    if (ty.kind === "string")
+        return "string";
+    if (ty.kind === "int" || ty.kind === "nat")
+        return "number";
+    if (ty.kind === "bool")
+        return "boolean";
+    return undefined;
+}
+function inferLambdaParamTypes(fn, rawArgs, ctx) {
     if (fn.kind === "field" && fn.obj.ty.kind === "array" &&
-        ["map", "filter", "every", "some", "find"].includes(fn.field) &&
+        ["map", "filter", "every", "some", "find", "findLast", "findIndex"].includes(fn.field) &&
         rawArgs.length >= 1 && rawArgs[0].kind === "lambda" &&
         rawArgs[0].params.length >= 1 && !rawArgs[0].params[0].tsType) {
         const elemTy = fn.obj.ty.elem;
-        const tsType = elemTy.kind === "user" ? elemTy.name
-            : elemTy.kind === "string" ? "string"
-                : elemTy.kind === "int" || elemTy.kind === "nat" ? "number"
-                    : elemTy.kind === "bool" ? "boolean" : undefined;
+        const tsType = tyToTsStr(elemTy);
         if (tsType) {
             const lam = rawArgs[0];
             const updatedParams = [{ ...lam.params[0], tsType }, ...lam.params.slice(1)];
             return [{ ...lam, params: updatedParams }, ...rawArgs.slice(1)];
         }
     }
+    // Named-callee propagation: when an argument position expects a function
+    // type, infer the lambda's param types from that function type. Aliases
+    // (e.g., `Comparator`) are expanded via the typeDecls.
+    if (fn.kind === "var" && ctx?.fnParams.has(fn.name)) {
+        const paramTys = ctx.fnParams.get(fn.name);
+        return rawArgs.map((a, i) => {
+            if (a.kind !== "lambda" || i >= paramTys.length)
+                return a;
+            let pTy = paramTys[i];
+            if (pTy.kind === "user") {
+                const decl = ctx.typeDecls.find(d => d.name === pTy.name);
+                if (decl?.kind === "alias" && decl.aliasOfTy)
+                    pTy = decl.aliasOfTy;
+                else if (decl?.kind === "alias" && decl.aliasOf)
+                    pTy = parseTsType(decl.aliasOf);
+            }
+            if (pTy.kind !== "fn")
+                return a;
+            const updatedParams = a.params.map((p, idx) => p.tsType || idx >= pTy.params.length ? p : { ...p, tsType: tyToTsStr(pTy.params[idx]) });
+            return { ...a, params: updatedParams };
+        });
+    }
     return rawArgs;
 }
 /** Coerce call arguments: string literals → user types, non-optional → Some, pad missing optional args. */
@@ -290,6 +477,11 @@ function coerceCallArgs(args, fn, ctx) {
 function inferMethodReturnTy(fn, args, ctx) {
     if (fn.kind !== "field")
         return { kind: "unknown" };
+    // `Array.isArray(x)` always returns boolean. narrow.ts recognizes this call as
+    // a discriminator predicate when `x` has type of a synthesized array-union.
+    if (fn.obj.kind === "var" && fn.obj.name === "Array" && fn.field === "isArray") {
+        return { kind: "bool" };
+    }
     const objTy = fn.obj.ty;
     if (objTy.kind === "map") {
         if (fn.field === "get")
@@ -312,12 +504,24 @@ function inferMethodReturnTy(fn, args, ctx) {
             return { kind: "int" };
         if (fn.field === "shift")
             return objTy.elem;
+        if (fn.field === "pop")
+            return { kind: "optional", inner: objTy.elem };
         if (fn.field === "push" || fn.field === "concat")
             return objTy;
         if (fn.field === "filter")
             return objTy;
         if (fn.field === "every" || fn.field === "some")
             return { kind: "bool" };
+        if (fn.field === "find" || fn.field === "findLast")
+            return { kind: "optional", inner: objTy.elem };
+        if (fn.field === "findIndex")
+            return { kind: "int" };
+        if (fn.field === "flat" && objTy.elem.kind === "array")
+            return { kind: "array", elem: objTy.elem.elem };
+        if (fn.field === "slice")
+            return objTy;
+        if (fn.field === "join" && objTy.elem.kind === "string")
+            return { kind: "string" };
         if (fn.field === "map" && args.length >= 1 && args[0].kind === "lambda") {
             const retTy = args[0].body.length > 0 && args[0].body[0].kind === "return"
                 ? args[0].body[0].value.ty : { kind: "unknown" };
@@ -325,13 +529,45 @@ function inferMethodReturnTy(fn, args, ctx) {
         }
     }
     else if (objTy.kind === "string") {
-        if (fn.field === "trim" || fn.field === "toLowerCase" || fn.field === "toUpperCase")
+        if (fn.field === "trim" || fn.field === "trimEnd" || fn.field === "trimStart" || fn.field === "toLowerCase" || fn.field === "toUpperCase")
             return { kind: "string" };
-        if (fn.field === "includes")
+        if (fn.field === "slice" || fn.field === "substring")
+            return { kind: "string" };
+        if (fn.field === "split")
+            return { kind: "array", elem: { kind: "string" } };
+        if (fn.field === "includes" || fn.field === "startsWith" || fn.field === "endsWith")
             return { kind: "bool" };
     }
     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) {
@@ -359,16 +595,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 === "!==") {
@@ -396,8 +628,21 @@ function resolveExpr(e, ctx) {
             return { kind: "unop", op: e.op, expr, ty: e.op === "!" ? { kind: "bool" } : expr.ty };
         }
         case "call": {
+            // Extern dispatch: `NS.method(args)` where NS.method is declared via
+            // `//@ extern`. Rewrite into a flat-name call (`NS_method(args)`) so the
+            // rest of the pipeline sees an ordinary pure function. The extern's
+            // declaration is emitted alongside the file as `function {:axiom} ...`.
+            if (e.fn.kind === "field" && e.fn.obj.kind === "var") {
+                const qualified = `${e.fn.obj.name}.${e.fn.field}`;
+                const ext = ctx.externs.get(qualified);
+                if (ext) {
+                    const args = e.args.map(a => resolveExpr(a, ctx));
+                    const fn = { kind: "var", name: ext.flat, ty: { kind: "unknown" } };
+                    return { kind: "call", fn, args, ty: ext.returnTy, callKind: "pure" };
+                }
+            }
             const fn = resolveExpr(e.fn, ctx);
-            const rawArgs = inferLambdaParamTypes(fn, e.args);
+            const rawArgs = inferLambdaParamTypes(fn, e.args, ctx);
             // For .push() on a typed array, resolve args with element type context
             let argCtx = ctx;
             if (fn.kind === "field" && fn.obj.ty.kind === "array" && fn.field === "push" &&
@@ -407,13 +652,20 @@ function resolveExpr(e, 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 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);
+            // Array method `.with(i, v)`: coerce the value arg to the element type
+            // so `arr[i] = v` on `(T|null)[]` wraps `T` → `Some(T)` (and similarly
+            // for synth array-unions). Same shape as the record-field coercion
+            // below: assigning a narrower value into a wider slot.
+            if (fn.kind === "field" && fn.field === "with" && fn.obj.ty.kind === "array" && args.length === 2) {
+                args = [args[0], coerceToTargetTy(args[1], fn.obj.ty.elem, ctx.typeDecls)];
+            }
             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)) {
@@ -424,56 +676,126 @@ function resolveExpr(e, 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" };
-            // Check narrowed field context (from conditional optional checks on field chains)
-            if (obj.kind === "var" && ctx.narrowedFields.length > 0) {
-                const nf = ctx.narrowedFields.find(n => n.objName === obj.name && n.fieldName === e.field);
-                if (nf)
-                    ty = nf.narrowedTy;
-            }
-            if (ty.kind === "unknown" && 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 (ty.kind === "unknown" && 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 (ty.kind === "unknown" && obj.ty.kind === "user") {
-                // Strip generic args for type lookup: "Result<Model, Err>" → "Result"
-                const baseTyName = obj.ty.name.includes("<") ? obj.ty.name.slice(0, obj.ty.name.indexOf("<")) : obj.ty.name;
-                if (getDiscriminant(ctx, baseTyName) === e.field)
-                    isDiscriminant = true;
-                const decl = findDecl(ctx, baseTyName);
-                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;
                 }
-                // 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 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;
+                }
+                else {
+                    // call: prev step yielded a callable (typically a method via field).
+                    // Build a fake fn TExpr from prev steps to reuse inferMethodReturnTy
+                    // and inferLambdaParamTypes — without the latter, a lambda arg buried
+                    // inside `obj?.filter(r => ...)` gets `int`-typed params instead of
+                    // the array's element type.
+                    const lastField = chain.length > 0 && chain[chain.length - 1].kind === "field"
+                        ? chain[chain.length - 1] : null;
+                    let callTy = { kind: "unknown" };
+                    let callKind = "unknown";
+                    let rawArgs = step.args;
+                    if (lastField) {
+                        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 };
+                        rawArgs = inferLambdaParamTypes(fakeFn, rawArgs);
+                        const args = rawArgs.map(a => resolveExpr(a, ctx));
+                        callTy = inferMethodReturnTy(fakeFn, args, ctx);
+                        callKind = "method";
+                        chain.push({ kind: "call", args, ty: callTy, callKind });
+                        stepInTy = callTy;
+                        continue;
                     }
+                    const args = rawArgs.map(a => resolveExpr(a, ctx));
+                    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;
             const ty = spread ? spread.ty : { kind: "unknown" };
-            // Infer record type: from spread, or from return type context
-            const recordTy = ty.kind === "user" ? ty : ctx.returnTy.kind === "user" ? ctx.returnTy : null;
+            // Record literal in map-typed context (e.g. `const M: Record<string, V> = {a: ...}`):
+            // attach the map type so transform/emit can produce a map literal.
+            if (!spread && ctx.returnTy.kind === "map") {
+                const mapTy = ctx.returnTy;
+                const fieldCtx = { ...ctx, returnTy: mapTy.value };
+                const fields = e.fields.map(f => ({ name: f.name, value: resolveExpr(f.value, fieldCtx) }));
+                return { kind: "record", spread: null, fields, ty: mapTy };
+            }
+            // Infer record type: from spread, or from return type context. Unwrap
+            // an outer Optional when looking at returnTy — `return {...} : null`
+            // has ctx.returnTy = Option<T>, but the record literal's natural type is T.
+            const returnTyUnwrapped = ctx.returnTy.kind === "optional" ? ctx.returnTy.inner : ctx.returnTy;
+            const recordTy = ty.kind === "user" ? ty : returnTyUnwrapped.kind === "user" ? returnTyUnwrapped : null;
             const decl = recordTy ? ctx.typeDecls.find(d => d.name === recordTy.name && d.kind === "record") : undefined;
             // Clear returnTy for field values — it applies to THIS record, not nested ones
             const fieldCtx = recordTy ? { ...ctx, returnTy: { kind: "unknown" } } : ctx;
@@ -492,19 +814,22 @@ function resolveExpr(e, ctx) {
                     if (value.kind === "record" && value.fields.length === 0 && !value.spread && declTy.kind === "map") {
                         value = { kind: "arrayLiteral", elems: [], ty: declTy };
                     }
-                    // Coerce non-optional to optional: wrap in Some (only when value type is concrete)
-                    if (declTy.kind === "optional" && value.ty.kind !== "optional" && value.ty.kind !== "void" && value.ty.kind !== "unknown") {
-                        value = wrapSome(value, declTy);
-                    }
+                    // Assignment-position upcasts: T → Option<T>, T[] → ArrayBranch(T[]),
+                    // U → NonArrayBranch(U). Handles both optional fields and fields
+                    // typed as a synth array-union (`T[] | U`).
+                    value = coerceToTargetTy(value, declTy, ctx.typeDecls);
                 }
                 return { name: f.name, value };
             });
             return { kind: "record", spread, fields, ty: recordTy ?? ty };
         }
         case "result":
+            // \result desugars to a regular var so all the variable-narrowing
+            // machinery (env lookup, optional checks, path matching) just works.
+            // The env in ensuresCtx is pre-seeded with "\result" → returnTy.
             if (!ctx.allowResult)
                 throw new Error("\\result is only valid in ensures");
-            return { kind: "result", ty: ctx.returnTy };
+            return { kind: "var", name: "\\result", ty: lookup(ctx.env, "\\result") ?? ctx.returnTy };
         case "forall": {
             const varTy = e.varType !== "int" ? parseTsType(e.varType)
                 : inferQuantVarType(e.var, e.body, ctx) ?? { kind: "int" };
@@ -539,87 +864,51 @@ function resolveExpr(e, ctx) {
         }
         case "conditional": {
             const cond = resolveExpr(e.cond, ctx);
-            let narrowedVar;
-            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));
-                }
-            }
-            // Phase 2/3: Explicit check — v !== undefined, or && with optional check.
-            // Resolve only narrows the type environment; transform handles all structural
-            // narrowing (match generation, variable binding, && splitting).
-            let narrowedExprResolved;
+            // 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;
-            if (!narrowedVar) {
-                const narrowed = detectOptionalCheck(e.cond, ctx)
-                    ?? (e.cond.kind === "binop" && e.cond.op === "&&" ? detectOptionalCheck(e.cond.left, ctx) : null);
-                if (narrowed && narrowed.inThen) {
-                    if (!narrowed.fieldExpr) {
-                        // Simple var: extend env only — transform will detect and generate match
-                        thenCtx = withEnv(thenCtx, extend(thenCtx.env, narrowed.varName, narrowed.innerTy));
-                    }
-                    else if (narrowed.fieldExpr.kind === "field" && narrowed.fieldExpr.obj.kind === "var") {
-                        // Simple field chain (obj.field): narrow via field context — no substitution
-                        thenCtx = {
-                            ...thenCtx,
-                            narrowedFields: [...thenCtx.narrowedFields, {
-                                    objName: narrowed.fieldExpr.obj.name,
-                                    fieldName: narrowed.fieldExpr.field,
-                                    narrowedTy: narrowed.innerTy,
-                                }],
-                        };
-                    }
-                    else {
-                        // Complex expression (call result, deep chain, etc.): transform can't detect these,
-                        // so keep old behavior — substitute + narrowedVar + narrowedExpr
-                        narrowedVar = narrowed.varName;
-                        narrowedExprResolved = narrowed.narrowedExpr ?? resolveExpr(narrowed.fieldExpr, ctx);
-                        rawThen = substituteRawExpr(e.then, narrowed.fieldExpr, { kind: "var", name: narrowed.varName });
-                        thenCtx = withEnv(thenCtx, extend(thenCtx.env, narrowed.varName, narrowed.innerTy));
-                    }
-                }
-                else if (narrowed && !narrowed.inThen && !narrowed.fieldExpr) {
-                    // v === undefined: narrow v in the else branch
-                    elseCtx = withEnv(elseCtx, extend(elseCtx.env, narrowed.varName, narrowed.innerTy));
-                }
-                // Compound || with === undefined: narrow all checked vars in else branch
-                // e.g. if (a === undefined || b === undefined) then X else Y → narrow a,b in Y
-                // TODO: resolve-time narrowing works but transform doesn't emit match unwrap
-                // for || conditions yet — decompose || into nested matches in transform.
-                // Workaround: split || into separate if guards in user code.
-                if (!narrowed && 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));
-                    }
+            // 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));
+            }
+            // 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);
+            // 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
-            if (narrowedExprResolved && (then_.ty.kind === "void" || else_.ty.kind === "void") && ty.kind !== "optional") {
-                ty = { kind: "optional", inner: ty };
+            else if (then_.ty.kind === "optional" && else_.ty.kind !== "optional" && else_.ty.kind !== "unknown") {
+                // Asymmetric optional: one branch returns Option<T>, the other returns T.
+                // Widen to Option<T> so callers/return-coercion see the wider type.
+                ty = then_.ty;
+            }
+            else if (else_.ty.kind === "optional" && then_.ty.kind !== "optional" && then_.ty.kind !== "unknown") {
+                ty = else_.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);
@@ -652,8 +941,10 @@ 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.
@@ -665,6 +956,20 @@ function resolveBlock(stmts, ctx) {
             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;
+            }
         }
     }
     return result;
@@ -672,20 +977,48 @@ function resolveBlock(stmts, ctx) {
 function resolveStmt(s, ctx) {
     switch (s.kind) {
         case "let": {
-            const declTy = resolveTsType(s.tsType, ctx.overrides, s.name);
+            // No source annotation → infer type from initializer (resolved first).
+            if (s.tsType === null) {
+                const init = resolveExpr(s.init, ctx);
+                const ty = init.ty;
+                const mutable = s.mutable || isRefMutableInTS(ty);
+                return [{ kind: "let", name: s.name, ty, mutable, init }, extend(ctx.env, s.name, ty)];
+            }
+            // expandAlias unwraps an array/collection alias (`type Board = number[]`)
+            // to its underlying type, so array methods / index-assignment on the
+            // local dispatch correctly (params get the same treatment, see makeParams).
+            const declTy = expandAlias(resolveTsType(s.tsType, ctx.overrides, s.name), ctx.typeDecls);
             // 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;
+            let ty;
+            if (isUnmodeledTy(declTy, ctx.typeDecls) && !isUnmodeledTy(init.ty, ctx.typeDecls)) {
+                // ts-morph's declared type is opaque to us (an expanded union it made
+                // by inlining an alias we shadow via declare-type, or any-laden), but
+                // LS resolved the initializer to something concrete. Take the structure
+                // from `init.ty`, keeping only the optionality ts-morph reported.
+                ty = declTy.kind === "optional" && init.ty.kind !== "optional"
+                    ? { kind: "optional", inner: init.ty }
+                    : init.ty;
+            }
+            else {
+                // Map indexing: TS says T, but access can fail → use Optional<T> from init
+                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)
             const mutable = s.mutable || isRefMutableInTS(ty);
             return [{ kind: "let", name: s.name, ty, mutable, init }, extend(ctx.env, s.name, ty)];
         }
         case "assign": {
             const targetTy = lookup(ctx.env, s.target) ?? { kind: "unknown" };
-            return [{ kind: "assign", target: s.target, value: coerceStr(resolveExpr(s.value, ctx), targetTy) }, ctx.env];
+            let value = coerceStr(resolveExpr(s.value, ctx), targetTy);
+            // Auto-wrap non-optional value in Some when target is optional
+            const isUndef = value.kind === "var" && value.name === "undefined";
+            if (targetTy.kind === "optional" && value.ty.kind !== "optional" && value.ty.kind !== "unknown" && !isUndef) {
+                value = wrapSome(value, targetTy);
+            }
+            return [{ kind: "assign", target: s.target, value }, ctx.env];
         }
         case "return": {
             let value = coerceStr(resolveExpr(s.value, ctx), ctx.returnTy);
@@ -705,30 +1038,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": {
@@ -791,7 +1130,7 @@ function resolveStmt(s, ctx) {
         case "assert": {
             const specCtx = { ...ctx, inSpec: true };
             const expr = resolveExpr(parseExpr(s.expr), specCtx);
-            return [{ kind: "assert", expr }, ctx.env];
+            return [{ kind: "assert", expr, assumed: s.assumed }, ctx.env];
         }
     }
 }
@@ -966,40 +1305,52 @@ function containsReturn(stmts) {
     return false;
 }
 // ── Resolve function / module ────────────────────────────────
-function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), opts) {
+function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), externs = new Map(), moduleConstants = 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");
+    const params = fn.params.map(p => ({ name: p.name, ty: expandAlias(resolveTsType(p.tsType, overrides, p.name), typeDecls) }));
+    const returnTy = expandAlias(resolveTsType(fn.returnType, overrides, "\\result"), typeDecls);
     let env = null;
+    // Module-level constants are in scope for every function body. Added before
+    // params so a param named the same as a const would shadow it (param wins).
+    for (const [name, ty] of moduleConstants)
+        env = extend(env, name, ty);
     if (opts?.thisBinding)
         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, fnReturns, inSpec: false, inLambda: false, narrowedFields: [] };
+    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, fnReturns, externs, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
     const requiresCtx = { ...baseCtx, inSpec: true };
-    const ensuresCtx = { ...baseCtx, allowResult: true, inSpec: true };
+    const ensuresCtx = { ...baseCtx, env: extend(env, "\\result", returnTy), allowResult: true, inSpec: true };
     // Apply type parameter constraints from //@ type T (==) annotations
     const typeParams = fn.typeParams.map(tp => {
         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),
         forcePure: fn.pure,
-        body: resolveBlock(fn.body, baseCtx),
+        body: resolveBlock(fn.body, bodyCtx),
     };
 }
-function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map()) {
+function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), externs = new Map(), moduleConstants = 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, fnReturns, {
+    const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, fnReturns, externs, moduleConstants, {
         thisBinding: { name: "this", ty: thisType },
         forcePure: false, // class methods are never pure (they access this)
     }));
@@ -1008,6 +1359,22 @@ function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns =
 /** Pre-compute Ty on all TypeDeclInfo fields/variants/aliases.
  *  Called once per module so consumers can read field.type instead of re-parsing tsType. */
 function precomputeFieldTypes(typeDecls) {
+    precomputeFieldTypesInner(typeDecls);
+    // Expand alias references inside record/variant field types so downstream
+    // code doesn't have to follow `user("Ruleset")` indirection at every lookup.
+    for (const d of typeDecls) {
+        if (d.fields)
+            for (const f of d.fields)
+                if (f.type)
+                    f.type = expandAlias(f.type, typeDecls);
+        if (d.variants)
+            for (const v of d.variants)
+                for (const f of v.fields)
+                    if (f.type)
+                        f.type = expandAlias(f.type, typeDecls);
+    }
+}
+function precomputeFieldTypesInner(typeDecls) {
     for (const d of typeDecls) {
         if (d.fields)
             for (const f of d.fields)
@@ -1021,7 +1388,6 @@ function precomputeFieldTypes(typeDecls) {
     }
 }
 export function resolveModule(raw) {
-    _synVarCounter = 0;
     precomputeFieldTypes(raw.typeDecls);
     const pureFns = computePureFns(raw.functions);
     // Pre-compute function parameter and return types
@@ -1029,20 +1395,76 @@ export function resolveModule(raw) {
     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"));
+        fnParams.set(fn.name, fn.params.map(p => expandAlias(resolveTsType(p.tsType, overrides, p.name), raw.typeDecls)));
+        fnReturns.set(fn.name, expandAlias(resolveTsType(fn.returnType, overrides, "\\result"), raw.typeDecls));
     }
-    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedFields: [] };
-    const constants = (raw.constants ?? []).map(c => ({
-        name: c.name,
-        ty: parseTsType(c.tsType),
-        value: resolveExpr(c.value, emptyCtx),
-    }));
+    // Externs: resolve param/return types once. For bare-name externs (no dot),
+    // also register in fnReturns so ordinary `foo(args)` calls get the right
+    // return type at resolution; dotted externs are handled in resolveExpr's
+    // call case via the externs map directly.
+    const externs = new Map();
+    // First pass: register signatures so spec resolution (below) can reference
+    // them — including the extern referring to itself, or specs that mention
+    // sibling externs.
+    for (const ext of raw.externs ?? []) {
+        const params = ext.params.map(p => parseTsType(p.tsType));
+        const returnTy = parseTsType(ext.returnType);
+        externs.set(ext.qualified, { flat: ext.flat, params, returnTy });
+        if (!ext.qualified.includes("."))
+            fnReturns.set(ext.qualified, returnTy);
+    }
+    // Second pass: resolve the lifted `requires`/`ensures` strings in each
+    // extern's own param scope. `\result` is in scope under `ensures`.
+    const tExterns = (raw.externs ?? []).map(ext => {
+        const sig = externs.get(ext.qualified);
+        let env = null;
+        for (let i = 0; i < ext.params.length; i++) {
+            env = extend(env, ext.params[i].name, sig.params[i]);
+        }
+        const baseCtx = { env, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: sig.returnTy, pureFns, fnParams, fnReturns, externs, inSpec: true, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
+        const ensuresCtx = { ...baseCtx, env: extend(env, "\\result", sig.returnTy), allowResult: true };
+        const requires = ext.requires.map(s => {
+            try {
+                return resolveSpec(s, baseCtx);
+            }
+            catch {
+                return null;
+            }
+        }).filter((e) => e !== null);
+        const ensures = ext.ensures.map(s => {
+            try {
+                return resolveSpec(s, ensuresCtx);
+            }
+            catch {
+                return null;
+            }
+        }).filter((e) => e !== null);
+        return {
+            qualified: ext.qualified,
+            flat: ext.flat,
+            typeParams: ext.typeParams,
+            params: ext.params.map((p, i) => ({ name: p.name, ty: sig.params[i] })),
+            returnTy: sig.returnTy,
+            requires,
+            ensures,
+        };
+    });
+    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, fnReturns, externs, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
+    const constants = (raw.constants ?? []).map(c => {
+        const ty = expandAlias(parseTsType(c.tsType), raw.typeDecls);
+        // Propagate the declared type into the value's resolution context so that
+        // record literals on map-typed constants (e.g. `Record<string, number>`)
+        // get their `ty` set to `map<...>` rather than `user("...")`.
+        const valueCtx = { ...emptyCtx, returnTy: ty };
+        return { name: c.name, ty, value: resolveExpr(c.value, valueCtx) };
+    });
+    const moduleConstants = new Map(constants.map(c => [c.name, c.ty]));
     return {
         file: raw.file,
         typeDecls: raw.typeDecls,
+        externs: tExterns,
         constants,
-        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams, fnReturns)),
-        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams, fnReturns)),
+        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams, fnReturns, externs, moduleConstants)),
+        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams, fnReturns, externs, moduleConstants)),
     };
 }