lemmascript 0.4.0 → 0.5.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -14,6 +14,15 @@ function lookup(env, name) {
14
14
  function extend(env, name, ty) {
15
15
  return { name, ty, parent: env };
16
16
  }
17
+ function envKeys(env) {
18
+ const out = [];
19
+ let e = env;
20
+ while (e) {
21
+ out.push(e.name);
22
+ e = e.parent;
23
+ }
24
+ return out;
25
+ }
17
26
  function asRawAccessPath(e) {
18
27
  if (e.kind === "var")
19
28
  return { rootVar: e.name, fields: [] };
@@ -55,6 +64,43 @@ function wrapSome(value, optionalTy) {
55
64
  args: [value], ty: optionalTy, callKind: "pure",
56
65
  };
57
66
  }
67
+ /** Find the synth array-union TypeDecl named `name` (discriminant `__isArray__`). */
68
+ function findSynthArrayUnion(name, typeDecls) {
69
+ const decl = typeDecls.find(d => d.name === name);
70
+ if (decl?.kind === "discriminated-union" && decl.discriminant === "__isArray__")
71
+ return decl;
72
+ return null;
73
+ }
74
+ /** Coerce `value` to `targetTy` at an assignment-position. Mirrors TS subtyping
75
+ * for the two upcast shapes LS synthesizes:
76
+ * - `T` into `optional<T>` slot → wrap with `Some(...)`
77
+ * - `T[]` into a synth `T[] | U` slot → wrap with `ArrayBranch(...)`
78
+ * - `U` into a synth `T[] | U` slot → wrap with `NonArrayBranch(...)`
79
+ * Returns `value` unchanged if no coercion applies (types already match,
80
+ * source is unknown, or no rule matches). */
81
+ function coerceToTargetTy(value, targetTy, typeDecls) {
82
+ if (value.ty.kind === "unknown" || value.ty.kind === "void")
83
+ return value;
84
+ if (targetTy.kind === "optional" && value.ty.kind !== "optional") {
85
+ return wrapSome(value, targetTy);
86
+ }
87
+ if (targetTy.kind === "user") {
88
+ const synth = findSynthArrayUnion(targetTy.name, typeDecls);
89
+ if (synth && synth.variants && synth.variants.length === 2) {
90
+ const arrVariant = synth.variants.find(v => v.name === "ArrayBranch");
91
+ const nonVariant = synth.variants.find(v => v.name === "NonArrayBranch");
92
+ if (value.ty.kind === "array" && arrVariant) {
93
+ return { kind: "call", fn: { kind: "var", name: "ArrayBranch", ty: targetTy },
94
+ args: [value], ty: targetTy, callKind: "pure" };
95
+ }
96
+ if (value.ty.kind !== "array" && nonVariant) {
97
+ return { kind: "call", fn: { kind: "var", name: "NonArrayBranch", ty: targetTy },
98
+ args: [value], ty: targetTy, callKind: "pure" };
99
+ }
100
+ }
101
+ }
102
+ return value;
103
+ }
58
104
  /** Detect optional checks: `v !== undefined` (positive narrows then-branch),
59
105
  * `v === undefined` (negative narrows else-branch), or `!v` (equivalent to
60
106
  * `=== undefined`).
@@ -94,6 +140,12 @@ function classifyOptExpr(e, ctx) {
94
140
  return null;
95
141
  return { varName: e.name, innerTy: ty.inner };
96
142
  }
143
+ if (e.kind === "result") {
144
+ const ty = lookup(ctx.env, "\\result");
145
+ if (!ty || ty.kind !== "optional")
146
+ return null;
147
+ return { varName: "\\result", innerTy: ty.inner };
148
+ }
97
149
  const resolved = resolveExpr(e, ctx);
98
150
  if (resolved.ty.kind !== "optional")
99
151
  return null;
@@ -193,11 +245,109 @@ function isRefMutableInTS(ty) {
193
245
  return ty.kind === "array" || ty.kind === "map" || ty.kind === "set";
194
246
  }
195
247
  function findDecl(ctx, name) {
196
- return ctx.typeDecls.find(d => d.name === name);
248
+ const direct = ctx.typeDecls.find(d => d.name === name);
249
+ if (direct)
250
+ return direct;
251
+ // Dotted names (e.g. `Agent.Info`, `Permission.Ruleset`): fall back to the
252
+ // last segment, so `//@ declare-type Info { ... }` matches a reference to
253
+ // `Agent.Info` without forcing the user to repeat the namespace.
254
+ const dotIdx = name.lastIndexOf(".");
255
+ if (dotIdx >= 0)
256
+ return ctx.typeDecls.find(d => d.name === name.slice(dotIdx + 1));
257
+ return undefined;
258
+ }
259
+ /** Expand alias-kind typeDecls when the alias target is structural (array,
260
+ * map, set, optional, or another user type). Primitive-typed aliases like
261
+ * `type TaskId = number` stay as `user("TaskId")` so the generated Dafny
262
+ * preserves the alias name. Recursive through compound types; cycle-safe. */
263
+ function expandAlias(ty, typeDecls, seen = new Set()) {
264
+ if (ty.kind === "user") {
265
+ if (seen.has(ty.name))
266
+ return ty;
267
+ let decl = typeDecls.find(d => d.name === ty.name);
268
+ if (!decl && ty.name.includes(".")) {
269
+ const tail = ty.name.slice(ty.name.lastIndexOf(".") + 1);
270
+ decl = typeDecls.find(d => d.name === tail);
271
+ }
272
+ if (decl?.kind === "alias" && decl.aliasOfTy) {
273
+ const target = decl.aliasOfTy;
274
+ if (target.kind === "array" || target.kind === "map" || target.kind === "set" || target.kind === "optional" || target.kind === "user") {
275
+ return expandAlias(target, typeDecls, new Set([...seen, ty.name]));
276
+ }
277
+ }
278
+ return ty;
279
+ }
280
+ if (ty.kind === "optional")
281
+ return { kind: "optional", inner: expandAlias(ty.inner, typeDecls, seen) };
282
+ if (ty.kind === "array")
283
+ return { kind: "array", elem: expandAlias(ty.elem, typeDecls, seen) };
284
+ if (ty.kind === "set")
285
+ return { kind: "set", elem: expandAlias(ty.elem, typeDecls, seen) };
286
+ if (ty.kind === "map")
287
+ return { kind: "map", key: expandAlias(ty.key, typeDecls, seen), value: expandAlias(ty.value, typeDecls, seen) };
288
+ return ty;
197
289
  }
198
290
  function getDiscriminant(ctx, typeName) {
199
291
  return findDecl(ctx, typeName)?.discriminant;
200
292
  }
293
+ // ── Equality hazard: structural in the proof vs reference at runtime ─────────
294
+ // `===`/`!==` is modeled as Dafny structural equality, but the SAME TypeScript
295
+ // runs `===` as JS *reference* equality on objects/arrays. The two agree only
296
+ // when the operand is a primitive at runtime: number / string / bool, or a
297
+ // string-union enum (which runs as a plain string). Records, discriminated
298
+ // unions, arrays, maps, sets, and unresolved generics are reference-compared at
299
+ // runtime, so a structural proof over them is unsound. Returns true for those.
300
+ function refEqHazard(ty, typeDecls) {
301
+ if (ty.kind === "array" || ty.kind === "map" || ty.kind === "set")
302
+ return true;
303
+ if (ty.kind === "user") {
304
+ let decl = typeDecls.find(d => d.name === ty.name);
305
+ if (!decl && ty.name.includes(".")) {
306
+ const tail = ty.name.slice(ty.name.lastIndexOf(".") + 1);
307
+ decl = typeDecls.find(d => d.name === tail);
308
+ }
309
+ if (!decl)
310
+ return true; // generic type parameter / unknown → assume reference
311
+ if (decl.kind === "string-union")
312
+ return false; // runs as a JS string → `===` is structural
313
+ if (decl.kind === "alias")
314
+ return decl.aliasOfTy ? refEqHazard(decl.aliasOfTy, typeDecls) : false;
315
+ return true; // record / discriminated-union → reference at runtime
316
+ }
317
+ return false; // primitives, optional, unknown, fn, void
318
+ }
319
+ const _warnedRefEq = new Set();
320
+ function warnRefEq(op, l, r) {
321
+ const label = (t) => t.kind === "user" ? t.name : t.kind;
322
+ const msg = `'${op}' compares non-primitive operands (${label(l)} ${op} ${label(r)}): structural equality in the proof, but reference equality when this TypeScript runs. Sound only if operands are primitives or a canonical (string/number) encoding; otherwise compare via an explicit structural equals.`;
323
+ if (_warnedRefEq.has(msg))
324
+ return;
325
+ _warnedRefEq.add(msg);
326
+ console.error(`WARNING: ${msg}`);
327
+ }
328
+ /** A type ts-morph handed us that LemmaScript hasn't modeled: contains
329
+ * `unknown` (TS `any`), or a `user` type whose name isn't a known declaration
330
+ * (an opaque expanded union like `"AssistantMsg | ToolMsg"` that ts-morph
331
+ * produced by expanding an alias LS shadows via declare-type). Used by
332
+ * `case "let"` to decide when LS's own `init.ty` is the better source of
333
+ * structure. */
334
+ function isUnmodeledTy(ty, typeDecls) {
335
+ if (ty.kind === "unknown")
336
+ return true;
337
+ if (ty.kind === "optional")
338
+ return isUnmodeledTy(ty.inner, typeDecls);
339
+ if (ty.kind === "array")
340
+ return isUnmodeledTy(ty.elem, typeDecls);
341
+ if (ty.kind === "set")
342
+ return isUnmodeledTy(ty.elem, typeDecls);
343
+ if (ty.kind === "map")
344
+ return isUnmodeledTy(ty.key, typeDecls) || isUnmodeledTy(ty.value, typeDecls);
345
+ if (ty.kind === "user") {
346
+ const base = ty.name.includes("<") ? ty.name.slice(0, ty.name.indexOf("<")) : ty.name;
347
+ return !typeDecls.some(d => d.name === base);
348
+ }
349
+ return false;
350
+ }
201
351
  /** Infer quantifier variable type from usage in body.
202
352
  * If the variable is used as a map/set key (e.g. map.has(k), map.get(k)),
203
353
  * return the collection's key type. Otherwise return null (default to int). */
@@ -264,8 +414,16 @@ function inferQuantVarType(varName, body, ctx) {
264
414
  function classifyCall(fn, ctx) {
265
415
  if (fn.kind === "field" && fn.obj.kind === "var" && fn.obj.name === "Math")
266
416
  return "pure";
417
+ if (fn.kind === "field" && fn.obj.kind === "var" && fn.obj.name === "Array" && fn.field === "isArray")
418
+ return "pure";
267
419
  if (fn.kind === "var" && (ctx.inSpec || ctx.inLambda) && ctx.pureFns.has(fn.name))
268
420
  return "spec-pure";
421
+ // Bare-name `//@ extern` declarations are emitted as `function {:axiom}` —
422
+ // pure from the verifier's perspective. Classify them as pure so callers
423
+ // don't get lifted to statement-level binds (which would force lambdas to
424
+ // become multi-statement, illegal in Dafny).
425
+ if (fn.kind === "var" && ctx.externs.has(fn.name))
426
+ return "pure";
269
427
  if (fn.kind === "var" && ctx.inSpec) {
270
428
  // Not a known pure function — could be external (Lean-defined spec helper).
271
429
  // Pass through as "pure" and let Lean catch any errors.
@@ -276,24 +434,56 @@ function classifyCall(fn, ctx) {
276
434
  return "unknown";
277
435
  }
278
436
  // ── Call resolution helpers ─────────────────────────────────
279
- /** Infer lambda param types from array method context (map, filter, etc.).
280
- * Returns updated rawArgs with inferred tsType on the first lambda param. */
281
- function inferLambdaParamTypes(fn, rawArgs) {
437
+ /** Infer lambda param types from array method context (map, filter, etc.)
438
+ * AND from function-typed parameters of named callees (e.g., a `Comparator =
439
+ * (a, b) => bool` parameter propagates `string, string` to the lambda's
440
+ * inline params). Returns updated rawArgs with inferred tsType. */
441
+ function tyToTsStr(ty) {
442
+ if (ty.kind === "user")
443
+ return ty.name;
444
+ if (ty.kind === "string")
445
+ return "string";
446
+ if (ty.kind === "int" || ty.kind === "nat")
447
+ return "number";
448
+ if (ty.kind === "bool")
449
+ return "boolean";
450
+ return undefined;
451
+ }
452
+ function inferLambdaParamTypes(fn, rawArgs, ctx) {
282
453
  if (fn.kind === "field" && fn.obj.ty.kind === "array" &&
283
- ["map", "filter", "every", "some", "find"].includes(fn.field) &&
454
+ ["map", "filter", "every", "some", "find", "findLast", "findIndex"].includes(fn.field) &&
284
455
  rawArgs.length >= 1 && rawArgs[0].kind === "lambda" &&
285
456
  rawArgs[0].params.length >= 1 && !rawArgs[0].params[0].tsType) {
286
457
  const elemTy = fn.obj.ty.elem;
287
- const tsType = elemTy.kind === "user" ? elemTy.name
288
- : elemTy.kind === "string" ? "string"
289
- : elemTy.kind === "int" || elemTy.kind === "nat" ? "number"
290
- : elemTy.kind === "bool" ? "boolean" : undefined;
458
+ const tsType = tyToTsStr(elemTy);
291
459
  if (tsType) {
292
460
  const lam = rawArgs[0];
293
461
  const updatedParams = [{ ...lam.params[0], tsType }, ...lam.params.slice(1)];
294
462
  return [{ ...lam, params: updatedParams }, ...rawArgs.slice(1)];
295
463
  }
296
464
  }
465
+ // Named-callee propagation: when an argument position expects a function
466
+ // type, infer the lambda's param types from that function type. Aliases
467
+ // (e.g., `Comparator`) are expanded via the typeDecls.
468
+ if (fn.kind === "var" && ctx?.fnParams.has(fn.name)) {
469
+ const paramTys = ctx.fnParams.get(fn.name);
470
+ return rawArgs.map((a, i) => {
471
+ if (a.kind !== "lambda" || i >= paramTys.length)
472
+ return a;
473
+ let pTy = paramTys[i];
474
+ if (pTy.kind === "user") {
475
+ const decl = ctx.typeDecls.find(d => d.name === pTy.name);
476
+ if (decl?.kind === "alias" && decl.aliasOfTy)
477
+ pTy = decl.aliasOfTy;
478
+ else if (decl?.kind === "alias" && decl.aliasOf)
479
+ pTy = parseTsType(decl.aliasOf);
480
+ }
481
+ if (pTy.kind !== "fn")
482
+ return a;
483
+ const updatedParams = a.params.map((p, idx) => p.tsType || idx >= pTy.params.length ? p : { ...p, tsType: tyToTsStr(pTy.params[idx]) });
484
+ return { ...a, params: updatedParams };
485
+ });
486
+ }
297
487
  return rawArgs;
298
488
  }
299
489
  /** Coerce call arguments: string literals → user types, non-optional → Some, pad missing optional args. */
@@ -322,6 +512,11 @@ function coerceCallArgs(args, fn, ctx) {
322
512
  function inferMethodReturnTy(fn, args, ctx) {
323
513
  if (fn.kind !== "field")
324
514
  return { kind: "unknown" };
515
+ // `Array.isArray(x)` always returns boolean. narrow.ts recognizes this call as
516
+ // a discriminator predicate when `x` has type of a synthesized array-union.
517
+ if (fn.obj.kind === "var" && fn.obj.name === "Array" && fn.field === "isArray") {
518
+ return { kind: "bool" };
519
+ }
325
520
  const objTy = fn.obj.ty;
326
521
  if (objTy.kind === "map") {
327
522
  if (fn.field === "get")
@@ -344,12 +539,24 @@ function inferMethodReturnTy(fn, args, ctx) {
344
539
  return { kind: "int" };
345
540
  if (fn.field === "shift")
346
541
  return objTy.elem;
542
+ if (fn.field === "pop")
543
+ return { kind: "optional", inner: objTy.elem };
347
544
  if (fn.field === "push" || fn.field === "concat")
348
545
  return objTy;
349
546
  if (fn.field === "filter")
350
547
  return objTy;
351
548
  if (fn.field === "every" || fn.field === "some")
352
549
  return { kind: "bool" };
550
+ if (fn.field === "find" || fn.field === "findLast")
551
+ return { kind: "optional", inner: objTy.elem };
552
+ if (fn.field === "findIndex")
553
+ return { kind: "int" };
554
+ if (fn.field === "flat" && objTy.elem.kind === "array")
555
+ return { kind: "array", elem: objTy.elem.elem };
556
+ if (fn.field === "slice")
557
+ return objTy;
558
+ if (fn.field === "join" && objTy.elem.kind === "string")
559
+ return { kind: "string" };
353
560
  if (fn.field === "map" && args.length >= 1 && args[0].kind === "lambda") {
354
561
  const retTy = args[0].body.length > 0 && args[0].body[0].kind === "return"
355
562
  ? args[0].body[0].value.ty : { kind: "unknown" };
@@ -357,9 +564,13 @@ function inferMethodReturnTy(fn, args, ctx) {
357
564
  }
358
565
  }
359
566
  else if (objTy.kind === "string") {
360
- if (fn.field === "trim" || fn.field === "toLowerCase" || fn.field === "toUpperCase")
567
+ if (fn.field === "trim" || fn.field === "trimEnd" || fn.field === "trimStart" || fn.field === "toLowerCase" || fn.field === "toUpperCase")
361
568
  return { kind: "string" };
362
- if (fn.field === "includes")
569
+ if (fn.field === "slice" || fn.field === "substring")
570
+ return { kind: "string" };
571
+ if (fn.field === "split")
572
+ return { kind: "array", elem: { kind: "string" } };
573
+ if (fn.field === "includes" || fn.field === "startsWith" || fn.field === "endsWith")
363
574
  return { kind: "bool" };
364
575
  }
365
576
  return { kind: "unknown" };
@@ -430,6 +641,11 @@ function resolveExpr(e, ctx) {
430
641
  if (e.op === "===" || e.op === "!==") {
431
642
  left = coerceStr(left, right.ty);
432
643
  right = coerceStr(right, left.ty);
644
+ // Spec (`//@`) comparisons are proof-only, so they can't diverge at
645
+ // runtime; only warn on executable code.
646
+ if (!ctx.inSpec && refEqHazard(left.ty, ctx.typeDecls) && refEqHazard(right.ty, ctx.typeDecls)) {
647
+ warnRefEq(e.op, left.ty, right.ty);
648
+ }
433
649
  }
434
650
  let ty = { kind: "unknown" };
435
651
  if (["===", "!==", ">=", "<=", ">", "<", "in"].includes(e.op))
@@ -452,8 +668,21 @@ function resolveExpr(e, ctx) {
452
668
  return { kind: "unop", op: e.op, expr, ty: e.op === "!" ? { kind: "bool" } : expr.ty };
453
669
  }
454
670
  case "call": {
671
+ // Extern dispatch: `NS.method(args)` where NS.method is declared via
672
+ // `//@ extern`. Rewrite into a flat-name call (`NS_method(args)`) so the
673
+ // rest of the pipeline sees an ordinary pure function. The extern's
674
+ // declaration is emitted alongside the file as `function {:axiom} ...`.
675
+ if (e.fn.kind === "field" && e.fn.obj.kind === "var") {
676
+ const qualified = `${e.fn.obj.name}.${e.fn.field}`;
677
+ const ext = ctx.externs.get(qualified);
678
+ if (ext) {
679
+ const args = e.args.map(a => resolveExpr(a, ctx));
680
+ const fn = { kind: "var", name: ext.flat, ty: { kind: "unknown" } };
681
+ return { kind: "call", fn, args, ty: ext.returnTy, callKind: "pure" };
682
+ }
683
+ }
455
684
  const fn = resolveExpr(e.fn, ctx);
456
- const rawArgs = inferLambdaParamTypes(fn, e.args);
685
+ const rawArgs = inferLambdaParamTypes(fn, e.args, ctx);
457
686
  // For .push() on a typed array, resolve args with element type context
458
687
  let argCtx = ctx;
459
688
  if (fn.kind === "field" && fn.obj.ty.kind === "array" && fn.field === "push" &&
@@ -463,13 +692,20 @@ function resolveExpr(e, ctx) {
463
692
  // Propagate parameter types to arguments for record literal resolution
464
693
  // (enables inline discriminated union construction in function arguments)
465
694
  const paramTypes = fn.kind === "var" && ctx.fnParams.has(fn.name) ? ctx.fnParams.get(fn.name) : null;
466
- const args = coerceCallArgs(rawArgs.map((a, i) => {
695
+ let args = coerceCallArgs(rawArgs.map((a, i) => {
467
696
  let aCtx = argCtx;
468
697
  if (paramTypes && i < paramTypes.length && paramTypes[i].kind === "user") {
469
698
  aCtx = { ...aCtx, returnTy: paramTypes[i] };
470
699
  }
471
700
  return resolveExpr(a, aCtx);
472
701
  }), fn, ctx);
702
+ // Array method `.with(i, v)`: coerce the value arg to the element type
703
+ // so `arr[i] = v` on `(T|null)[]` wraps `T` → `Some(T)` (and similarly
704
+ // for synth array-unions). Same shape as the record-field coercion
705
+ // below: assigning a narrower value into a wider slot.
706
+ if (fn.kind === "field" && fn.field === "with" && fn.obj.ty.kind === "array" && args.length === 2) {
707
+ args = [args[0], coerceToTargetTy(args[1], fn.obj.ty.elem, ctx.typeDecls)];
708
+ }
473
709
  let ty = inferMethodReturnTy(fn, args, ctx);
474
710
  // For same-file function calls, use the known return type
475
711
  if (ty.kind === "unknown" && fn.kind === "var" && ctx.fnReturns.has(fn.name)) {
@@ -553,22 +789,29 @@ function resolveExpr(e, ctx) {
553
789
  }
554
790
  else {
555
791
  // call: prev step yielded a callable (typically a method via field).
556
- // Build a fake fn TExpr from prev steps to reuse inferMethodReturnTy.
557
- const args = step.args.map(a => resolveExpr(a, ctx));
792
+ // Build a fake fn TExpr from prev steps to reuse inferMethodReturnTy
793
+ // and inferLambdaParamTypes without the latter, a lambda arg buried
794
+ // inside `obj?.filter(r => ...)` gets `int`-typed params instead of
795
+ // the array's element type.
558
796
  const lastField = chain.length > 0 && chain[chain.length - 1].kind === "field"
559
797
  ? chain[chain.length - 1] : null;
560
798
  let callTy = { kind: "unknown" };
561
799
  let callKind = "unknown";
800
+ let rawArgs = step.args;
562
801
  if (lastField) {
563
- // Build a synthetic field TExpr with the prior step's input type as obj
564
- // so inferMethodReturnTy can dispatch on the receiver type.
565
802
  const priorInTy = chain.length >= 2 ? chain[chain.length - 2].ty
566
803
  : (obj.ty.kind === "optional" ? obj.ty.inner : obj.ty);
567
804
  const fakeObj = { kind: "var", name: "_chain_recv", ty: priorInTy };
568
805
  const fakeFn = { kind: "field", obj: fakeObj, field: lastField.name, ty: lastField.ty };
806
+ rawArgs = inferLambdaParamTypes(fakeFn, rawArgs);
807
+ const args = rawArgs.map(a => resolveExpr(a, ctx));
569
808
  callTy = inferMethodReturnTy(fakeFn, args, ctx);
570
809
  callKind = "method";
810
+ chain.push({ kind: "call", args, ty: callTy, callKind });
811
+ stepInTy = callTy;
812
+ continue;
571
813
  }
814
+ const args = rawArgs.map(a => resolveExpr(a, ctx));
572
815
  chain.push({ kind: "call", args, ty: callTy, callKind });
573
816
  stepInTy = callTy;
574
817
  }
@@ -580,8 +823,19 @@ function resolveExpr(e, ctx) {
580
823
  case "record": {
581
824
  const spread = e.spread ? resolveExpr(e.spread, ctx) : null;
582
825
  const ty = spread ? spread.ty : { kind: "unknown" };
583
- // Infer record type: from spread, or from return type context
584
- const recordTy = ty.kind === "user" ? ty : ctx.returnTy.kind === "user" ? ctx.returnTy : null;
826
+ // Record literal in map-typed context (e.g. `const M: Record<string, V> = {a: ...}`):
827
+ // attach the map type so transform/emit can produce a map literal.
828
+ if (!spread && ctx.returnTy.kind === "map") {
829
+ const mapTy = ctx.returnTy;
830
+ const fieldCtx = { ...ctx, returnTy: mapTy.value };
831
+ const fields = e.fields.map(f => ({ name: f.name, value: resolveExpr(f.value, fieldCtx) }));
832
+ return { kind: "record", spread: null, fields, ty: mapTy };
833
+ }
834
+ // Infer record type: from spread, or from return type context. Unwrap
835
+ // an outer Optional when looking at returnTy — `return {...} : null`
836
+ // has ctx.returnTy = Option<T>, but the record literal's natural type is T.
837
+ const returnTyUnwrapped = ctx.returnTy.kind === "optional" ? ctx.returnTy.inner : ctx.returnTy;
838
+ const recordTy = ty.kind === "user" ? ty : returnTyUnwrapped.kind === "user" ? returnTyUnwrapped : null;
585
839
  const decl = recordTy ? ctx.typeDecls.find(d => d.name === recordTy.name && d.kind === "record") : undefined;
586
840
  // Clear returnTy for field values — it applies to THIS record, not nested ones
587
841
  const fieldCtx = recordTy ? { ...ctx, returnTy: { kind: "unknown" } } : ctx;
@@ -600,19 +854,22 @@ function resolveExpr(e, ctx) {
600
854
  if (value.kind === "record" && value.fields.length === 0 && !value.spread && declTy.kind === "map") {
601
855
  value = { kind: "arrayLiteral", elems: [], ty: declTy };
602
856
  }
603
- // Coerce non-optional to optional: wrap in Some (only when value type is concrete)
604
- if (declTy.kind === "optional" && value.ty.kind !== "optional" && value.ty.kind !== "void" && value.ty.kind !== "unknown") {
605
- value = wrapSome(value, declTy);
606
- }
857
+ // Assignment-position upcasts: T Option<T>, T[] ArrayBranch(T[]),
858
+ // U → NonArrayBranch(U). Handles both optional fields and fields
859
+ // typed as a synth array-union (`T[] | U`).
860
+ value = coerceToTargetTy(value, declTy, ctx.typeDecls);
607
861
  }
608
862
  return { name: f.name, value };
609
863
  });
610
864
  return { kind: "record", spread, fields, ty: recordTy ?? ty };
611
865
  }
612
866
  case "result":
867
+ // \result desugars to a regular var so all the variable-narrowing
868
+ // machinery (env lookup, optional checks, path matching) just works.
869
+ // The env in ensuresCtx is pre-seeded with "\result" → returnTy.
613
870
  if (!ctx.allowResult)
614
871
  throw new Error("\\result is only valid in ensures");
615
- return { kind: "result", ty: ctx.returnTy };
872
+ return { kind: "var", name: "\\result", ty: lookup(ctx.env, "\\result") ?? ctx.returnTy };
616
873
  case "forall": {
617
874
  const varTy = e.varType !== "int" ? parseTsType(e.varType)
618
875
  : inferQuantVarType(e.var, e.body, ctx) ?? { kind: "int" };
@@ -683,6 +940,14 @@ function resolveExpr(e, ctx) {
683
940
  else if (else_.ty.kind === "void" && then_.ty.kind !== "void" && then_.ty.kind !== "unknown") {
684
941
  ty = { kind: "optional", inner: then_.ty };
685
942
  }
943
+ else if (then_.ty.kind === "optional" && else_.ty.kind !== "optional" && else_.ty.kind !== "unknown") {
944
+ // Asymmetric optional: one branch returns Option<T>, the other returns T.
945
+ // Widen to Option<T> so callers/return-coercion see the wider type.
946
+ ty = then_.ty;
947
+ }
948
+ else if (else_.ty.kind === "optional" && then_.ty.kind !== "optional" && then_.ty.kind !== "unknown") {
949
+ ty = else_.ty;
950
+ }
686
951
  return { kind: "conditional", cond, then: then_, else: else_, ty };
687
952
  }
688
953
  case "emptyCollection": {
@@ -752,20 +1017,48 @@ function resolveBlock(stmts, ctx) {
752
1017
  function resolveStmt(s, ctx) {
753
1018
  switch (s.kind) {
754
1019
  case "let": {
755
- const declTy = resolveTsType(s.tsType, ctx.overrides, s.name);
1020
+ // No source annotation → infer type from initializer (resolved first).
1021
+ if (s.tsType === null) {
1022
+ const init = resolveExpr(s.init, ctx);
1023
+ const ty = init.ty;
1024
+ const mutable = s.mutable || isRefMutableInTS(ty);
1025
+ return [{ kind: "let", name: s.name, ty, mutable, init }, extend(ctx.env, s.name, ty)];
1026
+ }
1027
+ // expandAlias unwraps an array/collection alias (`type Board = number[]`)
1028
+ // to its underlying type, so array methods / index-assignment on the
1029
+ // local dispatch correctly (params get the same treatment, see makeParams).
1030
+ const declTy = expandAlias(resolveTsType(s.tsType, ctx.overrides, s.name), ctx.typeDecls);
756
1031
  // Propagate declared type as returnTy so nested record expressions
757
1032
  // resolve union variants correctly (e.g., EffectState → mode: EffectMode → { kind: 'Idle' })
758
1033
  const initCtx = declTy.kind === "user" ? { ...ctx, returnTy: declTy } : ctx;
759
1034
  const init = coerceStr(resolveExpr(s.init, initCtx), declTy);
760
- // Map indexing: TS says T, but access can fail → use Optional<T> from init
761
- const ty = (declTy.kind !== "optional" && init.ty.kind === "optional") ? init.ty : declTy;
1035
+ let ty;
1036
+ if (isUnmodeledTy(declTy, ctx.typeDecls) && !isUnmodeledTy(init.ty, ctx.typeDecls)) {
1037
+ // ts-morph's declared type is opaque to us (an expanded union it made
1038
+ // by inlining an alias we shadow via declare-type, or any-laden), but
1039
+ // LS resolved the initializer to something concrete. Take the structure
1040
+ // from `init.ty`, keeping only the optionality ts-morph reported.
1041
+ ty = declTy.kind === "optional" && init.ty.kind !== "optional"
1042
+ ? { kind: "optional", inner: init.ty }
1043
+ : init.ty;
1044
+ }
1045
+ else {
1046
+ // Map indexing: TS says T, but access can fail → use Optional<T> from init
1047
+ ty = (declTy.kind !== "optional" && init.ty.kind === "optional") ? init.ty : declTy;
1048
+ }
762
1049
  // const collections are mutable in value-semantics world (TS mutates in place, Dafny/Lean reassign)
763
1050
  const mutable = s.mutable || isRefMutableInTS(ty);
764
1051
  return [{ kind: "let", name: s.name, ty, mutable, init }, extend(ctx.env, s.name, ty)];
765
1052
  }
766
1053
  case "assign": {
767
1054
  const targetTy = lookup(ctx.env, s.target) ?? { kind: "unknown" };
768
- return [{ kind: "assign", target: s.target, value: coerceStr(resolveExpr(s.value, ctx), targetTy) }, ctx.env];
1055
+ let value = coerceStr(resolveExpr(s.value, ctx), targetTy);
1056
+ // Auto-wrap non-optional value in Some when target is optional
1057
+ const isUndef = value.kind === "var" && value.name === "undefined";
1058
+ if (targetTy.kind === "optional" && value.ty.kind !== "optional" && value.ty.kind !== "unknown" && !isUndef) {
1059
+ value = wrapSome(value, targetTy);
1060
+ }
1061
+ return [{ kind: "assign", target: s.target, value }, ctx.env];
769
1062
  }
770
1063
  case "return": {
771
1064
  let value = coerceStr(resolveExpr(s.value, ctx), ctx.returnTy);
@@ -877,7 +1170,7 @@ function resolveStmt(s, ctx) {
877
1170
  case "assert": {
878
1171
  const specCtx = { ...ctx, inSpec: true };
879
1172
  const expr = resolveExpr(parseExpr(s.expr), specCtx);
880
- return [{ kind: "assert", expr }, ctx.env];
1173
+ return [{ kind: "assert", expr, assumed: s.assumed }, ctx.env];
881
1174
  }
882
1175
  }
883
1176
  }
@@ -1052,18 +1345,22 @@ function containsReturn(stmts) {
1052
1345
  return false;
1053
1346
  }
1054
1347
  // ── Resolve function / module ────────────────────────────────
1055
- function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), opts) {
1348
+ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), externs = new Map(), moduleConstants = new Map(), opts) {
1056
1349
  const overrides = new Map(fn.typeAnnotations.map(a => [a.name, a.type]));
1057
- const params = fn.params.map(p => ({ name: p.name, ty: resolveTsType(p.tsType, overrides, p.name) }));
1058
- const returnTy = resolveTsType(fn.returnType, overrides, "\\result");
1350
+ const params = fn.params.map(p => ({ name: p.name, ty: expandAlias(resolveTsType(p.tsType, overrides, p.name), typeDecls) }));
1351
+ const returnTy = expandAlias(resolveTsType(fn.returnType, overrides, "\\result"), typeDecls);
1059
1352
  let env = null;
1353
+ // Module-level constants are in scope for every function body. Added before
1354
+ // params so a param named the same as a const would shadow it (param wins).
1355
+ for (const [name, ty] of moduleConstants)
1356
+ env = extend(env, name, ty);
1060
1357
  if (opts?.thisBinding)
1061
1358
  env = extend(env, opts.thisBinding.name, opts.thisBinding.ty);
1062
1359
  for (const p of params)
1063
1360
  env = extend(env, p.name, p.ty);
1064
- const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
1361
+ const baseCtx = { env, typeDecls, overrides, allowResult: false, returnTy, pureFns, fnParams, fnReturns, externs, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
1065
1362
  const requiresCtx = { ...baseCtx, inSpec: true };
1066
- const ensuresCtx = { ...baseCtx, allowResult: true, inSpec: true };
1363
+ const ensuresCtx = { ...baseCtx, env: extend(env, "\\result", returnTy), allowResult: true, inSpec: true };
1067
1364
  // Apply type parameter constraints from //@ type T (==) annotations
1068
1365
  const typeParams = fn.typeParams.map(tp => {
1069
1366
  const constraint = overrides.get(tp);
@@ -1087,13 +1384,13 @@ function resolveFunction(fn, typeDecls, pureFns, fnParams = new Map(), fnReturns
1087
1384
  body: resolveBlock(fn.body, bodyCtx),
1088
1385
  };
1089
1386
  }
1090
- function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map()) {
1387
+ function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns = new Map(), externs = new Map(), moduleConstants = new Map()) {
1091
1388
  const fields = cls.fields.map(f => ({ name: f.name, ty: parseTsType(f.tsType) }));
1092
1389
  // Create a synthetic record type for 'this' so field access resolves
1093
1390
  const thisType = { kind: "user", name: cls.name };
1094
1391
  const thisDecl = { name: cls.name, kind: "record", fields: cls.fields.map(f => ({ name: f.name, tsType: f.tsType, type: parseTsType(f.tsType) })) };
1095
1392
  const allTypeDecls = [...typeDecls, thisDecl];
1096
- const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, fnReturns, {
1393
+ const methods = cls.methods.map(fn => resolveFunction(fn, allTypeDecls, pureFns, fnParams, fnReturns, externs, moduleConstants, {
1097
1394
  thisBinding: { name: "this", ty: thisType },
1098
1395
  forcePure: false, // class methods are never pure (they access this)
1099
1396
  }));
@@ -1102,6 +1399,22 @@ function resolveClass(cls, typeDecls, pureFns, fnParams = new Map(), fnReturns =
1102
1399
  /** Pre-compute Ty on all TypeDeclInfo fields/variants/aliases.
1103
1400
  * Called once per module so consumers can read field.type instead of re-parsing tsType. */
1104
1401
  function precomputeFieldTypes(typeDecls) {
1402
+ precomputeFieldTypesInner(typeDecls);
1403
+ // Expand alias references inside record/variant field types so downstream
1404
+ // code doesn't have to follow `user("Ruleset")` indirection at every lookup.
1405
+ for (const d of typeDecls) {
1406
+ if (d.fields)
1407
+ for (const f of d.fields)
1408
+ if (f.type)
1409
+ f.type = expandAlias(f.type, typeDecls);
1410
+ if (d.variants)
1411
+ for (const v of d.variants)
1412
+ for (const f of v.fields)
1413
+ if (f.type)
1414
+ f.type = expandAlias(f.type, typeDecls);
1415
+ }
1416
+ }
1417
+ function precomputeFieldTypesInner(typeDecls) {
1105
1418
  for (const d of typeDecls) {
1106
1419
  if (d.fields)
1107
1420
  for (const f of d.fields)
@@ -1115,6 +1428,7 @@ function precomputeFieldTypes(typeDecls) {
1115
1428
  }
1116
1429
  }
1117
1430
  export function resolveModule(raw) {
1431
+ _warnedRefEq.clear();
1118
1432
  precomputeFieldTypes(raw.typeDecls);
1119
1433
  const pureFns = computePureFns(raw.functions);
1120
1434
  // Pre-compute function parameter and return types
@@ -1122,20 +1436,76 @@ export function resolveModule(raw) {
1122
1436
  const fnReturns = new Map();
1123
1437
  for (const fn of raw.functions) {
1124
1438
  const overrides = new Map(fn.typeAnnotations.map(a => [a.name, a.type]));
1125
- fnParams.set(fn.name, fn.params.map(p => resolveTsType(p.tsType, overrides, p.name)));
1126
- fnReturns.set(fn.name, resolveTsType(fn.returnType, overrides, "\\result"));
1439
+ fnParams.set(fn.name, fn.params.map(p => expandAlias(resolveTsType(p.tsType, overrides, p.name), raw.typeDecls)));
1440
+ fnReturns.set(fn.name, expandAlias(resolveTsType(fn.returnType, overrides, "\\result"), raw.typeDecls));
1127
1441
  }
1128
- const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, fnReturns, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
1129
- const constants = (raw.constants ?? []).map(c => ({
1130
- name: c.name,
1131
- ty: parseTsType(c.tsType),
1132
- value: resolveExpr(c.value, emptyCtx),
1133
- }));
1442
+ // Externs: resolve param/return types once. For bare-name externs (no dot),
1443
+ // also register in fnReturns so ordinary `foo(args)` calls get the right
1444
+ // return type at resolution; dotted externs are handled in resolveExpr's
1445
+ // call case via the externs map directly.
1446
+ const externs = new Map();
1447
+ // First pass: register signatures so spec resolution (below) can reference
1448
+ // them — including the extern referring to itself, or specs that mention
1449
+ // sibling externs.
1450
+ for (const ext of raw.externs ?? []) {
1451
+ const params = ext.params.map(p => parseTsType(p.tsType));
1452
+ const returnTy = parseTsType(ext.returnType);
1453
+ externs.set(ext.qualified, { flat: ext.flat, params, returnTy });
1454
+ if (!ext.qualified.includes("."))
1455
+ fnReturns.set(ext.qualified, returnTy);
1456
+ }
1457
+ // Second pass: resolve the lifted `requires`/`ensures` strings in each
1458
+ // extern's own param scope. `\result` is in scope under `ensures`.
1459
+ const tExterns = (raw.externs ?? []).map(ext => {
1460
+ const sig = externs.get(ext.qualified);
1461
+ let env = null;
1462
+ for (let i = 0; i < ext.params.length; i++) {
1463
+ env = extend(env, ext.params[i].name, sig.params[i]);
1464
+ }
1465
+ const baseCtx = { env, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: sig.returnTy, pureFns, fnParams, fnReturns, externs, inSpec: true, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
1466
+ const ensuresCtx = { ...baseCtx, env: extend(env, "\\result", sig.returnTy), allowResult: true };
1467
+ const requires = ext.requires.map(s => {
1468
+ try {
1469
+ return resolveSpec(s, baseCtx);
1470
+ }
1471
+ catch {
1472
+ return null;
1473
+ }
1474
+ }).filter((e) => e !== null);
1475
+ const ensures = ext.ensures.map(s => {
1476
+ try {
1477
+ return resolveSpec(s, ensuresCtx);
1478
+ }
1479
+ catch {
1480
+ return null;
1481
+ }
1482
+ }).filter((e) => e !== null);
1483
+ return {
1484
+ qualified: ext.qualified,
1485
+ flat: ext.flat,
1486
+ typeParams: ext.typeParams,
1487
+ params: ext.params.map((p, i) => ({ name: p.name, ty: sig.params[i] })),
1488
+ returnTy: sig.returnTy,
1489
+ requires,
1490
+ ensures,
1491
+ };
1492
+ });
1493
+ const emptyCtx = { env: null, typeDecls: raw.typeDecls, overrides: new Map(), allowResult: false, returnTy: { kind: "int" }, pureFns, fnParams, fnReturns, externs, inSpec: false, inLambda: false, narrowedPaths: [], narrowedIndices: [] };
1494
+ const constants = (raw.constants ?? []).map(c => {
1495
+ const ty = expandAlias(parseTsType(c.tsType), raw.typeDecls);
1496
+ // Propagate the declared type into the value's resolution context so that
1497
+ // record literals on map-typed constants (e.g. `Record<string, number>`)
1498
+ // get their `ty` set to `map<...>` rather than `user("...")`.
1499
+ const valueCtx = { ...emptyCtx, returnTy: ty };
1500
+ return { name: c.name, ty, value: resolveExpr(c.value, valueCtx) };
1501
+ });
1502
+ const moduleConstants = new Map(constants.map(c => [c.name, c.ty]));
1134
1503
  return {
1135
1504
  file: raw.file,
1136
1505
  typeDecls: raw.typeDecls,
1506
+ externs: tExterns,
1137
1507
  constants,
1138
- functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams, fnReturns)),
1139
- classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams, fnReturns)),
1508
+ functions: raw.functions.map(fn => resolveFunction(fn, raw.typeDecls, pureFns, fnParams, fnReturns, externs, moduleConstants)),
1509
+ classes: (raw.classes ?? []).map(cls => resolveClass(cls, raw.typeDecls, pureFns, fnParams, fnReturns, externs, moduleConstants)),
1140
1510
  };
1141
1511
  }