lemmascript 0.1.0 → 0.3.0

package/tools/dist/resolve.js

@@ -6,6 +6,72 @@
  */
 import { parseTsType } from "./types.js";
 import { parseExpr } from "./specparser.js";
+// ── Raw expression substitution ─────────────────────────────
+let _synVarCounter = 0;
+/**
+ * Structural equality for raw field-access chains (var and field nodes only).
+ * Exact within a single expression scope — raw IR has no bindings that could
+ * cause name collisions (those are introduced by resolve, which runs after).
+ */
+function rawExprEquals(a, b) {
+    if (a.kind === "var" && b.kind === "var")
+        return a.name === b.name;
+    if (a.kind === "field" && b.kind === "field")
+        return a.field === b.field && rawExprEquals(a.obj, b.obj);
+    if (a.kind === "call" && b.kind === "call")
+        return rawExprEquals(a.fn, b.fn) && a.args.length === b.args.length && a.args.every((arg, i) => rawExprEquals(arg, b.args[i]));
+    if (a.kind === "index" && b.kind === "index")
+        return rawExprEquals(a.obj, b.obj) && rawExprEquals(a.idx, b.idx);
+    return false;
+}
+/** Return the root variable name of a field-access chain, or null. */
+function rawChainRoot(e) {
+    if (e.kind === "var")
+        return e.name;
+    if (e.kind === "field")
+        return rawChainRoot(e.obj);
+    return null;
+}
+/**
+ * Replace all occurrences of `target` in `expr` with `replacement`.
+ * Only matches field-access chains (see rawExprEquals). Stops at lambda
+ * boundaries that shadow the chain's root variable.
+ */
+function substituteRawExpr(expr, target, replacement) {
+    if (rawExprEquals(expr, target))
+        return replacement;
+    const root = rawChainRoot(target);
+    const sub = (e) => substituteRawExpr(e, target, replacement);
+    switch (expr.kind) {
+        case "var":
+        case "num":
+        case "str":
+        case "bool":
+        case "result":
+        case "havoc":
+        case "emptyCollection":
+            return expr;
+        case "binop": return { ...expr, left: sub(expr.left), right: sub(expr.right) };
+        case "unop": return { ...expr, expr: sub(expr.expr) };
+        case "call": return { ...expr, fn: sub(expr.fn), args: expr.args.map(sub) };
+        case "field": return { ...expr, obj: sub(expr.obj) };
+        case "index": return { ...expr, obj: sub(expr.obj), idx: sub(expr.idx) };
+        case "record":
+            return { ...expr, spread: expr.spread ? sub(expr.spread) : null,
+                fields: expr.fields.map(f => ({ ...f, value: sub(f.value) })) };
+        case "arrayLiteral": return { ...expr, elems: expr.elems.map(sub) };
+        case "conditional": return { ...expr, cond: sub(expr.cond), then: sub(expr.then), else: sub(expr.else) };
+        case "nonNull": return { ...expr, expr: sub(expr.expr) };
+        case "forall":
+        case "exists":
+            return { ...expr, body: sub(expr.body) };
+        case "lambda":
+            // Don't cross lambda boundaries that shadow the chain's root variable
+            if (root && expr.params.some(p => p.name === root))
+                return expr;
+            return { ...expr, body: Array.isArray(expr.body) ? expr.body : sub(expr.body) };
+    }
+}
 function lookup(env, name) {
     if (!env)
         return undefined;
@@ -33,12 +99,96 @@ function coerceStr(expr, targetTy) {
     return expr;
 }
 // ── Helpers ──────────────────────────────────────────────────
+/** Detect `v !== undefined` or `undefined !== v` where v: optional<T>. */
+function narrowOptional(cond, env) {
+    if (cond.kind !== "binop" || (cond.op !== "!==" && cond.op !== "==="))
+        return null;
+    // v !== undefined  OR  undefined !== v
+    let varName = null;
+    if (cond.left.kind === "var" && cond.right.kind === "var" && cond.right.name === "undefined")
+        varName = cond.left.name;
+    if (cond.right.kind === "var" && cond.left.kind === "var" && cond.left.name === "undefined")
+        varName = cond.right.name;
+    if (!varName)
+        return null;
+    const ty = lookup(env, varName);
+    if (!ty || ty.kind !== "optional")
+        return null;
+    return { varName, innerTy: ty.inner, inThen: cond.op === "!==" };
+}
+/** 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";
+}
 function findDecl(ctx, name) {
     return ctx.typeDecls.find(d => d.name === name);
 }
 function getDiscriminant(ctx, typeName) {
     return findDecl(ctx, typeName)?.discriminant;
 }
+/** Infer quantifier variable type from usage in body.
+ *  If the variable is used as a map/set key (e.g. map.has(k), map.get(k)),
+ *  return the collection's key type. Otherwise return null (default to int). */
+function inferQuantVarType(varName, body, ctx) {
+    // Look for calls like map.has(k), map.get(k), or array.includes(k) where k is our variable
+    if (body.kind === "call" && body.fn.kind === "field" &&
+        (body.fn.field === "has" || body.fn.field === "get" || body.fn.field === "includes") &&
+        body.args.length === 1 && body.args[0].kind === "var" && body.args[0].name === varName) {
+        const objTy = lookup(ctx.env, body.fn.obj.kind === "var" ? body.fn.obj.name : "");
+        if (objTy?.kind === "map")
+            return objTy.key;
+        if (objTy?.kind === "set")
+            return objTy.elem;
+        if (objTy?.kind === "array")
+            return objTy.elem;
+    }
+    // Recurse into subexpressions
+    if (body.kind === "binop") {
+        return inferQuantVarType(varName, body.left, ctx) ?? inferQuantVarType(varName, body.right, ctx);
+    }
+    if (body.kind === "unop")
+        return inferQuantVarType(varName, body.expr, ctx);
+    if (body.kind === "call") {
+        for (const a of body.args) {
+            const r = inferQuantVarType(varName, a, ctx);
+            if (r)
+                return r;
+        }
+        return inferQuantVarType(varName, body.fn, ctx);
+    }
+    if (body.kind === "field")
+        return inferQuantVarType(varName, body.obj, ctx);
+    if (body.kind === "index") {
+        return inferQuantVarType(varName, body.obj, ctx) ?? inferQuantVarType(varName, body.idx, ctx);
+    }
+    if (body.kind === "conditional") {
+        return inferQuantVarType(varName, body.cond, ctx) ??
+            inferQuantVarType(varName, body.then, ctx) ?? inferQuantVarType(varName, body.else, ctx);
+    }
+    if ((body.kind === "forall" || body.kind === "exists") && body.var !== varName) {
+        return inferQuantVarType(varName, body.body, ctx);
+    }
+    if (body.kind === "arrayLiteral") {
+        for (const el of body.elems) {
+            const r = inferQuantVarType(varName, el, ctx);
+            if (r)
+                return r;
+        }
+    }
+    if (body.kind === "record") {
+        if (body.spread) {
+            const r = inferQuantVarType(varName, body.spread, ctx);
+            if (r)
+                return r;
+        }
+        for (const f of body.fields) {
+            const r = inferQuantVarType(varName, f.value, ctx);
+            if (r)
+                return r;
+        }
+    }
+    return null;
+}
 function classifyCall(fn, ctx) {
     if (fn.kind === "field" && fn.obj.kind === "var" && fn.obj.name === "Math")
         return "pure";
@@ -59,35 +209,170 @@ function resolveExpr(e, ctx) {
         case "var":
             return { kind: "var", name: e.name, ty: lookup(ctx.env, e.name) ?? { kind: "unknown" } };
         case "num":
+            if (!Number.isInteger(e.value))
+                return { kind: "num", value: e.value, ty: { kind: "real" } };
             return { kind: "num", value: e.value, ty: e.value >= 0 ? { kind: "nat" } : { kind: "int" } };
         case "str":
             return { kind: "str", value: e.value, ty: { kind: "string" } };
         case "bool":
             return { kind: "bool", value: e.value, ty: { kind: "bool" } };
+        case "nonNull": {
+            const expr = resolveExpr(e.expr, ctx);
+            // Unwrap optional type; for map.get()!, force to direct access type
+            if (expr.kind === "call" && expr.fn.kind === "field" &&
+                expr.fn.obj.ty.kind === "map" && expr.fn.field === "get") {
+                return { ...expr, ty: expr.fn.obj.ty.value };
+            }
+            const ty = expr.ty.kind === "optional" ? expr.ty.inner : expr.ty;
+            return { ...expr, ty };
+        }
         case "binop": {
             let left = resolveExpr(e.left, ctx);
-            let right = resolveExpr(e.right, ctx);
+            // && narrowing: if left is "x !== undefined", narrow x for right side
+            let rightCtx = ctx;
+            if (e.op === "&&") {
+                const narrowed = narrowOptional(e.left, ctx.env);
+                if (narrowed && narrowed.inThen) {
+                    rightCtx = withEnv(ctx, extend(ctx.env, narrowed.varName, narrowed.innerTy));
+                }
+            }
+            let right = resolveExpr(e.right, rightCtx);
             if (e.op === "===" || e.op === "!==") {
                 left = coerceStr(left, right.ty);
                 right = coerceStr(right, left.ty);
             }
             let ty = { kind: "unknown" };
-            if (["===", "!==", ">=", "<=", ">", "<", "&&", "||"].includes(e.op))
+            if (["===", "!==", ">=", "<=", ">", "<", "in"].includes(e.op))
                 ty = { kind: "bool" };
-            else if (["+", "-", "*", "/", "%"].includes(e.op))
-                ty = left.ty;
+            else if (e.op === "&&")
+                ty = right.ty;
+            else if (e.op === "||" && left.ty.kind === "optional") {
+                // || undefined is identity for optionals — keep the optional type
+                ty = (e.right.kind === "var" && e.right.name === "undefined") ? left.ty : left.ty.inner;
+            }
+            else if (e.op === "||")
+                ty = right.ty;
+            else if (["+", "-", "*", "/", "%"].includes(e.op)) {
+                ty = (left.ty.kind === "real" || right.ty.kind === "real") ? { kind: "real" } : left.ty;
+            }
             return { kind: "binop", op: e.op, left, right, ty };
         }
         case "unop": {
             const expr = resolveExpr(e.expr, ctx);
             return { kind: "unop", op: e.op, expr, ty: e.op === "!" ? { kind: "bool" } : expr.ty };
         }
-        case "call":
-            return { kind: "call", fn: resolveExpr(e.fn, ctx), args: e.args.map(a => resolveExpr(a, ctx)), ty: { kind: "unknown" }, callKind: classifyCall(e.fn, ctx) };
+        case "call": {
+            const fn = resolveExpr(e.fn, ctx);
+            // Infer lambda param types from array method context (map, filter, etc.)
+            let rawArgs = e.args;
+            if (fn.kind === "field" && fn.obj.ty.kind === "array" &&
+                ["map", "filter", "every", "some", "find"].includes(fn.field) &&
+                rawArgs.length >= 1 && rawArgs[0].kind === "lambda" &&
+                rawArgs[0].params.length >= 1 && !rawArgs[0].params[0].tsType) {
+                const elemTy = fn.obj.ty.elem;
+                const tsType = elemTy.kind === "user" ? elemTy.name
+                    : elemTy.kind === "string" ? "string"
+                        : elemTy.kind === "int" || elemTy.kind === "nat" ? "number"
+                            : elemTy.kind === "bool" ? "boolean" : undefined;
+                if (tsType) {
+                    const lam = rawArgs[0];
+                    const updatedParams = [{ ...lam.params[0], tsType }, ...lam.params.slice(1)];
+                    rawArgs = [{ ...lam, params: updatedParams }, ...rawArgs.slice(1)];
+                }
+            }
+            // For .push() on a typed array, resolve the argument with element type context
+            // so record expressions can match fields and coerce types
+            let argCtx = ctx;
+            if (fn.kind === "field" && fn.obj.ty.kind === "array" && fn.field === "push" &&
+                fn.obj.ty.elem.kind === "user") {
+                argCtx = { ...ctx, returnTy: fn.obj.ty.elem };
+            }
+            let args = rawArgs.map(a => resolveExpr(a, argCtx));
+            // Coerce args: string literals to user types, non-optional to Option, pad missing optional args
+            if (fn.kind === "var" && ctx.fnParams.has(fn.name)) {
+                const paramTys = ctx.fnParams.get(fn.name);
+                args = args.map((a, i) => {
+                    if (i >= paramTys.length)
+                        return a;
+                    // Coerce string literal to user type (e.g., 'MissingList' → Err constructor)
+                    a = coerceStr(a, paramTys[i]);
+                    // Wrap non-optional in Some when callee expects optional param
+                    if (a.ty.kind !== "optional" && paramTys[i].kind === "optional") {
+                        return {
+                            kind: "call",
+                            fn: { kind: "var", name: "Some", ty: paramTys[i] },
+                            args: [a],
+                            ty: paramTys[i],
+                            callKind: "pure",
+                        };
+                    }
+                    return a;
+                });
+                // Pad missing optional args with None
+                for (let i = args.length; i < paramTys.length; i++) {
+                    if (paramTys[i].kind === "optional") {
+                        args.push({ kind: "var", name: "undefined", ty: paramTys[i] });
+                    }
+                }
+            }
+            let ty = { kind: "unknown" };
+            // Infer return types for collection methods
+            if (fn.kind === "field" && fn.obj.ty.kind === "map") {
+                if (fn.field === "get")
+                    ty = ctx.inSpec ? fn.obj.ty.value : { kind: "optional", inner: fn.obj.ty.value };
+                else if (fn.field === "has")
+                    ty = { kind: "bool" };
+                else if (fn.field === "set")
+                    ty = fn.obj.ty;
+                else if (fn.field === "delete")
+                    ty = fn.obj.ty;
+            }
+            else if (fn.kind === "field" && fn.obj.ty.kind === "set") {
+                if (fn.field === "has")
+                    ty = { kind: "bool" };
+                else if (fn.field === "add")
+                    ty = fn.obj.ty;
+                else if (fn.field === "delete")
+                    ty = fn.obj.ty;
+            }
+            else if (fn.kind === "field" && fn.obj.ty.kind === "array") {
+                if (fn.field === "includes")
+                    ty = { kind: "bool" };
+                else if (fn.field === "shift")
+                    ty = fn.obj.ty.elem;
+                else if (fn.field === "push")
+                    ty = fn.obj.ty;
+                else if (fn.field === "concat")
+                    ty = fn.obj.ty;
+                else if (fn.field === "map" && args.length >= 1 && args[0].kind === "lambda") {
+                    const retTy = args[0].body.length > 0 && args[0].body[0].kind === "return"
+                        ? args[0].body[0].value.ty : { kind: "unknown" };
+                    ty = { kind: "array", elem: retTy };
+                }
+                else if (fn.field === "filter")
+                    ty = fn.obj.ty;
+                else if (fn.field === "every" || fn.field === "some")
+                    ty = { kind: "bool" };
+            }
+            else if (fn.kind === "field" && fn.obj.ty.kind === "string") {
+                if (fn.field === "trim")
+                    ty = { kind: "string" };
+                else if (fn.field === "toLowerCase")
+                    ty = { kind: "string" };
+                else if (fn.field === "toUpperCase")
+                    ty = { kind: "string" };
+                else if (fn.field === "includes")
+                    ty = { kind: "bool" };
+            }
+            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);
-            return { kind: "index", obj, idx, ty: obj.ty.kind === "array" ? obj.ty.elem : { kind: "unknown" } };
+            const idxTy = obj.ty.kind === "array" ? obj.ty.elem
+                : obj.ty.kind === "map" ? obj.ty.value
+                    : { kind: "unknown" };
+            return { kind: "index", obj, idx, ty: idxTy };
         }
         case "field": {
             const obj = resolveExpr(e.obj, ctx);
@@ -96,6 +381,9 @@ function resolveExpr(e, ctx) {
             if (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")) {
+                ty = { kind: "nat" };
+            }
             else if (obj.ty.kind === "user") {
                 if (getDiscriminant(ctx, obj.ty.name) === e.field)
                     isDiscriminant = true;
@@ -114,11 +402,23 @@ function resolveExpr(e, ctx) {
             // Infer record type: from spread, or from return type context
             const recordTy = ty.kind === "user" ? ty : ctx.returnTy.kind === "user" ? ctx.returnTy : null;
             const decl = recordTy ? ctx.typeDecls.find(d => d.name === recordTy.name && d.kind === "record") : undefined;
+            // 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, ctx);
+                let value = resolveExpr(f.value, fieldCtx);
                 const fieldDecl = decl?.fields?.find(df => df.name === f.name);
-                if (fieldDecl)
-                    value = coerceStr(value, parseTsType(fieldDecl.tsType));
+                if (fieldDecl) {
+                    const declTy = parseTsType(fieldDecl.tsType);
+                    value = coerceStr(value, 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 = {
+                            kind: "call",
+                            fn: { kind: "var", name: "Some", ty: declTy },
+                            args: [value], ty: declTy, callKind: "pure",
+                        };
+                    }
+                }
                 return { name: f.name, value };
             });
             return { kind: "record", spread, fields, ty: recordTy ?? ty };
@@ -128,11 +428,13 @@ function resolveExpr(e, ctx) {
                 throw new Error("\\result is only valid in ensures");
             return { kind: "result", ty: ctx.returnTy };
         case "forall": {
-            const varTy = e.varType === "nat" ? { kind: "nat" } : { kind: "int" };
+            const varTy = e.varType !== "int" ? parseTsType(e.varType)
+                : inferQuantVarType(e.var, e.body, ctx) ?? { kind: "int" };
             return { kind: "forall", var: e.var, varTy, body: resolveExpr(e.body, withEnv(ctx, extend(ctx.env, e.var, varTy))), ty: { kind: "bool" } };
         }
         case "exists": {
-            const varTy = e.varType === "nat" ? { kind: "nat" } : { kind: "int" };
+            const varTy = e.varType !== "int" ? parseTsType(e.varType)
+                : inferQuantVarType(e.var, e.body, ctx) ?? { kind: "int" };
             return { kind: "exists", var: e.var, varTy, body: resolveExpr(e.body, withEnv(ctx, extend(ctx.env, e.var, varTy))), ty: { kind: "bool" } };
         }
         case "arrayLiteral": {
@@ -159,13 +461,64 @@ function resolveExpr(e, ctx) {
         }
         case "conditional": {
             const cond = resolveExpr(e.cond, ctx);
-            let then_ = resolveExpr(e.then, ctx);
+            // Optional truthiness: opt ? X : Y
+            // Narrow the optional to its inner type in the then-branch so that
+            // field accesses resolve correctly (e.g. entry.decision.field).
+            let narrowedVar;
+            let narrowedExprResolved;
+            let thenCtx = ctx;
+            let rawThen = e.then;
+            if (cond.ty.kind === "optional") {
+                const innerTy = cond.ty.inner;
+                if (e.cond.kind === "var") {
+                    narrowedVar = e.cond.name;
+                    thenCtx = withEnv(ctx, extend(ctx.env, e.cond.name, innerTy));
+                }
+                else {
+                    narrowedVar = `_opt${_synVarCounter++}`;
+                    rawThen = substituteRawExpr(e.then, e.cond, { kind: "var", name: narrowedVar });
+                    thenCtx = withEnv(ctx, extend(ctx.env, narrowedVar, innerTy));
+                }
+            }
+            // Explicit optional check: x !== undefined ? expr(x) : undefined
+            // Narrow x to its inner type in the then-branch.
+            if (!narrowedVar) {
+                const narrowed = narrowOptional(e.cond, ctx.env);
+                if (narrowed && narrowed.inThen) {
+                    narrowedVar = narrowed.varName;
+                    thenCtx = withEnv(ctx, extend(ctx.env, narrowed.varName, narrowed.innerTy));
+                }
+                // Handle complex optional expressions: f() !== undefined ? f().field : undefined
+                if (!narrowedVar && e.cond.kind === "binop" && e.cond.op === "!==" &&
+                    e.cond.right.kind === "var" && e.cond.right.name === "undefined") {
+                    const optExpr = e.cond.left;
+                    const resolvedOpt = resolveExpr(optExpr, ctx);
+                    if (resolvedOpt.ty.kind === "optional") {
+                        narrowedVar = `_opt${_synVarCounter++}`;
+                        narrowedExprResolved = resolvedOpt;
+                        rawThen = substituteRawExpr(e.then, optExpr, { kind: "var", name: narrowedVar });
+                        thenCtx = withEnv(ctx, extend(ctx.env, narrowedVar, resolvedOpt.ty.inner));
+                    }
+                }
+            }
+            let then_ = resolveExpr(rawThen, thenCtx);
             let else_ = resolveExpr(e.else, ctx);
             then_ = coerceStr(then_, else_.ty);
             else_ = coerceStr(else_, then_.ty);
-            const ty = then_.ty.kind !== "unknown" ? then_.ty : else_.ty;
-            return { kind: "conditional", cond, then: then_, else: else_, ty };
+            let ty = then_.ty.kind !== "unknown" ? then_.ty : else_.ty;
+            // When narrowedExpr is set, the transform will emit a match producing Optional
+            if (narrowedExprResolved && ty.kind !== "optional") {
+                ty = { kind: "optional", inner: ty };
+            }
+            return { kind: "conditional", cond, then: then_, else: else_, ty, narrowedVar, narrowedExpr: narrowedExprResolved };
         }
+        case "emptyCollection": {
+            const ty = parseTsType(e.tsType);
+            const elems = e.initElems ? e.initElems.map(el => resolveExpr(el, ctx)) : [];
+            return { kind: "arrayLiteral", elems, ty };
+        }
+        case "havoc":
+            return { kind: "havoc", ty: resolveTsType(e.tsType, ctx.overrides) };
     }
 }
 // ── Resolve specs ────────────────────────────────────────────
@@ -202,22 +555,54 @@ function resolveStmt(s, ctx) {
         case "let": {
             const ty = resolveTsType(s.tsType, ctx.overrides, s.name);
             const init = coerceStr(resolveExpr(s.init, ctx), ty);
-            return [{ kind: "let", name: s.name, ty, mutable: s.mutable, init }, extend(ctx.env, s.name, ty)];
+            // 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)];
         }
         case "assign": {
             const targetTy = lookup(ctx.env, s.target) ?? { kind: "unknown" };
             return [{ kind: "assign", target: s.target, value: coerceStr(resolveExpr(s.value, ctx), targetTy) }, ctx.env];
         }
-        case "return":
-            return [{ kind: "return", value: coerceStr(resolveExpr(s.value, ctx), ctx.returnTy) }, ctx.env];
+        case "return": {
+            let value = coerceStr(resolveExpr(s.value, ctx), ctx.returnTy);
+            // Wrap non-optional return value in Some when function returns optional
+            // Skip if already optional, void, or undefined (which maps to None)
+            const isUndef = value.kind === "var" && value.name === "undefined";
+            if (ctx.returnTy.kind === "optional" && value.ty.kind !== "optional" && !isUndef) {
+                value = {
+                    kind: "call",
+                    fn: { kind: "var", name: "Some", ty: ctx.returnTy },
+                    args: [value], ty: ctx.returnTy, callKind: "pure",
+                };
+            }
+            return [{ kind: "return", value }, ctx.env];
+        }
         case "break":
             return [{ kind: "break" }, ctx.env];
         case "continue":
             return [{ kind: "continue" }, ctx.env];
         case "expr":
             return [{ kind: "expr", expr: resolveExpr(s.expr, ctx) }, ctx.env];
-        case "if":
-            return [{ kind: "if", cond: resolveExpr(s.cond, ctx), then: resolveBlock(s.then, ctx), else: resolveBlock(s.else, ctx) }, ctx.env];
+        case "if": {
+            // Narrow optional<T> → T when checking !== undefined or undefined !==
+            let thenCtx = ctx, elseCtx = ctx;
+            const narrowed = narrowOptional(s.cond, ctx.env);
+            if (narrowed) {
+                const env = extend(ctx.env, narrowed.varName, narrowed.innerTy);
+                if (narrowed.inThen)
+                    thenCtx = withEnv(ctx, env);
+                else
+                    elseCtx = withEnv(ctx, env);
+            }
+            // Also narrow from left side of && condition: if (x !== undefined && ...) { ... }
+            if (!narrowed && s.cond.kind === "binop" && s.cond.op === "&&") {
+                const leftNarrowed = narrowOptional(s.cond.left, ctx.env);
+                if (leftNarrowed && leftNarrowed.inThen) {
+                    thenCtx = withEnv(ctx, extend(ctx.env, leftNarrowed.varName, leftNarrowed.innerTy));
+                }
+            }
+            return [{ kind: "if", cond: resolveExpr(s.cond, ctx), then: resolveBlock(s.then, thenCtx), else: resolveBlock(s.else, elseCtx) }, ctx.env];
+        }
         case "while": {
             const whileSpecCtx = { ...ctx, inSpec: true };
             return [{
@@ -231,24 +616,66 @@ function resolveStmt(s, ctx) {
         }
         case "forof": {
             const iterable = resolveExpr(s.iterable, ctx);
-            const elemTy = iterable.ty.kind === "array" ? iterable.ty.elem : { kind: "unknown" };
-            const idxName = `_${s.varName}_idx`;
-            const withIdx = extend(ctx.env, idxName, { kind: "nat" });
-            const withElem = extend(withIdx, s.varName, elemTy);
-            const bodyCtx = withEnv(ctx, withElem);
+            // Determine element types for each destructured name
+            const nameTypes = [];
+            let env = ctx.env;
+            if (s.names.length === 1) {
+                // Single name: element type from array/set
+                const elemTy = iterable.ty.kind === "array" ? iterable.ty.elem
+                    : iterable.ty.kind === "set" ? iterable.ty.elem
+                        : { kind: "unknown" };
+                nameTypes.push(elemTy);
+            }
+            else if (s.names.length >= 2 && iterable.ty.kind === "map") {
+                // Map destructuring: [key, value]
+                nameTypes.push(iterable.ty.key, iterable.ty.value);
+            }
+            else {
+                // General tuple destructuring: all unknown
+                for (const _ of s.names)
+                    nameTypes.push({ kind: "unknown" });
+            }
+            const idxName = `_${s.names[0]}_idx`;
+            env = extend(env, idxName, { kind: "nat" });
+            for (let j = 0; j < s.names.length; j++) {
+                env = extend(env, s.names[j], nameTypes[j] ?? { kind: "unknown" });
+            }
+            const bodyCtx = withEnv(ctx, env);
             return [{
-                    kind: "forof", varName: s.varName, varTy: elemTy, iterable,
+                    kind: "forof", names: s.names, nameTypes, iterable,
                     invariants: resolveSpecs(s.invariants, { ...bodyCtx, inSpec: true }),
                     doneWith: s.doneWith ? resolveSpec(s.doneWith, { ...bodyCtx, inSpec: true }) : null,
                     body: resolveBlock(s.body, bodyCtx),
                 }, ctx.env];
         }
+        case "throw":
+            return [{ kind: "throw" }, ctx.env];
         case "switch":
             return [{
                     kind: "switch", expr: resolveExpr(s.expr, ctx), discriminant: s.discriminant,
                     cases: s.cases.map(c => ({ label: c.label, body: resolveBlock(c.body, ctx) })),
                     defaultBody: resolveBlock(s.defaultBody, ctx),
                 }, ctx.env];
+        case "ghostLet": {
+            const specCtx = { ...ctx, inSpec: true };
+            // Handle new Set<T>() / new Map<K,V>() constructors
+            const collMatch = s.init.match(/^new\s+(Set|Map)<(.+)>\(\)$/);
+            const init = collMatch
+                ? resolveExpr({ kind: "emptyCollection", collectionType: collMatch[1], tsType: `${collMatch[1]}<${collMatch[2]}>` }, specCtx)
+                : resolveExpr(parseExpr(s.init), specCtx);
+            const ty = s.tsType ? parseTsType(s.tsType) : init.ty;
+            return [{ kind: "ghostLet", name: s.name, ty, init }, extend(ctx.env, s.name, ty)];
+        }
+        case "ghostAssign": {
+            const specCtx = { ...ctx, inSpec: true };
+            const value = resolveExpr(parseExpr(s.value), specCtx);
+            return [{ kind: "ghostAssign", target: s.target, value }, ctx.env];
+        }
+        case "assert": {
+            const specCtx = { ...ctx, inSpec: true };
+            const expr = resolveExpr(parseExpr(s.expr), specCtx);
+            return [{ kind: "assert", expr }, ctx.env];
+        }
     }
 }
 // ── Pure / return-in-loop detection ──────────────────────────
@@ -259,7 +686,7 @@ function isSyntacticallyPure(stmts) {
             case "while":
             case "forof": return false;
             case "let":
-                if (s.mutable)
+                if (s.mutable || s.init.kind === "havoc")
                     return false;
                 break;
             case "if":
@@ -420,32 +847,71 @@ function containsReturn(stmts) {
     return false;
 }
 // ── Resolve function / module ────────────────────────────────
-function resolveFunction(fn, typeDecls, pureFns) {
-    if (hasReturnInLoop(fn.body)) {
-        throw new Error(`${fn.name}: return inside a loop is not supported.`);
-    }
+function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map()) {
     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");
     let env = null;
     for (const p of params)
         env = extend(env, p.name, p.ty);
-    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, inSpec: false, inLambda: false };
+    const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, inSpec: false, inLambda: false };
     const requiresCtx = { ...baseCtx, inSpec: true };
     const ensuresCtx = { ...baseCtx, allowResult: true, inSpec: true };
     return {
-        name: fn.name, params, returnTy,
+        name: fn.name, typeParams: fn.typeParams, params, returnTy,
         requires: resolveSpecs(fn.requires, requiresCtx),
         ensures: resolveSpecs(fn.ensures, ensuresCtx),
         isPure: pureFns.has(fn.name),
         body: resolveBlock(fn.body, baseCtx),
     };
 }
+function resolveClass(cls, typeDecls, pureFns, fnParams = 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 })) };
+    const allTypeDecls = [...typeDecls, thisDecl];
+    const methods = cls.methods.map(fn => {
+        // Add 'this' to the environment
+        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");
+        let env = null;
+        env = extend(env, "this", thisType);
+        for (const p of params)
+            env = extend(env, p.name, p.ty);
+        const baseCtx = { env, typeDecls: allTypeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, inSpec: false, inLambda: false };
+        const requiresCtx = { ...baseCtx, inSpec: true };
+        const ensuresCtx = { ...baseCtx, allowResult: true, inSpec: true };
+        return {
+            name: fn.name, typeParams: fn.typeParams, params, returnTy,
+            requires: resolveSpecs(fn.requires, requiresCtx),
+            ensures: resolveSpecs(fn.ensures, ensuresCtx),
+            isPure: false, // class methods are never pure (they access this)
+            body: resolveBlock(fn.body, baseCtx),
+        };
+    });
+    return { name: cls.name, fields, methods };
+}
 export function resolveModule(raw) {
     const pureFns = computePureFns(raw.functions);
+    // Pre-compute function parameter types for optional coercion
+    const fnParams = 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)));
+    }
+    const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, inSpec: false, inLambda: false };
+    const constants = (raw.constants ?? []).map(c => ({
+        name: c.name,
+        ty: parseTsType(c.tsType),
+        value: resolveExpr(c.value, emptyCtx),
+    }));
     return {
         file: raw.file,
         typeDecls: raw.typeDecls,
-        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns)),
+        constants,
+        functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams)),
+        classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams)),
     };
 }