lemmascript 0.1.0 → 0.3.0

package/tools/dist/extract.js

@@ -4,13 +4,63 @@
  * Produces structured AST nodes, not strings.
  * The only strings are //@ annotation expressions (parsed later by specparser).
  */
-import { Project, Node, SyntaxKind } from "ts-morph";
+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)) {
+        const fnExpr = node.getExpression();
+        const name = Node.isPropertyAccessExpression(fnExpr) ? fnExpr.getName()
+            : Node.isIdentifier(fnExpr) ? fnExpr.getText()
+                : null;
+        if (name === _havocKey) {
+            return { kind: "havoc", tsType: typeToString(node.getType()) };
+        }
+    }
     // Numeric literal
     if (Node.isNumericLiteral(node)) {
         return { kind: "num", value: Number(node.getLiteralValue()) };
     }
+    // BigInt literal (e.g. 32n, 0xffffn) — treat as integer
+    if (Node.isBigIntLiteral(node)) {
+        const text = node.getText().replace(/n$/, '');
+        return { kind: "num", value: Number(text) };
+    }
+    // Template literal: `foo${x}bar` → "foo" + x + "bar"
+    if (Node.isTemplateExpression(node)) {
+        const parts = [];
+        const head = node.getHead().getLiteralText();
+        if (head)
+            parts.push({ kind: "str", value: head });
+        for (const span of node.getTemplateSpans()) {
+            parts.push(extractExpr(span.getExpression()));
+            const text = span.getLiteral().getLiteralText();
+            if (text)
+                parts.push({ kind: "str", value: text });
+        }
+        return parts.reduce((left, right) => ({ kind: "binop", op: "+", left, right }));
+    }
+    // No-substitution template literal: `hello` → "hello"
+    if (Node.isNoSubstitutionTemplateLiteral(node)) {
+        return { kind: "str", value: node.getLiteralText() };
+    }
     // String literal
     if (Node.isStringLiteral(node)) {
         return { kind: "str", value: node.getLiteralValue() };
@@ -24,9 +74,28 @@ function extractExpr(node) {
     if (Node.isIdentifier(node)) {
         return { kind: "var", name: node.getText() };
     }
-    // Property access: x.foo
+    // this
+    if (node.getKind() === SyntaxKind.ThisKeyword) {
+        return { kind: "var", name: "this" };
+    }
+    // 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)) {
@@ -90,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)) };
         }
-        // [a, b, c] or [] → arrayLiteral
-        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));
+            }
+        }
+        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)) {
@@ -108,6 +200,10 @@ function extractExpr(node) {
             if (Node.isSpreadAssignment(prop)) {
                 spread = extractExpr(prop.getExpression());
             }
+            else if (Node.isShorthandPropertyAssignment(prop)) {
+                const name = prop.getName();
+                fields.push({ name, value: { kind: "var", name } });
+            }
             else if (Node.isPropertyAssignment(prop)) {
                 const init = prop.getInitializer();
                 if (init)
@@ -120,6 +216,51 @@ function extractExpr(node) {
     if (Node.isConditionalExpression(node)) {
         return { kind: "conditional", cond: extractExpr(node.getCondition()), then: extractExpr(node.getWhenTrue()), else: extractExpr(node.getWhenFalse()) };
     }
+    // Non-null assertion: expr!
+    if (Node.isNonNullExpression(node)) {
+        return { kind: "nonNull", expr: extractExpr(node.getExpression()) };
+    }
+    // new Map<K,V>() / new Set<T>() — with or without initializer
+    if (Node.isNewExpression(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(", ")}>`
+                : _eraseGenerics(typeToString(node.getType()));
+            const args = node.getArguments();
+            // 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 };
+        }
+    }
+    // As-expression: expr as T — strip the type assertion
+    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 ───────────────────────────────────────
@@ -149,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);
@@ -163,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)
@@ -172,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()) {
@@ -191,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 };
@@ -216,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;
@@ -224,39 +400,136 @@ function findDiscriminant(members) {
     return null;
 }
 function typeToString(type) {
-    if (type.isNumber())
+    if (type.isUndefined())
+        return "undefined";
+    if (type.isNumber() || type.isNumberLiteral())
         return "number";
-    if (type.isString())
+    if (type.isBigInt() || type.isBigIntLiteral())
+        return "bigint";
+    if (type.isString() || type.isStringLiteral())
         return "string";
     if (type.isBoolean())
         return "boolean";
+    // Named type alias (e.g. Priority = "low" | "medium" | "high") — use the alias name
+    if (type.getAliasSymbol()) {
+        const name = type.getAliasSymbol().getName();
+        const args = type.getAliasTypeArguments();
+        if (args.length > 0)
+            return `${name}<${args.map(t => typeToString(t)).join(", ")}>`;
+        return name;
+    }
+    if (type.isUnion()) {
+        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();
         return `${typeToString(elem)}[]`;
     }
     const symbol = type.getSymbol() ?? type.getAliasSymbol();
-    if (symbol)
-        return symbol.getName();
+    if (symbol) {
+        const name = symbol.getName();
+        const typeArgs = type.getTypeArguments();
+        if (typeArgs.length > 0) {
+            return `${name}<${typeArgs.map(t => typeToString(t)).join(", ")}>`;
+        }
+        return name;
+    }
     return type.getText();
 }
 const COMPOUND_OPS = {
     "+=": "+", "-=": "-", "*=": "*", "/=": "/", "%=": "%",
 };
 // ── Statement extraction ─────────────────────────────────────
+/** Parse ghost and assert annotations from comment ranges. */
+function parseSpecComments(ranges, line) {
+    const result = [];
+    for (const range of ranges) {
+        const text = range.getText().trim();
+        if (!text.startsWith(PREFIX))
+            continue;
+        const content = text.slice(PREFIX.length);
+        // assert expr
+        if (content.startsWith("assert ")) {
+            result.push({ kind: "assert", expr: content.slice(7).trim(), line });
+            continue;
+        }
+        if (!content.startsWith("ghost "))
+            continue;
+        const ghostBody = content.slice(6).trim();
+        // ghost let varName: type = expr  OR  ghost let varName = expr
+        const letMatch = ghostBody.match(/^let\s+(\w+)(?:\s*:\s*(\w+))?\s*=\s*(.+)$/);
+        if (letMatch) {
+            result.push({ kind: "ghostLet", name: letMatch[1], tsType: letMatch[2] ?? null, init: letMatch[3].trim(), line });
+            continue;
+        }
+        // ghost varName = expr
+        const assignMatch = ghostBody.match(/^(\w+)\s*=\s*(.+)$/);
+        if (assignMatch) {
+            result.push({ kind: "ghostAssign", target: assignMatch[1], value: assignMatch[2].trim(), line });
+        }
+    }
+    return result;
+}
 function extractStmts(stmts) {
     const result = [];
     for (const s of stmts) {
         const line = s.getStartLineNumber();
+        // Ghost annotations from leading comments → inject before this statement
+        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*(.+)|(?:\s+(\S+)))?$/))
+                .find(m => m !== 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();
-                const init = d.getInitializer();
+                let init;
+                if (isHavoc && !havocKey) {
+                    // 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();
+                    _havocKey = havocKey;
+                    init = initializer ? extractExpr(initializer) : { kind: "var", name: "default" };
+                    _havocKey = null;
+                }
                 result.push({
                     kind: "let",
                     name: d.getName(),
                     mutable: s.getDeclarationKind() === "let",
-                    tsType: typeToString(declType),
-                    init: init ? extractExpr(init) : { kind: "var", name: "default" },
+                    tsType: havocType ?? _eraseGenerics(typeToString(declType)),
+                    init,
                     line,
                 });
             }
@@ -279,13 +552,29 @@ function extractStmts(stmts) {
         }
         if (Node.isForOfStatement(s)) {
             const init = s.getInitializer();
-            const varName = Node.isVariableDeclarationList(init) ? init.getDeclarations()[0]?.getName() ?? "_" : "_";
+            const names = [];
+            if (Node.isVariableDeclarationList(init)) {
+                const decl = init.getDeclarations()[0];
+                const nameNode = decl?.getNameNode();
+                if (nameNode && Node.isArrayBindingPattern(nameNode)) {
+                    for (const elem of nameNode.getElements()) {
+                        if (Node.isBindingElement(elem))
+                            names.push(elem.getNameNode().getText());
+                    }
+                }
+                else {
+                    names.push(decl?.getName() ?? "_");
+                }
+            }
+            else {
+                names.push("_");
+            }
             const bodyNode = s.getStatement();
             const bodyStmts = Node.isBlock(bodyNode) ? bodyNode.getStatements() : [bodyNode];
             const annots = collectAnnotations(s, bodyStmts);
             result.push({
                 kind: "forof",
-                varName,
+                names,
                 iterable: extractExpr(s.getExpression()),
                 invariants: annots.filter(a => a.kind === "invariant").map(a => a.expr),
                 doneWith: annots.find(a => a.kind === "done_with")?.expr ?? null,
@@ -344,8 +633,18 @@ function extractStmts(stmts) {
         }
         if (Node.isExpressionStatement(s)) {
             const expr = s.getExpression();
-            // x = e
-            if (Node.isBinaryExpression(expr) && expr.getOperatorToken().getText() === "=") {
+            // arr[i] = v → arr = arr.with(i, v)
+            if (Node.isBinaryExpression(expr) && expr.getOperatorToken().getText() === "=" && Node.isElementAccessExpression(expr.getLeft())) {
+                const left = expr.getLeft();
+                const obj = extractExpr(left.getExpression());
+                const idx = extractExpr(left.getArgumentExpression());
+                const val = extractExpr(expr.getRight());
+                const target = left.getExpression().getText();
+                const withCall = { kind: "call", fn: { kind: "field", obj, field: "with" }, args: [idx, val] };
+                result.push({ kind: "assign", target, value: withCall, line });
+                // x = e
+            }
+            else if (Node.isBinaryExpression(expr) && expr.getOperatorToken().getText() === "=") {
                 result.push({ kind: "assign", target: expr.getLeft().getText(), value: extractExpr(expr.getRight()), line });
                 // x += e, x -= e, etc.
             }
@@ -371,17 +670,77 @@ function extractStmts(stmts) {
             }
             continue;
         }
+        if (Node.isThrowStatement(s)) {
+            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
+    if (stmts.length > 0) {
+        const last = stmts[stmts.length - 1];
+        const line = last.getStartLineNumber();
+        for (const sib of last.getNextSiblings()) {
+            const text = sib.getText().trim();
+            if (!text.startsWith(PREFIX))
+                continue;
+            const content = text.slice(PREFIX.length);
+            // assert expr
+            if (content.startsWith("assert ")) {
+                result.push({ kind: "assert", expr: content.slice(7).trim(), line });
+                continue;
+            }
+            if (!content.startsWith("ghost "))
+                continue;
+            const ghostBody = content.slice(6).trim();
+            const letMatch = ghostBody.match(/^let\s+(\w+)(?:\s*:\s*(\w+))?\s*=\s*(.+)$/);
+            if (letMatch) {
+                result.push({ kind: "ghostLet", name: letMatch[1], tsType: letMatch[2] ?? null, init: letMatch[3].trim(), line });
+                continue;
+            }
+            const assignMatch = ghostBody.match(/^(\w+)\s*=\s*(.+)$/);
+            if (assignMatch) {
+                result.push({ kind: "ghostAssign", target: assignMatch[1], value: assignMatch[2].trim(), line });
+            }
+        }
+    }
     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()}: 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") {
@@ -391,36 +750,478 @@ 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",
+        name: fn.getName?.() ?? "<anonymous>",
+        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);
         }
     }
+    // Extract module-level const declarations
+    const constants = [];
+    for (const stmt of sourceFile.getStatements()) {
+        if (Node.isVariableStatement(stmt)) {
+            for (const decl of stmt.getDeclarationList().getDeclarations()) {
+                if (stmt.getDeclarationList().getFlags() & 2 /* const */) {
+                    const init = decl.getInitializer();
+                    // Skip huge string constants — they crash the verifier and have no verification value
+                    const initType = decl.getType();
+                    const isHugeString = (initType.isString() || initType.isStringLiteral()) && init.getText().length > 200;
+                    if (init && !isHugeString && !Node.isArrowFunction(init)) {
+                        try {
+                            constants.push({
+                                name: decl.getName(),
+                                tsType: typeToString(decl.getType()),
+                                value: extractExpr(init),
+                            });
+                        }
+                        catch (e) {
+                            console.error(`WARNING: skipping const '${decl.getName()}': ${e.message}`);
+                        }
+                    }
+                }
+            }
+        }
+    }
+    // Collect all function-like declarations: function declarations + const arrow functions
+    const allFns = [];
+    for (const fn of sourceFile.getFunctions()) {
+        allFns.push({ name: fn.getName() ?? "<anonymous>", node: fn });
+    }
+    // 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, parentStmt: stmt });
+                }
+            }
+        }
+    }
+    // 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'))
+            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 => {
+        // 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.
+    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") {
+                    collectNamesExpr(s.value);
+                }
+                if (s.kind === "return") {
+                    collectNamesExpr(s.value);
+                }
+                if (s.kind === "if") {
+                    collectNamesExpr(s.cond);
+                    collectNames(s.then);
+                    collectNames(s.else);
+                }
+                if (s.kind === "while") {
+                    collectNamesExpr(s.cond);
+                    collectNames(s.body);
+                }
+                if (s.kind === "forof") {
+                    collectNamesExpr(s.iterable);
+                    collectNames(s.body);
+                }
+                if (s.kind === "expr") {
+                    collectNamesExpr(s.expr);
+                }
+            }
+        }
+        function collectNamesExpr(e) {
+            if (e.kind === "var")
+                referencedNames.add(e.name);
+            if (e.kind === "binop") {
+                collectNamesExpr(e.left);
+                collectNamesExpr(e.right);
+            }
+            if (e.kind === "unop") {
+                collectNamesExpr(e.expr);
+            }
+            if (e.kind === "call") {
+                collectNamesExpr(e.fn);
+                e.args.forEach(collectNamesExpr);
+            }
+            if (e.kind === "field") {
+                collectNamesExpr(e.obj);
+            }
+            if (e.kind === "index") {
+                collectNamesExpr(e.obj);
+                collectNamesExpr(e.idx);
+            }
+            if (e.kind === "record") {
+                if (e.spread)
+                    collectNamesExpr(e.spread);
+                e.fields.forEach(f => collectNamesExpr(f.value));
+            }
+            if (e.kind === "arrayLiteral") {
+                e.elems.forEach(collectNamesExpr);
+            }
+            if (e.kind === "conditional") {
+                collectNamesExpr(e.cond);
+                collectNamesExpr(e.then);
+                collectNamesExpr(e.else);
+            }
+        }
+        for (const fn of functions) {
+            for (const p of fn.params)
+                referencedNames.add(p.tsType);
+            referencedNames.add(fn.returnType);
+            collectNames(fn.body);
+            // Also scan spec annotations for identifier references
+            for (const spec of [...fn.requires, ...fn.ensures]) {
+                for (const m of spec.matchAll(/\b([a-zA-Z_]\w*)\b/g)) {
+                    referencedNames.add(m[1]);
+                }
+            }
+        }
+        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));
+    const builtins = new Set(["Map", "Set", "Array", "String", "Number", "Boolean", "Promise", "Date", "RegExp", "Error"]);
+    function resolveType(t, locationNode) {
+        // Unwrap arrays and generics to find user-defined types
+        if (t.isArray()) {
+            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 && !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);
+                }
+            }
+        }
+    }
+    for (const f of fnsToExtract) {
+        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()) {
+        const methods = [];
+        for (const method of cls.getMethods()) {
+            if (!method.getFullText().includes('//@ verify'))
+                continue;
+            methods.push(extractFunction(method));
+        }
+        if (methods.length === 0)
+            continue;
+        const fields = [];
+        for (const prop of cls.getProperties()) {
+            fields.push({ name: prop.getName(), tsType: typeToString(prop.getType()) });
+        }
+        classes.push({ name: cls.getName() ?? "Anonymous", fields, methods });
+    }
     return {
         file: sourceFile.getFilePath(),
         typeDecls,
-        functions: sourceFile.getFunctions().map(extractFunction),
+        constants,
+        functions,
+        classes,
     };
 }
 // ── Main ─────────────────────────────────────────────────────
@@ -430,6 +1231,6 @@ if (process.argv[1]?.match(/extract\.(ts|js)$/)) {
         console.error("Usage: extract <file.ts>");
         process.exit(1);
     }
-    const proj = new Project({ compilerOptions: { strict: true } });
+    const proj = new Project({ compilerOptions: { strict: true, target: ScriptTarget.ESNext, lib: ["lib.esnext.d.ts"] } });
     console.log(JSON.stringify(extractModule(proj.addSourceFileAtPath(file)), null, 2));
 }