lemmascript 0.5.1 → 0.5.3

package/tools/dist/dafny-emit.js

@@ -45,11 +45,16 @@ const DAFNY_KEYWORDS = new Set([
     "codatatype", "nameonly", "twostate", "opaque", "replaceable", "colemma",
     "copredicate", "inductive",
 ]);
+// The Dafny out-parameter name for the method currently being emitted. Default
+// `res`, but bumped (e.g. `res_`) when a parameter is named `res` — set by
+// methodHeader and reset per decl. `\result` in an ensures must use the *same*
+// name, so escapeName routes it here.
+let _resultName = "res";
 function escapeName(name) {
     // \result is carried through the IR as the var name "\\result"; render it
-    // as Dafny's canonical return-value identifier.
+    // as the current method's out-parameter name.
     if (name === "\\result")
-        return "res";
+        return _resultName;
     if (DAFNY_KEYWORDS.has(name))
         return `${name}_`;
     // Dafny doesn't allow identifiers starting with _
@@ -66,7 +71,17 @@ function paramList(params) {
  *  `returns (res: ())` on a void method fails verification. */
 function methodHeader(prefix, params, returnType) {
     const sig = `${prefix}(${paramList(params)})`;
-    return returnType.kind === "void" ? sig : `${sig} returns (res: ${tyToDafny(returnType)})`;
+    if (returnType.kind === "void")
+        return sig;
+    // The out-parameter is `res` by default, but a parameter named `res` (e.g. an
+    // Express handler's `(req, res)`) would collide; pick a fresh name and record
+    // it so `\result` references in the ensures/body resolve to the same name.
+    const taken = new Set(params.map(p => escapeName(p.name)));
+    let resName = "res";
+    while (taken.has(resName))
+        resName += "_";
+    _resultName = resName;
+    return `${sig} returns (${resName}: ${tyToDafny(returnType)})`;
 }
 // ── Lean op → Dafny op ─────────────────────────────────────
 const OP_MAP = {
@@ -166,6 +181,12 @@ function emitExpr(e) {
                     return `Std.Collections.Seq.Map(${args[0]}, ${obj})`;
                 if (e.method === "filter")
                     return `Std.Collections.Seq.Filter(${args[0]}, ${obj})`;
+                // filterMap (synthesized in resolve): drop Nones and unwrap to seq<T>.
+                if (e.method === "filterSome") {
+                    needPreamble("SeqFilterSome");
+                    needPreamble("OptionType");
+                    return `SeqFilterSome(${obj})`;
+                }
                 if (e.method === "every")
                     return `Std.Collections.Seq.All(${obj}, ${args[0]})`;
                 if (e.method === "findLast") {
@@ -350,16 +371,15 @@ function emitExpr(e) {
                 needPreamble("BitAnd");
                 return `BitAnd(${emitExpr(e.left)}, ${emitExpr(e.right)})`;
             }
-            // int * real coercion: wrap int side with "as real"
-            if (["+", "-", "*", "/"].includes(op)) {
-                const leftIsReal = e.left.kind === "num" && !Number.isInteger(e.left.value);
-                const rightIsReal = e.right.kind === "num" && !Number.isInteger(e.right.value);
-                if (leftIsReal !== rightIsReal) {
-                    const left = leftIsReal ? emitExpr(e.left) : `(${emitExpr(e.left)} as real)`;
-                    const right = rightIsReal ? emitExpr(e.right) : `(${emitExpr(e.right)} as real)`;
-                    return `(${left} ${op} ${right})`;
-                }
+            // x | y → BitOr(x, y) (recursive, mirrors BitAnd). Dafny has no `|` on int,
+            // only on bitvectors.
+            if (e.op === "|") {
+                needPreamble("BitOr");
+                return `BitOr(${emitExpr(e.left)}, ${emitExpr(e.right)})`;
             }
+            // int→real coercion is now injected upstream in transform (toReal nodes),
+            // which has full type information — including real-typed variables, not
+            // just literals — so no literal-based coercion is needed here.
             return `(${wrapQuantifier(e.left)} ${op} ${emitExpr(e.right)})`;
         }
         case "implies": {
@@ -399,6 +419,8 @@ function emitExpr(e) {
                 needPreamble("MathMin");
                 needPreamble("MinOfSeq");
             }
+            if (e.fn === "Perm")
+                needPreamble("Perm");
             return `${escapeName(e.fn)}(${args.join(", ")})`;
         }
         case "field": {
@@ -414,6 +436,8 @@ function emitExpr(e) {
         case "toNat":
             // Dafny doesn't need toNat — just emit the inner expression
             return emitExpr(e.expr);
+        case "toReal":
+            return `(${emitExpr(e.expr)} as real)`;
         case "index": {
             const obj = emitExpr(e.arr);
             const idx = emitExpr(e.idx);
@@ -522,6 +546,10 @@ function emitStmt(s, indent) {
                 return `${pad}var ${escapeName(s.name)}: ${tyToDafny(s.type)} := ${emitExpr(s.value)};`;
             return `${pad}var ${escapeName(s.name)} := ${emitExpr(s.value)};`;
         case "assign":
+            // Transform.ts lowers a bare expression statement to an assign with target _
+            // so special case this to Dafny's anonymous binding.
+            if (s.target === "_")
+                return `${pad}var _ := ${emitExpr(s.value)};`;
             return `${pad}${escapeName(s.target)} := ${emitExpr(s.value)};`;
         case "ghostLet":
             return `${pad}ghost var ${escapeName(s.name)}: ${tyToDafny(s.type)} := ${emitExpr(s.value)};`;
@@ -593,13 +621,33 @@ function emitStmt(s, indent) {
 }
 // ── Declaration emission ────────────────────────────────────
 function emitDecl(d) {
+    _resultName = "res"; // default; methodHeader bumps it if a param is named `res`
     switch (d.kind) {
         case "inductive": {
             const tp = d.typeParams?.length ? `<${d.typeParams.join(", ")}>` : "";
+            // Dafny requires a destructor shared across constructors to have a single
+            // type. Two TS variants can legitimately share a field name with different
+            // types (e.g. label.targetId: string vs leaf.targetId: string?). Detect
+            // such collisions and make those destructors per-constructor unique. Safe:
+            // variant-field reads lower to positional match bindings, never named
+            // destructors on the union (a name shared with differing types isn't even
+            // accessible on the union type in TS), so nothing references the old name.
+            const typesByField = new Map();
+            for (const c of d.constructors)
+                for (const f of c.fields) {
+                    let s = typesByField.get(f.name);
+                    if (!s) {
+                        s = new Set();
+                        typesByField.set(f.name, s);
+                    }
+                    s.add(tyToDafny(f.type));
+                }
+            const collides = new Set([...typesByField].filter(([, s]) => s.size > 1).map(([n]) => n));
             const ctors = d.constructors.map(c => {
                 if (c.fields.length === 0)
                     return escapeName(c.name);
-                return `${escapeName(c.name)}(${paramList(c.fields)})`;
+                const fields = c.fields.map(f => collides.has(f.name) ? { ...f, name: `${f.name}_${c.name}` } : f);
+                return `${escapeName(c.name)}(${paramList(fields)})`;
             });
             return `datatype ${d.name}${tp} = ${ctors.join(" | ")}`;
         }
@@ -609,6 +657,11 @@ function emitDecl(d) {
         case "type-alias": {
             return `type ${d.name} = ${tyToDafny(d.target)}`;
         }
+        case "opaque-type": {
+            // Abstract type — no definition. `(==)` so it can sit inside datatypes
+            // that derive structural equality. Never constructed or destructured.
+            return `type ${d.name}(==)`;
+        }
         case "def": {
             const tp = d.typeParams.length > 0 ? `<${d.typeParams.join(", ")}>` : "";
             const lines = [`function ${d.name}${tp}(${paramList(d.params)}): ${tyToDafny(d.returnType)}`];
@@ -724,6 +777,14 @@ const BIT_AND = `function BitAnd(x: int, y: int): int
   if x == 0 || y == 0 then 0
   else 2 * BitAnd(x / 2, y / 2) + (if x % 2 == 1 && y % 2 == 1 then 1 else 0)
 }`;
+const BIT_OR = `function BitOr(x: int, y: int): int
+  requires x >= 0 && y >= 0
+  decreases x
+{
+  if x == 0 then y
+  else if y == 0 then x
+  else 2 * BitOr(x / 2, y / 2) + (if x % 2 == 1 || y % 2 == 1 then 1 else 0)
+}`;
 const JS_FLOOR_DIV = `function JSFloorDiv(a: int, b: int): int
   requires b != 0
 {
@@ -743,6 +804,12 @@ const CEIL_REAL = `function CeilReal(x: real): int
   if x == (x.Floor as real) then x.Floor
   else x.Floor + 1
 }`;
+const SEQ_FILTER_SOME = `function SeqFilterSome<T>(xs: seq<Option<T>>): seq<T>
+  ensures |SeqFilterSome(xs)| <= |xs|
+{
+  if |xs| == 0 then []
+  else (if xs[0].Some? then [xs[0].value] else []) + SeqFilterSome(xs[1..])
+}`;
 const SEQ_FIND_INDEX = `function SeqFindIndex<T>(s: seq<T>, p: T -> bool): int
   ensures -1 <= SeqFindIndex(s, p) < |s|
   ensures SeqFindIndex(s, p) >= 0 ==> p(s[SeqFindIndex(s, p)])
@@ -956,6 +1023,10 @@ const NAT_TO_STRING = `function NatToString(n: nat): string
   else NatToString(n / 10) + [digit]
 }`;
 const MATH_ABS = `function MathAbs(x: int): nat { if x >= 0 then x else -x }`;
+// perm(a, b) — `a` and `b` are reorderings of each other (equal as multisets).
+// Transparent (Dafny unfolds it), so hand-proofs can reason with `multiset`
+// directly. The `(==)` bound requires the element type to support equality.
+const PERM = `predicate Perm<T(==)>(a: seq<T>, b: seq<T>) { multiset(a) == multiset(b) }`;
 const SET_TO_SEQ = `method SetToSeq<T>(s: set<T>) returns (res: seq<T>)
   ensures forall x :: x in s <==> x in res
   ensures |res| == |s|
@@ -981,11 +1052,13 @@ const PREAMBLE_CODE = [
     ["SetToSeq", SET_TO_SEQ],
     ["Pow2", POW2],
     ["BitAnd", BIT_AND],
+    ["BitOr", BIT_OR],
     ["JSFloorDiv", JS_FLOOR_DIV],
     ["CeilReal", CEIL_REAL],
     ["FloorReal", FLOOR_REAL],
     ["SeqIndexOf", SEQ_INDEX_OF],
     ["SeqFindIndex", SEQ_FIND_INDEX],
+    ["SeqFilterSome", SEQ_FILTER_SOME],
     ["SeqFindLast", SEQ_FIND_LAST],
     ["SeqFlatten", SEQ_FLATTEN],
     ["SeqJoin", SEQ_JOIN],
@@ -1001,6 +1074,7 @@ const PREAMBLE_CODE = [
     ["MathMax", MATH_MAX],
     ["MaxOfSeq", MAX_OF_SEQ],
     ["MinOfSeq", MIN_OF_SEQ],
+    ["Perm", PERM],
 ];
 // ── Constructor and record helpers ───────────────────────────
 let _recordCtors = new Map();

package/tools/dist/extract.js

@@ -69,6 +69,10 @@ function detectCrossFileExtern(callee, sourceFile) {
     if (decls.length === 0)
         return null;
     const currentPath = sourceFile.getFilePath();
+    // A declaration in the current file is authoritative — don't resolve to a
+    // same-named definition in another file.
+    if (decls.some(d => d.getSourceFile().getFilePath() === currentPath))
+        return null;
     const externalDecl = decls.find(d => d.getSourceFile().getFilePath() !== currentPath);
     if (!externalDecl)
         return null;
@@ -159,6 +163,26 @@ function _synthName(elemName, otherName) {
     const sanitize = (s) => s.replace(/[^A-Za-z0-9]+/g, "_").replace(/^_+|_+$/g, "");
     return `ArrayOf_${sanitize(elemName)}_Or_${sanitize(otherName)}`;
 }
+/**
+ * Fall-through for a union LS can't model as a tagged union (no runtime type
+ * test maps to a tag — e.g. members are unreachable imports with no visible
+ * discriminant). Registers an opaque-type TypeDeclInfo and returns its name, so
+ * the union becomes one abstract `type` rather than invalid raw-union Dafny.
+ *
+ * Sound because an opaque type has no constructor and no tag predicate: any
+ * attempt to build or type-test the value fails to lower, so it can only be
+ * passed through — the one sound use of a union we can't discriminate. Distinct
+ * from dropping the field (which collapses values and is unsound): the value is
+ * preserved, just uninspectable.
+ */
+function _synthOpaque(memberNames) {
+    const sanitize = (s) => s.replace(/[^A-Za-z0-9]+/g, "_").replace(/^_+|_+$/g, "");
+    const name = `Opaque_${memberNames.map(sanitize).join("_or_")}`;
+    if (_synthArrayUnions !== null && !_synthArrayUnions.some(d => d.name === name)) {
+        _synthArrayUnions.push({ name, kind: "opaque" });
+    }
+    return name;
+}
 /**
  * String-level fallback for synth detection on a `T[] | U` shape, used by
  * declare-type field parsing (where no ts-morph TypeNode is available). The
@@ -277,10 +301,10 @@ function extractExpr(node) {
     if (Node.isNumericLiteral(node)) {
         return { kind: "num", value: Number(node.getLiteralValue()) };
     }
-    // BigInt literal (e.g. 32n, 0xffffn) — treat as integer
+    // BigInt literal (e.g. 32n, 0xffffn) — integer, with bigint division semantics
     if (Node.isBigIntLiteral(node)) {
         const text = node.getText().replace(/n$/, '');
-        return { kind: "num", value: Number(text) };
+        return { kind: "num", value: Number(text), big: true };
     }
     // Template literal: `foo${x}bar` → "foo" + x + "bar"
     if (Node.isTemplateExpression(node)) {
@@ -431,12 +455,16 @@ function extractExpr(node) {
             const typeNode = p.getTypeNode();
             return { name: p.getName(), tsType: typeNode ? typeNode.getText() : undefined };
         });
+        // Capture an explicit return annotation (`(x): Out => …`) so resolve can type
+        // return-position record literals to their named type instead of a tuple.
+        const retNode = node.getReturnTypeNode();
+        const returnTsType = retNode ? typeToString(node.getReturnType()) : undefined;
         const body = node.getBody();
         if (Node.isExpression(body)) {
-            return { kind: "lambda", params, body: extractExpr(body) };
+            return { kind: "lambda", params, body: extractExpr(body), returnTsType };
         }
         if (Node.isBlock(body)) {
-            return { kind: "lambda", params, body: extractStmts(body.getStatements()) };
+            return { kind: "lambda", params, body: extractStmts(body.getStatements()), returnTsType };
         }
         throw new Error(`Unsupported arrow function body: ${node.getText().slice(0, 80)}`);
     }
@@ -552,6 +580,12 @@ function extractExpr(node) {
     if (Node.isNullLiteral(node)) {
         return { kind: "var", name: "undefined" };
     }
+    // `typeof X` — only meaningful in a `typeof X === "string"` discriminator over
+    // a synth `U | T[]` union (see narrow's parseTypeofStringCheck); any other use
+    // survives to emit and errors there.
+    if (Node.isTypeOfExpression(node)) {
+        return { kind: "unop", op: "typeof", expr: extractExpr(node.getExpression()) };
+    }
     throw new Error(`Unsupported expression: ${node.getText()}`);
 }
 // ── Annotation parsing ───────────────────────────────────────
@@ -796,7 +830,10 @@ function typeToString(type) {
         return "bigint";
     if (type.isString() || type.isStringLiteral())
         return "string";
-    if (type.isBoolean())
+    // TS expands `boolean` to the literal union `false | true`; normalize the
+    // literals back so the union dedupes to a single `boolean` rather than being
+    // mistaken for an unmodelable multi-member union.
+    if (type.isBoolean() || type.isBooleanLiteral())
         return "boolean";
     // Named type alias (e.g. Priority = "low" | "medium" | "high") — use the alias name
     if (type.getAliasSymbol()) {
@@ -838,7 +875,15 @@ function typeToString(type) {
             }
         }
         const parts = [...new Set(unionTypes.map(typeToString))];
-        return parts.join(" | ");
+        // `undefined`/`null` are optional markers; a single real member with them
+        // is an Option, left as `X | undefined` for the optional lowering.
+        const real = parts.filter(p => p !== "undefined" && p !== "null");
+        if (real.length <= 1)
+            return parts.join(" | ");
+        // A genuine multi-member union with no tagged-union shape LS can model →
+        // a single opaque type. (See _synthOpaque.) Preserve an outer optional.
+        const opaque = _synthOpaque(real);
+        return real.length === parts.length ? opaque : `${opaque} | undefined`;
     }
     if (type.isTuple()) {
         return `[${type.getTupleElements().map(t => typeToString(t)).join(", ")}]`;
@@ -1704,6 +1749,22 @@ function extractFunctionInner(fn, parentAnnotations) {
             return [{ name: p.getName(), tsType }];
         }),
         returnType: (() => {
+            // `async` with no `await`: the `Promise<T>` wrapper is just the calling
+            // convention (the body returns T-typed values), so unwrap to T. Gated on
+            // no `await` — that's the suspension point we can't model, and it only
+            // type-checks inside `async`, so the gate is self-justifying. With `await`
+            // present we leave `Promise<...>` (unmodellable) rather than atomize it.
+            const isAsync = fn.isAsync?.() ?? false;
+            const noAwait = fn.getDescendantsOfKind(SyntaxKind.AwaitExpression).length === 0;
+            if (isAsync && noAwait) {
+                const args = fn.getReturnType().getTypeArguments();
+                if (args.length === 1) {
+                    if (args[0].isAny())
+                        return "unknown";
+                    return _eraseGenerics(typeToString(args[0]));
+                }
+                return "void"; // Promise<void>
+            }
             const node = fn.getReturnTypeNode();
             if (node && Node.isUnionTypeNode(node))
                 return _eraseGenerics(_tsTypeFromUnionNode(node));
@@ -1718,6 +1779,7 @@ function extractFunctionInner(fn, parentAnnotations) {
         ensures: annots.filter(a => a.kind === "ensures").map(a => a.expr),
         decreases: annots.find(a => a.kind === "decreases")?.expr ?? null,
         pure: hasPureAnnotation(fn, body && Node.isBlock(body) ? body.getStatements() : undefined),
+        autohavoc: false, // set in extractModule (file-level directive or per-function)
         typeAnnotations,
         body: extractedBody,
         line: fn.getStartLineNumber(),
@@ -1869,63 +1931,144 @@ export function extractModule(sourceFile) {
             }
         }
     }
+    // Top-level inline closures: a handler passed directly to a module-level call,
+    // e.g. `app.get("/x", (req, res) => { //@ verify ... })`. "Move" each such
+    // closure to the top level by extracting it as a synthetic named function.
+    // Only top-level call arguments are considered (not nested lambdas), and only
+    // closures carrying a //@ verify (so ordinary callbacks aren't pulled in). The
+    // name is derived from the call's method and route literal (e.g. get_x).
+    const usedNames = new Set(allFns.map(f => f.name));
+    const sanitizeIdent = (s) => s.replace(/[^A-Za-z0-9]+/g, "_").replace(/^_+|_+$/g, "") || "handler";
+    for (const stmt of sourceFile.getStatements()) {
+        if (!Node.isExpressionStatement(stmt))
+            continue;
+        const call = stmt.getExpression();
+        if (!Node.isCallExpression(call))
+            continue;
+        const callee = call.getExpression();
+        const method = Node.isPropertyAccessExpression(callee) ? callee.getName()
+            : Node.isIdentifier(callee) ? callee.getText() : "handler";
+        const routeArg = call.getArguments().find(a => Node.isStringLiteral(a));
+        const route = routeArg && Node.isStringLiteral(routeArg) ? routeArg.getLiteralValue() : "";
+        for (const arg of call.getArguments()) {
+            if (!Node.isArrowFunction(arg))
+                continue;
+            if (!hasLineDirective(arg.getFullText(), "verify"))
+                continue;
+            const base = sanitizeIdent(route ? `${method}_${route}` : method);
+            let name = base, n = 2;
+            while (usedNames.has(name))
+                name = `${base}_${n++}`;
+            usedNames.add(name);
+            allFns.push({ name, node: arg, parentStmt: stmt });
+        }
+    }
     // `//@ extern` on a same-file declaration: register the function as an
     // opaque axiom (signature + any //@ requires/ensures), skip its body. Use
     // when the function is outside LS's verification model — e.g., wraps a
     // regex — but its callers should still be verifiable against an
     // uninterpreted predicate. Parallel to auto-extern for cross-file calls,
     // and emitted the same way (`function {:axiom} foo(...)` in Dafny).
+    // Match a `//@ <kw>` directive only as the first non-whitespace on a line, so
+    // a mention mid-line in prose or inside a block/JSDoc comment (e.g. "the
+    // `//@ extern` annotation", or ` * //@ extern`) doesn't falsely trigger it.
+    function hasLineDirective(text, kw) {
+        return new RegExp(String.raw `^[ \t]*//@ ${kw}\b`, "m").test(text);
+    }
     function hasExtern(f) {
-        if (f.node.getFullText().includes('//@ extern'))
+        if (hasLineDirective(f.node.getFullText(), "extern"))
             return true;
         if (f.parentStmt) {
             for (const r of f.parentStmt.getLeadingCommentRanges()) {
-                if (r.getText().includes('//@ extern'))
+                if (hasLineDirective(r.getText(), "extern"))
                     return true;
             }
         }
         return false;
     }
+    // `//@ extern NS.method` registers the extern under a *dotted* qualified name,
+    // so a real `NS.method(args)` call dispatches to it (resolve.ts) with no
+    // wrapper — e.g. `//@ extern fs.readFileSync` lets you call `fs.readFileSync`
+    // directly while still discharging its `//@ requires`. The function declaration
+    // just carries the signature/contract; its own name is unused.
+    function externName(f) {
+        const re = /^[ \t]*\/\/@ extern[ \t]+(\S+)/m;
+        const m = f.node.getFullText().match(re);
+        if (m)
+            return m[1];
+        if (f.parentStmt) {
+            for (const r of f.parentStmt.getLeadingCommentRanges()) {
+                const m2 = r.getText().match(re);
+                if (m2)
+                    return m2[1];
+            }
+        }
+        return null;
+    }
     for (const f of allFns) {
         if (!hasExtern(f))
             continue;
-        if (_externs.has(f.name))
+        const qualified = externName(f) ?? f.name;
+        const flat = qualified.replace(/\./g, "_");
+        if (_externs.has(qualified))
             continue;
         const sig = f.node.getType().getCallSignatures()[0];
         if (!sig)
             continue;
         const typeParams = sig.getTypeParameters().map(tp => tp.getText());
+        // Normalize via typeToString (not raw getText): resolves declare-type
+        // shadows and yields bare names, so a param typed by an unreachable import
+        // becomes `AgentMessage`, not `import("/abs/path").AgentMessage`.
         const params = sig.getParameters().map(p => ({
             name: p.getName(),
-            tsType: p.getTypeAtLocation(f.node).getText(),
+            tsType: typeToString(p.getTypeAtLocation(f.node)),
         }));
-        const returnType = sig.getReturnType().getText();
+        const returnType = typeToString(sig.getReturnType());
         const annots = collectFunctionAnnotations(f.node);
         const requires = annots.filter(a => a.kind === "requires").map(a => a.expr);
         const ensures = annots.filter(a => a.kind === "ensures").map(a => a.expr);
-        _externs.set(f.name, { qualified: f.name, flat: f.name, typeParams, params, returnType, requires, ensures });
+        _externs.set(qualified, { qualified, flat, typeParams, params, returnType, requires, ensures });
     }
     // If any function has //@ verify, only extract those (brownfield mode).
     // For expression-body arrows, //@ verify may be on the parent variable statement.
     function hasVerify(f) {
-        if (f.node.getFullText().includes('//@ verify'))
+        if (hasLineDirective(f.node.getFullText(), "verify"))
             return true;
         if (f.parentStmt) {
             for (const r of f.parentStmt.getLeadingCommentRanges()) {
-                if (r.getText().includes('//@ verify'))
+                if (hasLineDirective(r.getText(), "verify"))
                     return true;
             }
         }
         return false;
     }
-    const hasVerifyDirective = sourceFile.getFullText().includes('//@ verify');
+    const hasVerifyDirective = hasLineDirective(sourceFile.getFullText(), "verify");
     const nonExternFns = allFns.filter(f => !hasExtern(f));
     const fnsToExtract = hasVerifyDirective ? nonExternFns.filter(hasVerify) : nonExternFns;
+    // `//@ autohavoc` — enable the auto-havoc abstraction (see autohavoc.ts).
+    // File-level: a directive at column 0 (top of file) enables it for every
+    // function. Per-function: the annotation attached to a function (or its
+    // parent variable statement), mirroring `//@ verify`.
+    const fileAutohavoc = /^\/\/@ autohavoc\b/m.test(sourceFile.getFullText());
+    function hasAutohavoc(f) {
+        if (fileAutohavoc)
+            return true;
+        if (hasLineDirective(f.node.getFullText(), "autohavoc"))
+            return true;
+        if (f.parentStmt) {
+            for (const r of f.parentStmt.getLeadingCommentRanges()) {
+                if (hasLineDirective(r.getText(), "autohavoc"))
+                    return true;
+            }
+        }
+        return false;
+    }
     const functions = fnsToExtract.map(f => {
         // For expression-body arrows, annotations come from the parent variable statement
         const parentAnnots = f.parentStmt ? parseAnnotations(f.parentStmt) : undefined;
         const raw = extractFunction(f.node, parentAnnots);
         raw.name = f.name; // use the const name, not "<anonymous>"
+        raw.autohavoc = hasAutohavoc(f);
         return raw;
     });
     // Resolve type references in function signatures via ts-morph's type
@@ -2121,10 +2264,17 @@ export function extractModule(sourceFile) {
                 collectNamesExpr(e.else);
             }
         }
+        // Signature types (params + return) get base-name stripping below; body /
+        // spec references stay exact-match (so a body `let xs: Hunk[]` doesn't pull
+        // `Hunk` into the filter early and reorder output that resolveType re-adds).
+        const sigTypes = new Set();
         for (const fn of functions) {
-            for (const p of fn.params)
+            for (const p of fn.params) {
                 referencedNames.add(p.tsType);
+                sigTypes.add(p.tsType);
+            }
             referencedNames.add(fn.returnType);
+            sigTypes.add(fn.returnType);
             collectNames(fn.body);
             // Also scan spec annotations for identifier references
             for (const spec of [...fn.requires, ...fn.ensures]) {
@@ -2153,6 +2303,15 @@ export function extractModule(sourceFile) {
         }
         for (const name of referencedNames)
             markType(name);
+        // Signature types also mark their base after stripping array/optional
+        // WRAPPERS (`Out[]`/`Msg | undefined` → `Out`/`Msg`), so a function returning
+        // a local `Out[]` keeps `Out`. Wrappers only — never dig into generic args
+        // (`Omit<FilePathOptions, …>` must not pull in the inner type).
+        for (const name of sigTypes) {
+            const base = name.replace(/\s*\|\s*(undefined|null)\s*$/, "").replace(/(\[\])+$/, "").trim();
+            if (base !== name && /^[A-Za-z_]\w*$/.test(base))
+                markType(base);
+        }
         typeDecls.splice(0, typeDecls.length, ...typeDecls.filter(d => neededTypes.has(d.name) || declaredNames.has(d.name)));
     }
     // Resolve imported types: extract types referenced in function signatures but not in this file