lemmascript 0.2.0 → 0.3.0

package/tools/dist/extract.js

@@ -8,6 +8,21 @@ import { Project, Node, SyntaxKind, ScriptTarget } from "ts-morph";
 // ── Expression extraction ────────────────────────────────────
 /** When set, calls whose function/method name matches this key are replaced with havoc. */
 let _havocKey = null;
+/** Generic bounds erasure map — set during extractFunction, applied in extractStmts. */
+let _typeParamMap = new Map();
+function _eraseGenerics(tsType) {
+    if (_typeParamMap.size === 0)
+        return tsType;
+    if (tsType.includes(" | ")) {
+        const arms = [...new Set(tsType.split(" | ").map(a => _eraseGenerics(a.trim())))];
+        return arms.length === 1 ? arms[0] : arms.join(" | ");
+    }
+    if (tsType.endsWith("[]"))
+        return _eraseGenerics(tsType.slice(0, -2)) + "[]";
+    if (_typeParamMap.has(tsType))
+        return _typeParamMap.get(tsType);
+    return tsType;
+}
 function extractExpr(node) {
     // Havoc key matching: replace matching calls with havoc expression
     if (_havocKey && Node.isCallExpression(node)) {
@@ -63,9 +78,24 @@ function extractExpr(node) {
     if (node.getKind() === SyntaxKind.ThisKeyword) {
         return { kind: "var", name: "this" };
     }
-    // Property access: x.foo
+    // Property access: x.foo or x?.foo
     if (Node.isPropertyAccessExpression(node)) {
-        return { kind: "field", obj: extractExpr(node.getExpression()), field: node.getName() };
+        const obj = extractExpr(node.getExpression());
+        const field = node.getName();
+        // Optional chaining on data access: x?.foo → x !== undefined ? x.foo : undefined
+        // Skip if this is a method call (parent is CallExpression) — handled differently.
+        if (node.hasQuestionDotToken()) {
+            const parent = node.getParent();
+            const isMethodCall = parent && Node.isCallExpression(parent) && parent.getExpression() === node;
+            if (!isMethodCall) {
+                return { kind: "conditional",
+                    cond: { kind: "binop", op: "!==", left: obj, right: { kind: "var", name: "undefined" } },
+                    then: { kind: "field", obj, field },
+                    else: { kind: "var", name: "undefined" },
+                };
+            }
+        }
+        return { kind: "field", obj, field };
     }
     // Element access: arr[i]
     if (Node.isElementAccessExpression(node)) {
@@ -129,15 +159,38 @@ function extractExpr(node) {
         }
         throw new Error(`Unsupported arrow function body: ${node.getText().slice(0, 80)}`);
     }
-    // Array literal: [a, b, c] → arrayLiteral, [...arr, elem] → push(arr, elem)
+    // Array literal: [a, b, c] → arrayLiteral, with spreads → concatenation
     if (Node.isArrayLiteralExpression(node)) {
         const elems = node.getElements();
-        // [...arr, elem] → push(arr, elem)
-        if (elems.length === 2 && Node.isSpreadElement(elems[0])) {
-            return { kind: "call", fn: { kind: "field", obj: extractExpr(elems[0].getExpression()), field: "push" }, args: [extractExpr(elems[1])] };
+        const hasSpread = elems.some(e => Node.isSpreadElement(e));
+        if (!hasSpread) {
+            return { kind: "arrayLiteral", elems: elems.map(e => extractExpr(e)) };
+        }
+        // Build concatenation: [a, ...b, c] → [a] + b + [c]
+        // Group consecutive non-spread elements into array literals, spreads are bare
+        const segments = [];
+        let currentLiterals = [];
+        for (const e of elems) {
+            if (Node.isSpreadElement(e)) {
+                if (currentLiterals.length > 0) {
+                    segments.push({ kind: "arrayLiteral", elems: currentLiterals });
+                    currentLiterals = [];
+                }
+                segments.push(extractExpr(e.getExpression()));
+            }
+            else {
+                currentLiterals.push(extractExpr(e));
+            }
         }
-        // [a, b, c] or [] → arrayLiteral
-        return { kind: "arrayLiteral", elems: elems.map(e => extractExpr(e)) };
+        if (currentLiterals.length > 0) {
+            segments.push({ kind: "arrayLiteral", elems: currentLiterals });
+        }
+        // Fold segments with arrayConcat
+        let result = segments[0];
+        for (let i = 1; i < segments.length; i++) {
+            result = { kind: "binop", op: "arrayConcat", left: result, right: segments[i] };
+        }
+        return result;
     }
     // Object literal: { res: true, done: false } or { ...obj, res: true }
     if (Node.isObjectLiteralExpression(node)) {
@@ -172,14 +225,31 @@ function extractExpr(node) {
         const name = node.getExpression().getText();
         if (name === "Map" || name === "Set") {
             const typeArgs = node.getTypeArguments();
+            // Use explicit type args if present, otherwise infer from TS type system
             const tsType = typeArgs && typeArgs.length > 0
                 ? `${name}<${typeArgs.map(t => t.getText()).join(", ")}>`
-                : name;
+                : _eraseGenerics(typeToString(node.getType()));
             const args = node.getArguments();
-            // new Map(arr.map(fn)) — map-from-array constructor
+            // new Map(source) — clone existing map or build from entries
             if (name === "Map" && args && args.length === 1) {
+                const argType = args[0].getType();
+                const argSymbol = argType.getSymbol()?.getName() ?? argType.getAliasSymbol()?.getName();
+                if (argSymbol === "Map") {
+                    // new Map(existingMap) — identity (Dafny maps are value types)
+                    return extractExpr(args[0]);
+                }
+                // new Map(entries) — map-from-array constructor
                 return { kind: "call", fn: { kind: "var", name: "__mapFromArray" }, args: [extractExpr(args[0])] };
             }
+            // new Set([a, b, c]) — set with initial elements
+            if (name === "Set" && args && args.length === 1) {
+                const arg = args[0];
+                if (Node.isArrayLiteralExpression(arg)) {
+                    return { kind: "emptyCollection", collectionType: "Set", tsType, initElems: arg.getElements().map(e => extractExpr(e)) };
+                }
+                // new Set(existingSet) — pass through
+                return extractExpr(arg);
+            }
             return { kind: "emptyCollection", collectionType: name, tsType };
         }
     }
@@ -187,6 +257,10 @@ function extractExpr(node) {
     if (Node.isAsExpression(node)) {
         return extractExpr(node.getExpression());
     }
+    // null → undefined (both map to None in backends)
+    if (Node.isNullLiteral(node)) {
+        return { kind: "var", name: "undefined" };
+    }
     throw new Error(`Unsupported expression: ${node.getText()}`);
 }
 // ── Annotation parsing ───────────────────────────────────────
@@ -216,13 +290,15 @@ function collectAnnotations(node, body) {
     return own;
 }
 // ── Type declaration extraction ──────────────────────────────
-function extractTypeDecl(decl) {
+function extractTypeDecl(decl, extraDecls) {
     const name = decl.getName();
     const type = decl.getType();
+    const typeParams = decl.getTypeParameters().map(tp => tp.getName());
+    const tpField = typeParams.length > 0 ? typeParams : undefined;
     if (type.isUnion()) {
         const members = type.getUnionTypes();
         if (members.every(m => m.isStringLiteral())) {
-            return { name, kind: "string-union", values: members.map(m => m.getLiteralValue()) };
+            return { name, typeParams: tpField, kind: "string-union", values: members.map(m => m.getLiteralValue()) };
         }
         if (members.every(m => m.isObject())) {
             const discriminant = findDiscriminant(members);
@@ -230,7 +306,8 @@ function extractTypeDecl(decl) {
                 const variants = members.map(m => {
                     const tagProp = m.getProperty(discriminant);
                     const tagType = tagProp?.getTypeAtLocation(decl);
-                    const tag = tagType?.getLiteralValue();
+                    const tag = tagType?.isStringLiteral() ? String(tagType.getLiteralValue())
+                        : tagType?.getText() ?? "unknown";
                     const fields = [];
                     for (const prop of m.getProperties()) {
                         if (prop.getName() === discriminant)
@@ -239,15 +316,19 @@ function extractTypeDecl(decl) {
                     }
                     return { name: tag, fields };
                 });
-                return { name, kind: "discriminated-union", discriminant, variants };
+                return { name, typeParams: tpField, kind: "discriminated-union", discriminant, variants };
             }
         }
     }
-    if (type.isObject())
-        return extractRecord(name, type, decl);
+    if (type.isObject() || type.isIntersection())
+        return extractRecord(name, type, decl, undefined, extraDecls);
+    // Primitive type alias: type TaskId = number → alias
+    const tsType = typeToString(type);
+    if (tsType !== name)
+        return { name, kind: "alias", aliasOf: tsType };
     return null;
 }
-function extractInterface(decl) {
+function extractInterface(decl, extraDecls) {
     // Collect field type overrides from trailing //@ type annotations
     const overrides = new Map();
     for (const member of decl.getMembers()) {
@@ -258,16 +339,44 @@ function extractInterface(decl) {
                 overrides.set(member.getName(), match[1]);
         }
     }
-    return extractRecord(decl.getName(), decl.getType(), decl, overrides);
+    return extractRecord(decl.getName(), decl.getType(), decl, overrides, extraDecls);
 }
-function extractRecord(name, type, locationNode, overrides) {
+function extractRecord(name, type, locationNode, overrides, extraDecls) {
     const props = type.getProperties();
     if (props.length === 0)
         return null;
     const fields = [];
     for (const prop of props) {
         const override = overrides?.get(prop.getName());
-        const tsType = override ?? typeToString(prop.getTypeAtLocation(locationNode));
+        if (override) {
+            fields.push({ name: prop.getName(), tsType: override });
+            continue;
+        }
+        const propType = prop.getTypeAtLocation(locationNode);
+        let tsType = typeToString(propType);
+        // Inline anonymous object types: ts-morph names them __type.
+        // Generate a synthetic named record and reference it by name instead.
+        if (extraDecls && tsType.includes("__type")) {
+            let innerType = propType;
+            let isOptional = false;
+            if (propType.isUnion()) {
+                const uTypes = propType.getUnionTypes();
+                const nonUndef = uTypes.filter(t => !t.isUndefined());
+                if (uTypes.some(t => t.isUndefined()) && nonUndef.length === 1) {
+                    innerType = nonUndef[0];
+                    isOptional = true;
+                }
+            }
+            if (innerType.isObject() && !innerType.isArray()) {
+                const fName = prop.getName();
+                const synName = name + fName.charAt(0).toUpperCase() + fName.slice(1);
+                const extracted = extractRecord(synName, innerType, locationNode, undefined, extraDecls);
+                if (extracted) {
+                    extraDecls.push(extracted);
+                    tsType = isOptional ? `${synName} | undefined` : synName;
+                }
+            }
+        }
         fields.push({ name: prop.getName(), tsType });
     }
     return { name, kind: "record", fields };
@@ -283,7 +392,7 @@ function findDiscriminant(members) {
             if (!p)
                 return false;
             const t = p.getDeclarations()[0] ? p.getTypeAtLocation(p.getDeclarations()[0]) : null;
-            return t?.isStringLiteral() ?? false;
+            return (t?.isStringLiteral() || t?.isBooleanLiteral()) ?? false;
         });
         if (allHave)
             return name;
@@ -310,7 +419,11 @@ function typeToString(type) {
         return name;
     }
     if (type.isUnion()) {
-        return type.getUnionTypes().map(typeToString).join(" | ");
+        const parts = [...new Set(type.getUnionTypes().map(typeToString))];
+        return parts.join(" | ");
+    }
+    if (type.isTuple()) {
+        return `[${type.getTupleElements().map(t => typeToString(t)).join(", ")}]`;
     }
     if (type.isArray()) {
         const elem = type.getArrayElementTypeOrThrow();
@@ -366,18 +479,44 @@ function extractStmts(stmts) {
     for (const s of stmts) {
         const line = s.getStartLineNumber();
         // Ghost annotations from leading comments → inject before this statement
-        result.push(...parseSpecComments(s.getLeadingCommentRanges(), line));
+        const leadingComments = s.getLeadingCommentRanges();
+        result.push(...parseSpecComments(leadingComments, line));
+        // //@ skip — omit this statement from the verification model
+        if (leadingComments.some(r => r.getText().trim() === "//@ skip")) {
+            continue;
+        }
         if (Node.isVariableStatement(s)) {
             const havocMatch = s.getLeadingCommentRanges()
-                .map(r => r.getText().trim().match(/^\/\/@ havoc(?:\s+(\S+))?$/))
+                .map(r => r.getText().trim().match(/^\/\/@ havoc(?:\s*:\s*(.+)|(?:\s+(\S+)))?$/))
                 .find(m => m !== null);
-            const havocKey = havocMatch?.[1] ?? null;
+            const havocType = havocMatch?.[1]?.trim() ?? null; // //@ havoc : Type
+            const havocKey = havocMatch?.[2] ?? null; // //@ havoc key
             const isHavoc = !!havocMatch;
             for (const d of s.getDeclarations()) {
+                // Havoc on destructuring: emit each named binding as a separate havoced variable
+                const nameNode = d.getNameNode();
+                if (isHavoc && !havocKey && Node.isObjectBindingPattern(nameNode)) {
+                    const rhsType = d.getType();
+                    const havocTypes = havocType?.split(",").map(t => t.trim()) ?? [];
+                    const elements = nameNode.getElements();
+                    for (let ei = 0; ei < elements.length; ei++) {
+                        const el = elements[ei];
+                        const name = el.getName();
+                        const propType = rhsType.getProperty(name)?.getTypeAtLocation(d);
+                        const tsType = havocTypes[ei] ?? (propType ? _eraseGenerics(typeToString(propType)) : "unknown");
+                        result.push({
+                            kind: "let", name, mutable: s.getDeclarationKind() === "let",
+                            tsType, init: { kind: "havoc", tsType }, line,
+                        });
+                    }
+                    continue;
+                }
                 const declType = d.getType();
                 let init;
                 if (isHavoc && !havocKey) {
-                    init = { kind: "havoc", tsType: typeToString(declType) };
+                    // Use explicit type from //@ havoc : Type, or initializer type (pre-cast), or declared type
+                    const initType = d.getInitializer()?.getType();
+                    init = { kind: "havoc", tsType: havocType ?? typeToString(initType ?? declType) };
                 }
                 else {
                     const initializer = d.getInitializer();
@@ -389,7 +528,7 @@ function extractStmts(stmts) {
                     kind: "let",
                     name: d.getName(),
                     mutable: s.getDeclarationKind() === "let",
-                    tsType: typeToString(declType),
+                    tsType: havocType ?? _eraseGenerics(typeToString(declType)),
                     init,
                     line,
                 });
@@ -535,6 +674,11 @@ function extractStmts(stmts) {
             result.push({ kind: "throw", line });
             continue;
         }
+        // Block statement: { ... } — flatten into parent
+        if (Node.isBlock(s)) {
+            result.push(...extractStmts(s.getStatements()));
+            continue;
+        }
         throw new Error(`Unsupported statement at line ${line}: ${s.getText().slice(0, 80)}`);
     }
     // Ghost comments after the last statement (before closing brace) appear as sibling trivia nodes
@@ -568,12 +712,35 @@ function extractStmts(stmts) {
     return result;
 }
 // ── Function extraction ──────────────────────────────────────
-function extractFunction(fn) {
+function extractFunction(fn, parentAnnotations) {
+    // Generic bounds erasure: <T extends Base> → substitute T with Base everywhere
+    // Unbounded type params are preserved as Dafny type parameters
+    _typeParamMap = new Map();
+    const unboundedTypeParams = [];
+    for (const tp of fn.getTypeParameters?.() ?? []) {
+        const constraint = tp.getConstraint();
+        if (constraint)
+            _typeParamMap.set(tp.getName(), constraint.getText());
+        else
+            unboundedTypeParams.push(tp.getName());
+    }
     const body = fn.getBody();
-    if (!body || !Node.isBlock(body))
-        throw new Error(`${fn.getName?.() ?? "arrow"}: function body is not a block`);
-    const bodyStmts = body.getStatements();
-    const annots = collectAnnotations(fn, bodyStmts);
+    // Expression-body arrow: wrap in implicit return
+    let extractedBody;
+    let annots;
+    if (body && !Node.isBlock(body)) {
+        const expr = extractExpr(body);
+        extractedBody = [{ kind: "return", value: expr, line: body.getStartLineNumber() }];
+        annots = parentAnnotations ?? collectAnnotations(fn);
+    }
+    else if (body && Node.isBlock(body)) {
+        const bodyStmts = body.getStatements();
+        extractedBody = extractStmts(bodyStmts);
+        annots = collectAnnotations(fn, bodyStmts);
+    }
+    else {
+        throw new Error(`${fn.getName?.() ?? "arrow"}: function has no body`);
+    }
     const typeAnnotations = [];
     for (const a of annots) {
         if (a.kind === "type") {
@@ -584,27 +751,94 @@ function extractFunction(fn) {
     }
     return {
         name: fn.getName?.() ?? "<anonymous>",
-        params: fn.getParameters().map(p => ({ name: p.getName(), tsType: p.getTypeNode()?.getText() ?? "unknown" })),
-        returnType: fn.getReturnTypeNode()?.getText() ?? "unknown",
+        typeParams: unboundedTypeParams,
+        params: fn.getParameters().flatMap(p => {
+            // Flatten destructured object params into individual params
+            const nameNode = p.getNameNode();
+            if (Node.isObjectBindingPattern(nameNode)) {
+                const type = p.getType();
+                return nameNode.getElements().map(el => {
+                    const name = el.getName();
+                    const propType = type.getProperty(name)?.getTypeAtLocation(p);
+                    return { name, tsType: propType ? typeToString(propType) : "unknown" };
+                });
+            }
+            let tsType = _eraseGenerics(p.getTypeNode()?.getText() ?? "unknown");
+            // Optional parameters (foo?: T) need | undefined in the type string
+            if (p.hasQuestionToken())
+                tsType = `${tsType} | undefined`;
+            return [{ name: p.getName(), tsType }];
+        }),
+        returnType: (() => {
+            const node = fn.getReturnTypeNode();
+            if (node)
+                return _eraseGenerics(node.getText());
+            const inferred = fn.getReturnType();
+            if (inferred.isAny())
+                return "unknown";
+            return _eraseGenerics(typeToString(inferred));
+        })(),
         requires: annots.filter(a => a.kind === "requires").map(a => a.expr),
         ensures: annots.filter(a => a.kind === "ensures").map(a => a.expr),
         typeAnnotations,
-        body: extractStmts(bodyStmts),
+        body: extractedBody,
         line: fn.getStartLineNumber(),
     };
 }
 // ── Module extraction ────────────────────────────────────────
 export function extractModule(sourceFile) {
     const typeDecls = [];
+    // Parse //@ declare-type directives from file comments
+    for (const range of sourceFile.getLeadingCommentRanges()) {
+        const text = range.getText().trim();
+        if (!text.startsWith("//@ declare-type "))
+            continue;
+        const body = text.slice("//@ declare-type ".length);
+        const match = body.match(/^(\w+)\s*\{(.+)\}$/);
+        if (!match)
+            continue;
+        const name = match[1];
+        const fieldsStr = match[2];
+        const fields = fieldsStr.split(",").map(f => f.trim()).filter(Boolean).map(f => {
+            const [fname, ftype] = f.split(":").map(s => s.trim());
+            return { name: fname, tsType: ftype };
+        });
+        typeDecls.push({ name, kind: "record", fields });
+    }
+    // Also scan statement-level comments for declare-type
+    for (const stmt of sourceFile.getStatements()) {
+        for (const range of stmt.getLeadingCommentRanges()) {
+            const text = range.getText().trim();
+            if (!text.startsWith("//@ declare-type "))
+                continue;
+            const body = text.slice("//@ declare-type ".length);
+            const match = body.match(/^(\w+)\s*\{(.+)\}$/);
+            if (!match)
+                continue;
+            const name = match[1];
+            const fields = match[2].split(",").map(f => f.trim()).filter(Boolean).map(f => {
+                const [fname, ftype] = f.split(":").map(s => s.trim());
+                return { name: fname, tsType: ftype };
+            });
+            typeDecls.push({ name, kind: "record", fields });
+        }
+    }
     // Extract type declarations in source order to respect dependencies
+    // Skip types already declared via //@ declare-type
+    const declaredNames = new Set(typeDecls.map(d => d.name));
     for (const stmt of sourceFile.getStatements()) {
-        if (Node.isTypeAliasDeclaration(stmt)) {
-            const info = extractTypeDecl(stmt);
+        if (Node.isTypeAliasDeclaration(stmt) && !declaredNames.has(stmt.getName())) {
+            const extra = [];
+            const info = extractTypeDecl(stmt, extra);
+            typeDecls.push(...extra);
             if (info)
                 typeDecls.push(info);
         }
-        else if (Node.isInterfaceDeclaration(stmt)) {
-            const info = extractInterface(stmt);
+        else if (Node.isInterfaceDeclaration(stmt) && !declaredNames.has(stmt.getName())) {
+            const extra = [];
+            const info = extractInterface(stmt, extra);
+            // Synthetic types from inline objects must precede the parent type
+            typeDecls.push(...extra);
             if (info)
                 typeDecls.push(info);
         }
@@ -640,35 +874,134 @@ export function extractModule(sourceFile) {
     for (const fn of sourceFile.getFunctions()) {
         allFns.push({ name: fn.getName() ?? "<anonymous>", node: fn });
     }
-    // const f = (...) => { ... } — treat as named function
+    // const f = (...) => expr  OR  const f = (...) => { ... }
     for (const stmt of sourceFile.getStatements()) {
         if (Node.isVariableStatement(stmt)) {
             for (const decl of stmt.getDeclarationList().getDeclarations()) {
                 const init = decl.getInitializer();
                 if (init && Node.isArrowFunction(init)) {
-                    allFns.push({ name: decl.getName(), node: init });
+                    allFns.push({ name: decl.getName(), node: init, parentStmt: stmt });
                 }
             }
         }
     }
     // If any function has //@ verify, only extract those (brownfield mode).
-    // Otherwise extract all functions (backwards-compatible with existing examples).
-    const hasVerifyDirective = allFns.some(f => f.node.getFullText().includes('//@ verify'));
-    const fnsToExtract = hasVerifyDirective
-        ? allFns.filter(f => f.node.getFullText().includes('//@ verify'))
-        : allFns;
+    // For expression-body arrows, //@ verify may be on the parent variable statement.
+    function hasVerify(f) {
+        if (f.node.getFullText().includes('//@ verify'))
+            return true;
+        if (f.parentStmt) {
+            for (const r of f.parentStmt.getLeadingCommentRanges()) {
+                if (r.getText().includes('//@ verify'))
+                    return true;
+            }
+        }
+        return false;
+    }
+    const hasVerifyDirective = allFns.some(hasVerify);
+    const fnsToExtract = hasVerifyDirective ? allFns.filter(hasVerify) : allFns;
     const functions = fnsToExtract.map(f => {
-        const raw = extractFunction(f.node);
+        // 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>"
         return raw;
     });
+    // Resolve imported type names referenced in function signatures and type fields
+    const knownTypeNames = new Set(typeDecls.map(d => d.name));
+    const primitives = new Set(["number", "string", "boolean", "void", "unknown", "undefined"]);
+    const builtinTypes = new Set(["Map", "Set", "Array", "Record", "Promise", "Date", "RegExp", "Error"]);
+    function resolveTypeName(name) {
+        if (knownTypeNames.has(name) || primitives.has(name) || builtinTypes.has(name))
+            return;
+        for (const sf2 of sourceFile.getProject().getSourceFiles()) {
+            for (const stmt of sf2.getStatements()) {
+                if (Node.isTypeAliasDeclaration(stmt) && stmt.getName() === name) {
+                    const extra = [];
+                    const info = extractTypeDecl(stmt, extra);
+                    typeDecls.push(...extra);
+                    if (info) {
+                        typeDecls.push(info);
+                        knownTypeNames.add(name);
+                    }
+                }
+                if (Node.isInterfaceDeclaration(stmt) && stmt.getName() === name && !knownTypeNames.has(name)) {
+                    const extra = [];
+                    const info = extractInterface(stmt, extra);
+                    typeDecls.push(...extra);
+                    if (info) {
+                        typeDecls.push(info);
+                        knownTypeNames.add(name);
+                    }
+                }
+            }
+            if (knownTypeNames.has(name))
+                break;
+        }
+        // Recursively resolve types referenced by the newly added type's fields
+        const decl = typeDecls.find(d => d.name === name);
+        if (decl?.fields) {
+            for (const f of decl.fields) {
+                for (const m of f.tsType.matchAll(/\b([A-Z]\w*)\b/g))
+                    resolveTypeName(m[1]);
+            }
+        }
+        if (decl?.variants) {
+            for (const v of decl.variants) {
+                for (const f of v.fields) {
+                    for (const m of f.tsType.matchAll(/\b([A-Z]\w*)\b/g))
+                        resolveTypeName(m[1]);
+                }
+            }
+        }
+    }
+    for (const fn of functions) {
+        const refs = [fn.returnType, ...fn.params.map(p => p.tsType)];
+        for (const ref of refs) {
+            for (const m of ref.matchAll(/\b([A-Z]\w*)\b/g))
+                resolveTypeName(m[1]);
+        }
+    }
+    // Resolve union param types: A | B → intersection of fields
+    const typeDeclMap = new Map(typeDecls.map(d => [d.name, d]));
+    for (const fn of functions) {
+        for (const p of fn.params) {
+            if (!p.tsType.includes(" | "))
+                continue;
+            const arms = p.tsType.split(" | ").map(a => a.trim());
+            const armDecls = arms.map(a => typeDeclMap.get(a)).filter((d) => !!d && d.kind === "record");
+            if (armDecls.length < 2 || armDecls.length !== arms.length)
+                continue;
+            // Compute field name intersection
+            const fieldSets = armDecls.map(d => new Set(d.fields.map(f => f.name)));
+            const common = [...fieldSets[0]].filter(name => fieldSets.every(s => s.has(name)));
+            // Find an existing type that matches, or use the first arm's fields
+            const match = armDecls.find(d => d.fields.length === common.length && d.fields.every(f => common.includes(f.name)));
+            if (match) {
+                p.tsType = match.name;
+            }
+            else {
+                // Generate synthetic union type with intersected fields
+                const synName = arms.join("Or");
+                if (!typeDeclMap.has(synName)) {
+                    const fields = common.map(name => {
+                        const f = armDecls[0].fields.find(f => f.name === name);
+                        return { name: f.name, tsType: f.tsType };
+                    });
+                    typeDecls.push({ name: synName, kind: "record", fields });
+                    typeDeclMap.set(synName, typeDecls[typeDecls.length - 1]);
+                }
+                p.tsType = synName;
+            }
+        }
+    }
     // In brownfield mode, filter consts to only those referenced by verified functions.
-    // Types are NOT filtered — they may be needed transitively (e.g. Option from T | undefined).
     if (hasVerifyDirective) {
         const referencedNames = new Set();
         function collectNames(stmts) {
             for (const s of stmts) {
                 if (s.kind === "let") {
+                    referencedNames.add(s.tsType);
                     collectNamesExpr(s.init);
                 }
                 if (s.kind === "assign") {
@@ -743,6 +1076,26 @@ export function extractModule(sourceFile) {
             }
         }
         constants.splice(0, constants.length, ...constants.filter(c => referencedNames.has(c.name)));
+        // Filter types to only those referenced by verified functions (transitive)
+        const neededTypes = new Set();
+        function markType(name) {
+            if (neededTypes.has(name))
+                return;
+            const d = typeDecls.find(t => t.name === name);
+            if (!d)
+                return;
+            neededTypes.add(name);
+            for (const f of d.fields ?? [])
+                for (const m of f.tsType.matchAll(/\b([A-Z]\w*)\b/g))
+                    markType(m[1]);
+            for (const v of d.variants ?? [])
+                for (const f of v.fields)
+                    for (const m of f.tsType.matchAll(/\b([A-Z]\w*)\b/g))
+                        markType(m[1]);
+        }
+        for (const name of referencedNames)
+            markType(name);
+        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
     const knownTypes = new Set(typeDecls.map(d => d.name));
@@ -753,15 +1106,53 @@ export function extractModule(sourceFile) {
             resolveType(t.getArrayElementTypeOrThrow(), locationNode);
             return;
         }
+        // Resolve type aliases (e.g. string unions imported from other files)
+        const alias = t.getAliasSymbol();
+        if (alias) {
+            const aliasName = alias.getName();
+            if (!knownTypes.has(aliasName) && !builtins.has(aliasName) && !aliasName.startsWith("__")) {
+                const decls = alias.getDeclarations();
+                if (decls.length > 0 && Node.isTypeAliasDeclaration(decls[0])) {
+                    const extra = [];
+                    const info = extractTypeDecl(decls[0], extra);
+                    if (info) {
+                        typeDecls.push(...extra);
+                        typeDecls.push(info);
+                        knownTypes.add(aliasName);
+                    }
+                }
+                else if (t.getProperties().length > 0) {
+                    // Alias declaration not available (e.g. intersection type) — extract from properties
+                    const extra = [];
+                    const info = extractRecord(aliasName, t, locationNode, undefined, extra);
+                    if (info) {
+                        typeDecls.push(...extra);
+                        typeDecls.push(info);
+                        knownTypes.add(aliasName);
+                    }
+                }
+            }
+        }
+        if (t.isUnion()) {
+            for (const u of t.getUnionTypes())
+                resolveType(u, locationNode);
+            return;
+        }
         for (const arg of t.getTypeArguments())
             resolveType(arg, locationNode);
         const sym = t.getSymbol() ?? t.getAliasSymbol();
         const name = sym?.getName();
-        if (name && !knownTypes.has(name) && !builtins.has(name) && t.isObject()) {
-            const info = extractRecord(name, t, locationNode);
+        if (name && !name.startsWith("__") && !knownTypes.has(name) && !builtins.has(name) && (t.isObject() || t.isIntersection())) {
+            const extra = [];
+            const info = extractRecord(name, t, locationNode, undefined, extra);
             if (info) {
+                typeDecls.push(...extra);
                 typeDecls.push(info);
                 knownTypes.add(name);
+                // Recursively resolve types referenced in this type's fields
+                for (const prop of t.getProperties()) {
+                    resolveType(prop.getTypeAtLocation(locationNode), locationNode);
+                }
             }
         }
     }
@@ -769,6 +1160,45 @@ export function extractModule(sourceFile) {
         for (const p of f.node.getParameters())
             resolveType(p.getType(), p);
     }
+    // Resolve anonymous object return types into synthetic named types
+    for (let i = 0; i < fnsToExtract.length; i++) {
+        const f = fnsToExtract[i];
+        const fn = functions[i];
+        const retType = f.node.getReturnType();
+        // Prefer alias symbol (named type aliases) over underlying object symbol (__type)
+        const aliasSym = retType.getAliasSymbol();
+        if (aliasSym && !aliasSym.getName().startsWith("__")) {
+            // Named type alias — resolve it instead of generating a synthetic name
+            resolveType(retType, f.node);
+            const aliasName = aliasSym.getName();
+            if (knownTypes.has(aliasName)) {
+                // Preserve type arguments: Result<Model, Err> not just Result
+                const typeArgs = retType.getAliasTypeArguments();
+                fn.returnType = typeArgs.length > 0
+                    ? `${aliasName}<${typeArgs.map(t => typeToString(t)).join(", ")}>`
+                    : aliasName;
+            }
+            continue;
+        }
+        const sym = retType.getSymbol();
+        if (sym?.getName() === "__type" && retType.isObject() && !retType.isArray()) {
+            const synName = fn.name.charAt(0).toUpperCase() + fn.name.slice(1) + "Result";
+            if (!knownTypes.has(synName)) {
+                const extra = [];
+                const info = extractRecord(synName, retType, f.node, undefined, extra);
+                if (info) {
+                    typeDecls.push(...extra);
+                    typeDecls.push(info);
+                    knownTypes.add(synName);
+                }
+            }
+            fn.returnType = synName;
+            // Also resolve imported types referenced in the return type's fields
+            for (const prop of retType.getProperties()) {
+                resolveType(prop.getTypeAtLocation(f.node), f.node);
+            }
+        }
+    }
     // Extract classes with //@ verify methods
     const classes = [];
     for (const cls of sourceFile.getClasses()) {