lemmascript 0.1.0 → 0.2.0

package/tools/dist/extract.js

@@ -4,13 +4,48 @@
  * 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;
 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,6 +59,10 @@ function extractExpr(node) {
     if (Node.isIdentifier(node)) {
         return { kind: "var", name: node.getText() };
     }
+    // this
+    if (node.getKind() === SyntaxKind.ThisKeyword) {
+        return { kind: "var", name: "this" };
+    }
     // Property access: x.foo
     if (Node.isPropertyAccessExpression(node)) {
         return { kind: "field", obj: extractExpr(node.getExpression()), field: node.getName() };
@@ -108,6 +147,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 +163,30 @@ 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();
+            const tsType = typeArgs && typeArgs.length > 0
+                ? `${name}<${typeArgs.map(t => t.getText()).join(", ")}>`
+                : name;
+            const args = node.getArguments();
+            // new Map(arr.map(fn)) — map-from-array constructor
+            if (name === "Map" && args && args.length === 1) {
+                return { kind: "call", fn: { kind: "var", name: "__mapFromArray" }, args: [extractExpr(args[0])] };
+            }
+            return { kind: "emptyCollection", collectionType: name, tsType };
+        }
+    }
+    // As-expression: expr as T — strip the type assertion
+    if (Node.isAsExpression(node)) {
+        return extractExpr(node.getExpression());
+    }
     throw new Error(`Unsupported expression: ${node.getText()}`);
 }
 // ── Annotation parsing ───────────────────────────────────────
@@ -224,39 +291,106 @@ 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()) {
+        return type.getUnionTypes().map(typeToString).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
+        result.push(...parseSpecComments(s.getLeadingCommentRanges(), line));
         if (Node.isVariableStatement(s)) {
+            const havocMatch = s.getLeadingCommentRanges()
+                .map(r => r.getText().trim().match(/^\/\/@ havoc(?:\s+(\S+))?$/))
+                .find(m => m !== null);
+            const havocKey = havocMatch?.[1] ?? null;
+            const isHavoc = !!havocMatch;
             for (const d of s.getDeclarations()) {
                 const declType = d.getType();
-                const init = d.getInitializer();
+                let init;
+                if (isHavoc && !havocKey) {
+                    init = { kind: "havoc", tsType: typeToString(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" },
+                    init,
                     line,
                 });
             }
@@ -279,13 +413,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 +494,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,15 +531,47 @@ function extractStmts(stmts) {
             }
             continue;
         }
+        if (Node.isThrowStatement(s)) {
+            result.push({ kind: "throw", line });
+            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) {
     const body = fn.getBody();
     if (!body || !Node.isBlock(body))
-        throw new Error(`${fn.getName()}: function body is not a block`);
+        throw new Error(`${fn.getName?.() ?? "arrow"}: function body is not a block`);
     const bodyStmts = body.getStatements();
     const annots = collectAnnotations(fn, bodyStmts);
     const typeAnnotations = [];
@@ -391,7 +583,7 @@ function extractFunction(fn) {
         }
     }
     return {
-        name: fn.getName() ?? "<anonymous>",
+        name: fn.getName?.() ?? "<anonymous>",
         params: fn.getParameters().map(p => ({ name: p.getName(), tsType: p.getTypeNode()?.getText() ?? "unknown" })),
         returnType: fn.getReturnTypeNode()?.getText() ?? "unknown",
         requires: annots.filter(a => a.kind === "requires").map(a => a.expr),
@@ -417,10 +609,189 @@ export function extractModule(sourceFile) {
                 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 = (...) => { ... } — treat as named function
+    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 });
+                }
+            }
+        }
+    }
+    // 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;
+    const functions = fnsToExtract.map(f => {
+        const raw = extractFunction(f.node);
+        raw.name = f.name; // use the const name, not "<anonymous>"
+        return raw;
+    });
+    // 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") {
+                    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)));
+    }
+    // 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;
+        }
+        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 (info) {
+                typeDecls.push(info);
+                knownTypes.add(name);
+            }
+        }
+    }
+    for (const f of fnsToExtract) {
+        for (const p of f.node.getParameters())
+            resolveType(p.getType(), p);
+    }
+    // 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 +801,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));
 }

package/tools/dist/ir.js

@@ -1,7 +1,7 @@
 /**
- * Lean IR — the intermediate representation between transform and emit.
+ * IR — the intermediate representation between transform and emit.
  *
  * The transform phase produces these types.
- * The emit phase pretty-prints them to Lean syntax.
+ * The emit phase pretty-prints them to backend syntax (Lean or Dafny).
  */
 export {};

package/tools/dist/lean-commands.js

@@ -0,0 +1,35 @@
+/**
+ * Lean backend commands: gen, check.
+ */
+import { existsSync, writeFileSync } from "fs";
+import { execSync } from "child_process";
+import path from "path";
+export function leanGen(typesPath, defPath, typesText, defText) {
+    if (typesPath && typesText) {
+        writeFileSync(typesPath, typesText);
+        console.log(`Generated: ${typesPath}`);
+    }
+    writeFileSync(defPath, defText);
+    console.log(`Generated: ${defPath}`);
+}
+export function leanCheck(dir, base) {
+    let lakeDir = dir;
+    while (lakeDir !== path.dirname(lakeDir)) {
+        if (existsSync(path.join(lakeDir, "lakefile.lean")))
+            break;
+        lakeDir = path.dirname(lakeDir);
+    }
+    const proofPath = path.join(dir, `${base}.proof.lean`);
+    if (!existsSync(proofPath)) {
+        console.error(`No proof file: ${proofPath}`);
+        return false;
+    }
+    console.log("Running lake build...");
+    try {
+        execSync(`lake build`, { cwd: lakeDir, stdio: "inherit" });
+        return true;
+    }
+    catch {
+        return false;
+    }
+}