lemmascript 0.3.1 → 0.3.3

package/tools/dist/resolve.js

@@ -108,7 +108,9 @@ function wrapSome(value, optionalTy) {
 }
 /** Detect `v !== undefined` or `undefined !== v` where v: optional<T>.
  *  Handles simple variables, field access chains, and arbitrary expressions.
- *  When `fieldExpr` is returned, callers must use `substituteRawExpr` to narrow.
+ *  When `fieldExpr` is returned for a simple field chain (obj.field), callers
+ *  can narrow via `narrowedFields` context. For complex expressions (calls, etc.),
+ *  callers should fall back to `substituteRawExpr`.
  *
  *  Does NOT recurse into `&&` — callers that need to detect optional checks
  *  inside `&&` conditions should check `cond.left` explicitly. */
@@ -141,6 +143,19 @@ function detectOptionalCheck(cond, ctx) {
     }
     return null;
 }
+/** Collect all optional narrowings from an early-return condition.
+ *  Handles single checks (x === undefined) and compound || chains
+ *  (x === undefined || y === undefined). */
+function collectEarlyReturnNarrowings(cond, ctx) {
+    if (cond.kind === "binop" && cond.op === "||") {
+        return [...collectEarlyReturnNarrowings(cond.left, ctx), ...collectEarlyReturnNarrowings(cond.right, ctx)];
+    }
+    const narrowed = detectOptionalCheck(cond, ctx);
+    if (narrowed && !narrowed.inThen && !narrowed.fieldExpr) {
+        return [{ varName: narrowed.varName, innerTy: narrowed.innerTy }];
+    }
+    return [];
+}
 /** TS reference types that become value types in Dafny/Lean — const bindings need mutable var. */
 function isRefMutableInTS(ty) {
     return ty.kind === "array" || ty.kind === "map" || ty.kind === "set";
@@ -293,6 +308,8 @@ function inferMethodReturnTy(fn, args, ctx) {
     else if (objTy.kind === "array") {
         if (fn.field === "includes")
             return { kind: "bool" };
+        if (fn.field === "indexOf")
+            return { kind: "int" };
         if (fn.field === "shift")
             return objTy.elem;
         if (fn.field === "push" || fn.field === "concat")
@@ -387,15 +404,28 @@ function resolveExpr(e, ctx) {
                 fn.obj.ty.elem.kind === "user") {
                 argCtx = { ...ctx, returnTy: fn.obj.ty.elem };
             }
-            const args = coerceCallArgs(rawArgs.map(a => resolveExpr(a, argCtx)), fn, ctx);
-            const ty = inferMethodReturnTy(fn, args, ctx);
+            // Propagate parameter types to arguments for record literal resolution
+            // (enables inline discriminated union construction in function arguments)
+            const paramTypes = fn.kind === "var" && ctx.fnParams.has(fn.name) ? ctx.fnParams.get(fn.name) : null;
+            const args = coerceCallArgs(rawArgs.map((a, i) => {
+                let aCtx = argCtx;
+                if (paramTypes && i < paramTypes.length && paramTypes[i].kind === "user") {
+                    aCtx = { ...aCtx, returnTy: paramTypes[i] };
+                }
+                return resolveExpr(a, aCtx);
+            }), fn, ctx);
+            let ty = inferMethodReturnTy(fn, args, ctx);
+            // For same-file function calls, use the known return type
+            if (ty.kind === "unknown" && fn.kind === "var" && ctx.fnReturns.has(fn.name)) {
+                ty = ctx.fnReturns.get(fn.name);
+            }
             return { kind: "call", fn, args, ty, callKind: classifyCall(e.fn, ctx) };
         }
         case "index": {
             const obj = resolveExpr(e.obj, ctx);
             const idx = resolveExpr(e.idx, ctx);
             const idxTy = obj.ty.kind === "array" ? obj.ty.elem
-                : obj.ty.kind === "map" ? obj.ty.value
+                : obj.ty.kind === "map" ? { kind: "optional", inner: obj.ty.value }
                     : { kind: "unknown" };
             return { kind: "index", obj, idx, ty: idxTy };
         }
@@ -403,16 +433,24 @@ function resolveExpr(e, ctx) {
             const obj = resolveExpr(e.obj, ctx);
             let isDiscriminant = false;
             let ty = { kind: "unknown" };
-            if (e.field === "length" && (obj.ty.kind === "array" || obj.ty.kind === "string")) {
+            // Check narrowed field context (from conditional optional checks on field chains)
+            if (obj.kind === "var" && ctx.narrowedFields.length > 0) {
+                const nf = ctx.narrowedFields.find(n => n.objName === obj.name && n.fieldName === e.field);
+                if (nf)
+                    ty = nf.narrowedTy;
+            }
+            if (ty.kind === "unknown" && e.field === "length" && (obj.ty.kind === "array" || obj.ty.kind === "string")) {
                 ty = { kind: "nat" };
             }
-            else if (e.field === "size" && (obj.ty.kind === "map" || obj.ty.kind === "set")) {
+            else if (ty.kind === "unknown" && e.field === "size" && (obj.ty.kind === "map" || obj.ty.kind === "set")) {
                 ty = { kind: "nat" };
             }
-            else if (obj.ty.kind === "user") {
-                if (getDiscriminant(ctx, obj.ty.name) === e.field)
+            else if (ty.kind === "unknown" && obj.ty.kind === "user") {
+                // Strip generic args for type lookup: "Result<Model, Err>" → "Result"
+                const baseTyName = obj.ty.name.includes("<") ? obj.ty.name.slice(0, obj.ty.name.indexOf("<")) : obj.ty.name;
+                if (getDiscriminant(ctx, baseTyName) === e.field)
                     isDiscriminant = true;
-                const decl = findDecl(ctx, obj.ty.name);
+                const decl = findDecl(ctx, baseTyName);
                 if (decl?.kind === "record") {
                     const f = decl.fields?.find(f => f.name === e.field);
                     if (f)
@@ -440,11 +478,20 @@ function resolveExpr(e, ctx) {
             // Clear returnTy for field values — it applies to THIS record, not nested ones
             const fieldCtx = recordTy ? { ...ctx, returnTy: { kind: "unknown" } } : ctx;
             const fields = e.fields.map(f => {
-                let value = resolveExpr(f.value, fieldCtx);
                 const fieldDecl = decl?.fields?.find(df => df.name === f.name);
+                // Propagate declared field type into context so nested records resolve
+                // their union variant correctly (e.g., { kind: 'Idle' } → EffectMode.Idle)
+                const valueCtx = (fieldDecl?.type?.kind === "user")
+                    ? { ...fieldCtx, returnTy: fieldDecl.type }
+                    : fieldCtx;
+                let value = resolveExpr(f.value, valueCtx);
                 if (fieldDecl) {
                     const declTy = fieldDecl.type;
                     value = coerceStr(value, declTy);
+                    // Empty {} for map-typed fields → empty map (arrayLiteral with map type → emptyMap in transform)
+                    if (value.kind === "record" && value.fields.length === 0 && !value.spread && declTy.kind === "map") {
+                        value = { kind: "arrayLiteral", elems: [], ty: declTy };
+                    }
                     // Coerce non-optional to optional: wrap in Some (only when value type is concrete)
                     if (declTy.kind === "optional" && value.ty.kind !== "optional" && value.ty.kind !== "void" && value.ty.kind !== "unknown") {
                         value = wrapSome(value, declTy);
@@ -493,7 +540,6 @@ function resolveExpr(e, ctx) {
         case "conditional": {
             const cond = resolveExpr(e.cond, ctx);
             let narrowedVar;
-            let narrowedExprResolved;
             let thenCtx = ctx;
             let rawThen = e.then;
             // Phase 1: Optional truthiness — cond itself is optional (e.g. opt ? X : Y)
@@ -509,23 +555,56 @@ function resolveExpr(e, ctx) {
                     thenCtx = withEnv(ctx, extend(ctx.env, narrowedVar, innerTy));
                 }
             }
-            // Phase 2: Explicit check — v !== undefined (simple vars, field chains,
-            // complex expressions all handled uniformly by detectOptionalCheck)
+            // Phase 2/3: Explicit check — v !== undefined, or && with optional check.
+            // Resolve only narrows the type environment; transform handles all structural
+            // narrowing (match generation, variable binding, && splitting).
+            let narrowedExprResolved;
+            let elseCtx = ctx;
             if (!narrowedVar) {
                 const narrowed = detectOptionalCheck(e.cond, ctx)
-                    // Phase 3: && with optional check — (v !== undefined && ...) ? ... : ...
                     ?? (e.cond.kind === "binop" && e.cond.op === "&&" ? detectOptionalCheck(e.cond.left, ctx) : null);
                 if (narrowed && narrowed.inThen) {
-                    narrowedVar = narrowed.varName;
-                    thenCtx = withEnv(ctx, extend(ctx.env, narrowed.varName, narrowed.innerTy));
-                    if (narrowed.fieldExpr) {
+                    if (!narrowed.fieldExpr) {
+                        // Simple var: extend env only — transform will detect and generate match
+                        thenCtx = withEnv(thenCtx, extend(thenCtx.env, narrowed.varName, narrowed.innerTy));
+                    }
+                    else if (narrowed.fieldExpr.kind === "field" && narrowed.fieldExpr.obj.kind === "var") {
+                        // Simple field chain (obj.field): narrow via field context — no substitution
+                        thenCtx = {
+                            ...thenCtx,
+                            narrowedFields: [...thenCtx.narrowedFields, {
+                                    objName: narrowed.fieldExpr.obj.name,
+                                    fieldName: narrowed.fieldExpr.field,
+                                    narrowedTy: narrowed.innerTy,
+                                }],
+                        };
+                    }
+                    else {
+                        // Complex expression (call result, deep chain, etc.): transform can't detect these,
+                        // so keep old behavior — substitute + narrowedVar + narrowedExpr
+                        narrowedVar = narrowed.varName;
                         narrowedExprResolved = narrowed.narrowedExpr ?? resolveExpr(narrowed.fieldExpr, ctx);
                         rawThen = substituteRawExpr(e.then, narrowed.fieldExpr, { kind: "var", name: narrowed.varName });
+                        thenCtx = withEnv(thenCtx, extend(thenCtx.env, narrowed.varName, narrowed.innerTy));
+                    }
+                }
+                else if (narrowed && !narrowed.inThen && !narrowed.fieldExpr) {
+                    // v === undefined: narrow v in the else branch
+                    elseCtx = withEnv(elseCtx, extend(elseCtx.env, narrowed.varName, narrowed.innerTy));
+                }
+                // Compound || with === undefined: narrow all checked vars in else branch
+                // e.g. if (a === undefined || b === undefined) then X else Y → narrow a,b in Y
+                // TODO: resolve-time narrowing works but transform doesn't emit match unwrap
+                // for || conditions yet — decompose || into nested matches in transform.
+                // Workaround: split || into separate if guards in user code.
+                if (!narrowed && e.cond.kind === "binop" && e.cond.op === "||") {
+                    for (const n of collectEarlyReturnNarrowings(e.cond, ctx)) {
+                        elseCtx = withEnv(elseCtx, extend(elseCtx.env, n.varName, n.innerTy));
                     }
                 }
             }
             let then_ = resolveExpr(rawThen, thenCtx);
-            let else_ = resolveExpr(e.else, ctx);
+            let else_ = resolveExpr(e.else, elseCtx);
             then_ = coerceStr(then_, else_.ty);
             else_ = coerceStr(else_, then_.ty);
             let ty = then_.ty.kind !== "unknown" ? then_.ty : else_.ty;
@@ -537,8 +616,7 @@ function resolveExpr(e, ctx) {
                 ty = { kind: "optional", inner: then_.ty };
             }
             // When narrowedExpr is set AND a branch is void, the match produces Optional
-            const hasVoidBranch = then_.ty.kind === "void" || else_.ty.kind === "void";
-            if (narrowedExprResolved && hasVoidBranch && ty.kind !== "optional") {
+            if (narrowedExprResolved && (then_.ty.kind === "void" || else_.ty.kind === "void") && ty.kind !== "optional") {
                 ty = { kind: "optional", inner: ty };
             }
             return { kind: "conditional", cond, then: then_, else: else_, ty, narrowedVar, narrowedExpr: narrowedExprResolved };
@@ -579,12 +657,13 @@ function resolveBlock(stmts, ctx) {
         result.push(typed);
         env = nextEnv;
         // Flow narrowing: if (x === undefined) { return } narrows x for rest of block.
+        // Also handles compound: if (x === undefined || y === undefined) { return }
         // Field chains are excluded — resolve can't substitute in statement lists;
         // transform's emitOptionalMatch handles field chains in statement contexts.
         if (s.kind === "if" && s.then.length > 0 && s.then[s.then.length - 1].kind === "return" && s.else.length === 0) {
-            const narrowed = detectOptionalCheck(s.cond, withEnv(ctx, env));
-            if (narrowed && !narrowed.inThen && !narrowed.fieldExpr) {
-                env = extend(env, narrowed.varName, narrowed.innerTy);
+            const narrowings = collectEarlyReturnNarrowings(s.cond, withEnv(ctx, env));
+            for (const n of narrowings) {
+                env = extend(env, n.varName, n.innerTy);
             }
         }
     }
@@ -593,8 +672,13 @@ function resolveBlock(stmts, ctx) {
 function resolveStmt(s, ctx) {
     switch (s.kind) {
         case "let": {
-            const ty = resolveTsType(s.tsType, ctx.overrides, s.name);
-            const init = coerceStr(resolveExpr(s.init, ctx), ty);
+            const declTy = resolveTsType(s.tsType, ctx.overrides, s.name);
+            // Propagate declared type as returnTy so nested record expressions
+            // resolve union variants correctly (e.g., EffectState → mode: EffectMode → { kind: 'Idle' })
+            const initCtx = declTy.kind === "user" ? { ...ctx, returnTy: declTy } : ctx;
+            const init = coerceStr(resolveExpr(s.init, initCtx), declTy);
+            // Map indexing: TS says T, but access can fail → use Optional<T> from init
+            const ty = (declTy.kind !== "optional" && init.ty.kind === "optional") ? init.ty : declTy;
             // const collections are mutable in value-semantics world (TS mutates in place, Dafny/Lean reassign)
             const mutable = s.mutable || isRefMutableInTS(ty);
             return [{ kind: "let", name: s.name, ty, mutable, init }, extend(ctx.env, s.name, ty)];
@@ -825,6 +909,8 @@ function collectCallsStmts(stmts, fns, out) {
 }
 function computePureFns(functions) {
     const allFnNames = new Set(functions.map(fn => fn.name));
+    // //@ pure functions are always considered pure — never taint callers
+    const forcePure = new Set(functions.filter(fn => fn.pure).map(fn => fn.name));
     // Build call graph: fn → set of same-file functions it calls
     const callGraph = new Map();
     for (const fn of functions) {
@@ -832,8 +918,8 @@ function computePureFns(functions) {
         collectCallsStmts(fn.body, allFnNames, calls);
         callGraph.set(fn.name, calls);
     }
-    // Seed: syntactically non-pure functions
-    const nonPure = new Set(functions.filter(fn => !isSyntacticallyPure(fn.body)).map(fn => fn.name));
+    // Seed: syntactically non-pure functions (skip //@ pure)
+    const nonPure = new Set(functions.filter(fn => !forcePure.has(fn.name) && !isSyntacticallyPure(fn.body)).map(fn => fn.name));
     // Build reverse graph: fn → set of functions that call it
     const callers = new Map();
     for (const name of allFnNames)
@@ -842,12 +928,12 @@ function computePureFns(functions) {
         for (const callee of callees)
             callers.get(callee).add(caller);
     }
-    // Propagate impurity through reverse call graph
+    // Propagate impurity through reverse call graph (skip //@ pure)
     const worklist = [...nonPure];
     while (worklist.length > 0) {
         const fn = worklist.pop();
         for (const caller of callers.get(fn) ?? []) {
-            if (!nonPure.has(caller)) {
+            if (!nonPure.has(caller) && !forcePure.has(caller)) {
                 nonPure.add(caller);
                 worklist.push(caller);
             }
@@ -880,7 +966,7 @@ function containsReturn(stmts) {
     return false;
 }
 // ── Resolve function / module ────────────────────────────────
-function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), opts) {
+function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), opts) {
     const overrides = new Map(fn.typeAnnotations.map(a => [a.name, a.type]));
     const params = fn.params.map(p => ({ name: p.name, ty: resolveTsType(p.tsType, overrides, p.name) }));
     const returnTy = resolveTsType(fn.returnType, overrides, "\\result");
@@ -889,7 +975,7 @@ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), opts) {
         env = extend(env, opts.thisBinding.name, opts.thisBinding.ty);
     for (const p of params)
         env = extend(env, p.name, p.ty);
-    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, inSpec: false, inLambda: false };
+    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedFields: [] };
     const requiresCtx = { ...baseCtx, inSpec: true };
     const ensuresCtx = { ...baseCtx, allowResult: true, inSpec: true };
     // Apply type parameter constraints from //@ type T (==) annotations
@@ -901,17 +987,19 @@ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), opts) {
         name: fn.name, typeParams, params, returnTy,
         requires: resolveSpecs(fn.requires, requiresCtx),
         ensures: resolveSpecs(fn.ensures, ensuresCtx),
+        decreases: fn.decreases ? resolveSpec(fn.decreases, requiresCtx) : null,
         isPure: opts?.forcePure !== undefined ? opts.forcePure : pureFns.has(fn.name),
+        forcePure: fn.pure,
         body: resolveBlock(fn.body, baseCtx),
     };
 }
-function resolveClass(cls, typeDecls, pureFns, fnParams = new Map()) {
+function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map()) {
     const fields = cls.fields.map(f => ({ name: f.name, ty: parseTsType(f.tsType) }));
     // Create a synthetic record type for 'this' so field access resolves
     const thisType = { kind: "user", name: cls.name };
     const thisDecl = { name: cls.name, kind: "record", fields: cls.fields.map(f => ({ name: f.name, tsType: f.tsType, type: parseTsType(f.tsType) })) };
     const allTypeDecls = [...typeDecls, thisDecl];
-    const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, {
+    const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, fnReturns, {
         thisBinding: { name: "this", ty: thisType },
         forcePure: false, // class methods are never pure (they access this)
     }));
@@ -933,15 +1021,18 @@ function precomputeFieldTypes(typeDecls) {
     }
 }
 export function resolveModule(raw) {
+    _synVarCounter = 0;
     precomputeFieldTypes(raw.typeDecls);
     const pureFns = computePureFns(raw.functions);
-    // Pre-compute function parameter types for optional coercion
+    // Pre-compute function parameter and return types
     const fnParams = new Map();
+    const fnReturns = new Map();
     for (const fn of raw.functions) {
         const overrides = new Map(fn.typeAnnotations.map(a => [a.name, a.type]));
         fnParams.set(fn.name, fn.params.map(p => resolveTsType(p.tsType, overrides, p.name)));
+        fnReturns.set(fn.name, resolveTsType(fn.returnType, overrides, "\\result"));
     }
-    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, inSpec: false, inLambda: false };
+    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedFields: [] };
     const constants = (raw.constants ?? []).map(c => ({
         name: c.name,
         ty: parseTsType(c.tsType),
@@ -951,7 +1042,7 @@ export function resolveModule(raw) {
         file: raw.file,
         typeDecls: raw.typeDecls,
         constants,
-        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams)),
-        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams)),
+        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams, fnReturns)),
+        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams, fnReturns)),
     };
 }