lemmascript 0.4.0 → 0.5.0

package/tools/dist/resolve.js

@@ -14,6 +14,15 @@ 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: [] };
@@ -55,6 +64,43 @@ function wrapSome(value, optionalTy) {
         args: [value], ty: optionalTy, callKind: "pure",
     };
 }
+/** 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`).
@@ -94,6 +140,12 @@ function classifyOptExpr(e, ctx) {
             return null;
         return { varName: e.name, innerTy: ty.inner };
     }
+    if (e.kind === "result") {
+        const ty = lookup(ctx.env, "\\result");
+        if (!ty || ty.kind !== "optional")
+            return null;
+        return { varName: "\\result", innerTy: ty.inner };
+    }
     const resolved = resolveExpr(e, ctx);
     if (resolved.ty.kind !== "optional")
         return null;
@@ -193,11 +245,74 @@ 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). */
@@ -264,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.
@@ -276,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. */
@@ -322,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")
@@ -344,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" };
@@ -357,9 +529,13 @@ 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" };
@@ -452,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" &&
@@ -463,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)) {
@@ -553,22 +749,29 @@ function resolveExpr(e, ctx) {
                 }
                 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));
+                    // 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) {
-                        // 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 };
+                        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;
                 }
@@ -580,8 +783,19 @@ function resolveExpr(e, ctx) {
         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;
@@ -600,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" };
@@ -683,6 +900,14 @@ function resolveExpr(e, ctx) {
             else if (else_.ty.kind === "void" && then_.ty.kind !== "void" && then_.ty.kind !== "unknown") {
                 ty = { kind: "optional", inner: then_.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 };
         }
         case "emptyCollection": {
@@ -752,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);
@@ -877,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];
         }
     }
 }
@@ -1052,18 +1305,22 @@ 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, narrowedPaths: [], narrowedIndices: [] };
+    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);
@@ -1087,13 +1344,13 @@ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns
         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)
     }));
@@ -1102,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)
@@ -1122,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, narrowedPaths: [], narrowedIndices: [] };
-    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)),
     };
 }

package/tools/dist/specparser.js

@@ -231,6 +231,12 @@ class Parser {
                 this.advance();
                 return { kind: "bool", value: false };
             }
+            // Match the body extractor (extract.ts NullLiteral): `null` and
+            // `undefined` are interchangeable in LS, both map to None.
+            if (t.value === "null") {
+                this.advance();
+                return { kind: "var", name: "undefined" };
+            }
             // new Set<T>() / new Map<K,V>()
             if (t.value === "new") {
                 this.advance();