From dfcb2c2d49e0ef544b01ad26e866e92eb75ed282 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Sat, 4 Jul 2026 02:23:32 +0200 Subject: [PATCH 01/12] Petrinaut HIR: compile user code to buffer-native programs, remove Babel Introduce a typed, source-spanned, JSON-serializable HIR as the single compiler for dynamics/lambda/kernel user code: - TS->HIR lowering (const destructuring, guard if/early-return, .map comprehensions, first-class distributions) with positioned diagnostics - Schema-checked typing against Color/parameter definitions; analyses for dependency sets, distribution DAGs (nodes/edges/sinks/shared draws), binding usage and constant folding - Buffer-ABI emitters: lambdas read token attributes at statically resolved offsets (tokenValues[slotBases[slot] + attr]); kernels write place-major staging floats with deferred distribution sampling ordered by output index (exact legacy RNG-stream parity); dynamics compile to allocation-free Float64Array loops - Engine hot loops (single-run + Monte Carlo) prefer buffer programs and fall back to object-convention HIR programs per item - HirArtifacts v2 compiled in the LSP worker (sdcpn/compileHirArtifacts request) and threaded through simulation/experiment configs; the engine no longer bundles any compiler (simulation workers: ~3MB -> ~40kB) - HIR semantic lints in the LSP checker (source "hir", codes 99001+); out-of-subset code is an error and Play gates on error severity only - Remove @babel/standalone and compile-user-code.ts; fix H-6519 violations in the supply-chain example; shipped examples are the buffer-ABI coverage gate Design docs: src/hir/README.md (architecture) and src/hir/dsl-sketch.md (planned OCaml-flavoured DSL lowering to the same HIR). Co-Authored-By: Claude Fable 5 --- libs/@hashintel/petrinaut-core/package.json | 2 - .../examples/supply-chain-with-disruption.ts | 7 +- .../petrinaut-core/src/hir-runtime.ts | 24 + libs/@hashintel/petrinaut-core/src/hir.ts | 86 ++ .../petrinaut-core/src/hir/README.md | 165 +++ .../petrinaut-core/src/hir/analyze.test.ts | 223 +++ .../petrinaut-core/src/hir/analyze.ts | 742 ++++++++++ .../petrinaut-core/src/hir/compile.test.ts | 74 + .../petrinaut-core/src/hir/compile.ts | 278 ++++ .../petrinaut-core/src/hir/dsl-sketch.md | 231 +++ .../src/hir/emit-buffer-js.test.ts | 264 ++++ .../petrinaut-core/src/hir/emit-buffer-js.ts | 585 ++++++++ .../petrinaut-core/src/hir/emit-js.test.ts | 194 +++ .../petrinaut-core/src/hir/emit-js.ts | 475 +++++++ libs/@hashintel/petrinaut-core/src/hir/hir.ts | 385 +++++ .../petrinaut-core/src/hir/instantiate.ts | 189 +++ .../petrinaut-core/src/hir/lint.test.ts | 319 +++++ .../@hashintel/petrinaut-core/src/hir/lint.ts | 178 +++ .../src/hir/lower-typescript.test.ts | 530 +++++++ .../src/hir/lower-typescript.ts | 1253 +++++++++++++++++ .../petrinaut-core/src/hir/surface-context.ts | 339 +++++ .../petrinaut-core/src/hir/typecheck.ts | 797 +++++++++++ libs/@hashintel/petrinaut-core/src/index.ts | 9 + .../petrinaut-core/src/lsp/language-client.ts | 35 +- .../petrinaut-core/src/lsp/lib/check-hir.ts | 128 ++ .../src/lsp/lib/checker.test.ts | 149 +- .../petrinaut-core/src/lsp/lib/checker.ts | 86 +- .../petrinaut-core/src/lsp/lib/ts-to-lsp.ts | 2 +- .../src/lsp/worker/language-server.worker.ts | 10 + .../petrinaut-core/src/lsp/worker/protocol.ts | 9 + .../src/simulation/ARCHITECTURE.md | 12 +- .../petrinaut-core/src/simulation/api.ts | 8 + .../user-code/compile-user-code.test.ts | 368 ----- .../authoring/user-code/compile-user-code.ts | 121 -- .../authoring/user-code/distribution.ts | 27 +- .../simulation/engine/buffer-transition.ts | 131 ++ .../engine/build-simulation-hir.test.ts | 208 +++ .../engine/build-simulation.test.ts | 28 +- .../src/simulation/engine/build-simulation.ts | 279 +++- .../engine/compute-next-frame.test.ts | 17 +- .../compute-possible-transition.test.ts | 1 + .../engine/execute-transitions.test.ts | 1 + .../src/simulation/engine/types.ts | 34 + .../monte-carlo/monte-carlo-simulator.test.ts | 17 +- .../src/simulation/monte-carlo/run-state.ts | 1 + .../monte-carlo/runtime/experiment.ts | 6 + .../src/simulation/monte-carlo/types.ts | 4 + .../simulation/monte-carlo/worker/messages.ts | 4 + .../monte-carlo/worker/monte-carlo.worker.ts | 1 + .../src/simulation/runtime/simulation.ts | 1 + .../src/simulation/worker/messages.ts | 4 + .../simulation/worker/simulation.worker.ts | 1 + .../petrinaut/docs/petri-net-extensions.md | 17 +- libs/@hashintel/petrinaut/docs/simulation.md | 2 +- .../src/react/experiments/provider.tsx | 31 +- .../petrinaut/src/react/lsp/context.ts | 24 + .../petrinaut/src/react/lsp/provider.tsx | 11 +- .../src/react/simulation/provider.tsx | 9 + .../components/BottomBar/bottom-bar.tsx | 6 +- .../BottomBar/diagnostics-indicator.tsx | 20 +- 60 files changed, 8528 insertions(+), 634 deletions(-) create mode 100644 libs/@hashintel/petrinaut-core/src/hir-runtime.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/README.md create mode 100644 libs/@hashintel/petrinaut-core/src/hir/analyze.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/analyze.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/compile.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/compile.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/dsl-sketch.md create mode 100644 libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/emit-js.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/emit-js.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/hir.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/instantiate.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/lint.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/lint.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/lower-typescript.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/lower-typescript.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/surface-context.ts create mode 100644 libs/@hashintel/petrinaut-core/src/hir/typecheck.ts create mode 100644 libs/@hashintel/petrinaut-core/src/lsp/lib/check-hir.ts delete mode 100644 libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.test.ts delete mode 100644 libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.ts create mode 100644 libs/@hashintel/petrinaut-core/src/simulation/engine/buffer-transition.ts create mode 100644 libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation-hir.test.ts diff --git a/libs/@hashintel/petrinaut-core/package.json b/libs/@hashintel/petrinaut-core/package.json index d96ad8faef2..352b36f113f 100644 --- a/libs/@hashintel/petrinaut-core/package.json +++ b/libs/@hashintel/petrinaut-core/package.json @@ -60,7 +60,6 @@ "test:unit": "vitest" }, "dependencies": { - "@babel/standalone": "7.28.5", "elkjs": "0.11.0", "immer": "10.1.3", "uuid": "14.0.0", @@ -68,7 +67,6 @@ "zod": "4.4.3" }, "devDependencies": { - "@types/babel__standalone": "7.1.9", "@types/node": "22.18.13", "@typescript/native-preview": "7.0.0-dev.20260511.1", "oxlint": "1.63.0", diff --git a/libs/@hashintel/petrinaut-core/src/examples/supply-chain-with-disruption.ts b/libs/@hashintel/petrinaut-core/src/examples/supply-chain-with-disruption.ts index 178b81fdcd7..25f214a80e9 100644 --- a/libs/@hashintel/petrinaut-core/src/examples/supply-chain-with-disruption.ts +++ b/libs/@hashintel/petrinaut-core/src/examples/supply-chain-with-disruption.ts @@ -1127,7 +1127,7 @@ export default TransitionKernel(() => ({}));`, { elementId: "order_priority", name: "priority", - type: "integer", + type: "real", }, { elementId: "order_promise", @@ -1193,10 +1193,9 @@ export default TransitionKernel(() => ({}));`, // reaches 0 the derivative becomes 0 (it holds, ready for arrival predicates). // All other attributes are constant during transit (derivative 0). export default Dynamics((tokens, parameters) => { - return tokens.map(({ eta, risk_score, source, cost }) => ({ + return tokens.map(({ eta }) => ({ eta: eta > 0 ? -1 : 0, risk_score: 0, - source: 0, cost: 0, })); });`, @@ -1208,7 +1207,7 @@ export default Dynamics((tokens, parameters) => { code: `// Every waiting order ages at a constant rate (age derivative = 1), driving // the backorder-conversion and cancellation hazards that depend on age. export default Dynamics((tokens, parameters) => { - return tokens.map(({ age, priority, promised_lead_time }) => ({ + return tokens.map(() => ({ age: 1, priority: 0, promised_lead_time: 0, diff --git a/libs/@hashintel/petrinaut-core/src/hir-runtime.ts b/libs/@hashintel/petrinaut-core/src/hir-runtime.ts new file mode 100644 index 00000000000..190eec0fdfa --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir-runtime.ts @@ -0,0 +1,24 @@ +/** + * Runtime-only entry point for HIR-compiled artifacts. + * + * The simulation engine and its workers import from here — it instantiates + * precompiled artifact sources without pulling the TS→HIR compiler (and its + * `typescript` dependency) into worker bundles. The full pipeline (lowering, + * analyses, linting, artifact compilation) lives in `./hir.ts`. + */ +export { + hirDistributionRuntime, + instantiateHirBufferDynamics, + instantiateHirBufferKernel, + instantiateHirBufferLambda, + instantiateHirUserFn, + type HirArtifacts, + type HirCompiledBufferDynamics, + type HirCompiledBufferKernel, + type HirCompiledBufferLambda, + type HirCompiledUserFn, + type HirDynamicsArtifact, + type HirKernelArtifact, + type HirLambdaArtifact, + type HirParameterValues, +} from "./hir/instantiate"; diff --git a/libs/@hashintel/petrinaut-core/src/hir.ts b/libs/@hashintel/petrinaut-core/src/hir.ts new file mode 100644 index 00000000000..981667a8807 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir.ts @@ -0,0 +1,86 @@ +/** + * Public entry point for the Petrinaut HIR (high-level intermediate + * representation) — see `hir/README.md` for the design document. + * + * The HIR pipeline: TypeScript user code → `lowerTypeScriptToHir` → + * analyses (`typecheckHir`, `analyzeHir`) / linting (`lintHirUserCode`) / + * compilation (`emit-js` via `tryCompileHir*`). + */ +export { + analyzeHir, + foldHir, + type DistributionDag, + type DistributionDagEdge, + type DistributionDagNode, + type DistributionSink, + type HirAnalysis, + type HirBindingInfo, + type HirDependencies, + type HirTokenRead, +} from "./hir/analyze"; +export { + compileHirArtifacts, + tryCompileHirBufferDynamics, + tryCompileHirKernel, + tryCompileHirLambda, + type HirCompileFailure, + type HirCompileResult, +} from "./hir/compile"; +export { + hirDistributionRuntime, + instantiateHirBufferDynamics, + instantiateHirBufferKernel, + instantiateHirBufferLambda, + instantiateHirUserFn, + type HirArtifacts, + type HirCompiledBufferDynamics, + type HirCompiledBufferKernel, + type HirCompiledBufferLambda, + type HirCompiledUserFn, + type HirDynamicsArtifact, + type HirKernelArtifact, + type HirLambdaArtifact, + type HirParameterValues, +} from "./hir/instantiate"; +export { + emitBufferKernelJs, + emitBufferLambdaJs, + type BufferKernelProgram, + type BufferProgram, +} from "./hir/emit-buffer-js"; +export { emitBufferDynamicsJs, emitUserFunctionJs } from "./hir/emit-js"; +export { + formatHirType, + hirChildren, + walkHir, + type HirDiagnostic, + type HirDiagnosticSeverity, + type HirExpr, + type HirFunction, + type HirNodeId, + type HirSurfaceKind, + type HirType, + type Span, +} from "./hir/hir"; +export { + lintHirUserCode, + type HirLintOptions, + type HirLintResult, +} from "./hir/lint"; +export { + lowerTypeScriptToHir, + type LowerTypeScriptResult, +} from "./hir/lower-typescript"; +export { + buildDynamicsContext, + buildKernelContext, + buildLambdaContext, + type HirDynamicsContext, + type HirKernelContext, + type HirLambdaContext, + type HirParameterInfo, + type HirPlaceBinding, + type HirSurfaceContext, + type HirTokenElementInfo, +} from "./hir/surface-context"; +export { typecheckHir, type HirTypecheckResult } from "./hir/typecheck"; diff --git a/libs/@hashintel/petrinaut-core/src/hir/README.md b/libs/@hashintel/petrinaut-core/src/hir/README.md new file mode 100644 index 00000000000..233393db317 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/README.md @@ -0,0 +1,165 @@ +# Petrinaut HIR — high-level intermediate representation + +The HIR is a **typed, source-spanned, JSON-serializable expression tree** for +Petrinaut user code, and its pipeline is the **only compiler** for dynamics, +lambda and transition-kernel code (the Babel `new Function` path is gone): + +``` + (surface languages) (backends) + + TypeScript ──lower-typescript.ts──┐ ┌──emit-buffer-js.ts──▶ buffer-ABI JS + ├──▶ HIR ──▶──┤ (direct packed-buffer reads/writes) + Petrinaut DSL (planned) ──────────┘ │ ├──emit-js.ts─────────▶ object-convention JS + │ │ (structural fallback shapes) + ▼ ├──▶ WASM (planned) + analyses └──▶ GPU kernels (planned) + (typecheck, dependencies, + distribution DAG, semantic lints) +``` + +## Why an HIR at all + +TypeScript is a great authoring surface but a poor analysis substrate. With +the HIR, the things Petrinaut needs to know are one call away: + +- **Distribution DAG of a kernel** — which distributions exist, how they + derive via `.map`, which output attributes their samples flow into, which + outputs share one draw (`analyzeHir(fn).distributionDag`). +- **Dependency sets** — does this lambda depend on parameters? Which token + attributes does it actually read? (`analyzeHir(fn).dependencies`). +- **Determinism** — pure function of (tokens, parameters) or sampling? +- **Compilation** — token attributes are validated against the color schema + (`typecheckHir` + `HirSurfaceContext`) and compiled to **statically + resolved buffer offsets**: `tokenValues[slotBases[slot] + attrIndex]` + instead of decode-to-records / re-encode. + +## Design invariants + +Defined in [`hir.ts`](./hir.ts): + +1. **Expression-oriented and pure** (OCaml-like): `let` bindings, `cond`, + `arrayMap` comprehensions — no mutation, loops or recursion. Guard-clause + `if (c) return a;` and if/else-returns lower to `cond`; `const { a, b } = +parameters` and array destructuring lower to plain bindings. Every + analysis is a single structural pass or one symbolic evaluation. +2. **Every node carries a `Span`** into the user-visible source text — + diagnostics land on exact editor ranges with no sourcemap machinery. +3. **Every node has a stable `id`** — analyses return side tables. +4. **Distributions are first-class node kinds** (`distribution`, + `distributionMap`). +5. **Plain JSON** — `HirFunction` round-trips through `JSON.stringify`. + +### The accepted subset + +`const` (with object/array destructuring and renames), guard `if`s and +early returns, ternaries, arithmetic/comparison/logic, `Math.*`, +`parameters.x` (and destructured aliases), token access +(`input.Place[i].attr`, `.length`), `.map(...)` (including `(token, index)` +and destructured params), `Distribution.*` + `.map`, record/array literals, +`Infinity`/`NaN`/`Math.PI`/`Math.E`, type assertions (erased). + +Rejected with positioned errors (out-of-subset code **cannot run**): loops, +`let`/`var`, spread, strings-as-values, computed keys, arbitrary calls, +helper functions, unreachable code. The shipped example models are the +coverage gate (`compile.test.ts`): all of them must compile fully — and all +of them reach the buffer ABI. + +## Pipeline stages + +| Stage | Module | Contract | +| -------------- | --------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Lowering | `lower-typescript.ts` | TS module → `HirFunction`, or short-circuit with one positioned diagnostic. | +| Context | `surface-context.ts` | Model-derived environment per surface: parameters, token attribute schemas, and the **slot layout invariant** — one slot per colored non-inhibitor input arc token (arc order), mirroring the engine's `inputPlacesWithTokenValues` filter. Duplicate place names resolve last-arc-wins (runtime key-overwrite parity). Net-aware for subnet items. | +| Typecheck | `typecheck.ts` | Bottom-up inference + surface checks against the Color schema: unknown/missing attributes, discrete derivatives, Distribution-into-discrete (H-6519), index vs arc weight, output place coverage and token counts, predicate/stochastic returns. A clean typecheck **gates buffer-ABI emission**. | +| Analyses | `analyze.ts` | Symbolic evaluation → dependencies, distribution DAG (nodes/edges/sinks/shared draws), binding usage; `foldHir` constant folding. | +| Lint | `lint.ts` | All of the above as `HirDiagnostic[]` + semantic rules (`hir:math-random`, `hir:transition-never-fires`, `hir:shared-sample`, `hir:unused-binding`). Out-of-subset = error by default. | +| Buffer backend | `emit-buffer-js.ts` | Symbolic-evaluation emitter producing the **buffer ABI** (below). Structural shapes it cannot scalarize return `null`. | +| Object backend | `emit-js.ts` | Object-convention functions (legacy calling shape) — the fallback for non-scalarizable shapes, plus the buffer-native dynamics loop emitter. | +| Instantiation | `instantiate.ts` | Compiler-free (`typescript`-free) — safe in worker bundles. Defines `HirArtifacts` (v2). | +| Batch compile | `compile.ts` | `compileHirArtifacts(sdcpn, extensions)` → `{ artifacts, failures }` over the root net **and all subnets**, availability-gated like the engine. | + +## The buffer ABI + +Design goal: zero per-combination allocation, token attributes at statically +resolved offsets. All scratch lives on `CompiledTransition.buffer` and is +reused (the engine is single-threaded per simulation instance). + +``` +dynamics: (__params) => (currentState, dimensions, numberOfTokens) => Float64Array +lambda: (tokenValues: Float64Array, slotBases: Int32Array) => number | boolean +kernel: (tokenValues, slotBases, out: Float64Array, + distSink: (floatIndex, distribution) => void) => void +``` + +- **`slotBases`** — one float base offset per input token slot (per colored + non-inhibitor arc, `weight` slots each, arc order). The engine fills it per + enumerated combination with `place.offset + tokenIndex * dimensions` + (`engine/buffer-transition.ts#fillSlotBases`). +- **Reads** compile to `tokenValues[slotBases[k] + attrIndex]` (booleans + `!== 0`); `.length` folds to the static arc weight; `.map` over tuples is + unrolled (weights are static and small). +- **Kernel writes** go place-major into the `out` staging floats: integers + pre-rounded, booleans 0/1. Distribution-valued attributes are deferred via + `distSink`; the engine samples them **ordered by output float index**, + reproducing the legacy (place, token, element) sampling order and hence the + exact RNG stream. Shared draws work by object identity (a `const` + distribution feeding two slots samples once via the sample cache). +- Artifacts carry `inputSlotCount` / `outputFloatCount` so `buildSimulation` + can reject stale buffer programs (it then runs the object program; a + missing artifact altogether is a per-item `SDCPNItemError`). + +Equivalence is enforced by `engine/build-simulation-hir.test.ts`: buffer and +object programs produce **bit-identical frames** over 50 steps of a +stochastic model, including RNG state evolution. + +## Where things run + +- **Compilation** (needs `typescript`): the LSP worker. The worker answers a + `sdcpn/compileHirArtifacts` request (`LanguageClient.requestHirArtifacts`); + the React simulation/experiments providers await it and pass + `hirArtifacts` through `SimulationConfig` / the Monte-Carlo config → worker + init messages → `buildSimulation`. +- **Instantiation** (dependency-free): the simulation and monte-carlo + workers import only [`src/hir-runtime.ts`](../hir-runtime.ts). Removing + Babel shrank those worker bundles from ~3 MB to ~40 kB each. +- **Editor diagnostics**: `lsp/lib/checker.ts` runs the HIR linter per item + after TypeScript (only when TS reports no errors), with `source: "hir"` + and stable numeric codes (99001+). Out-of-subset code is an **error** and + blocks Play (`DiagnosticsSnapshot.errorCount` gates the run controls). + +## Current limitations + +- Structural shapes the buffer emitter cannot scalarize (conditionals over + whole records, dynamic token indices, `.map` over >16 tokens) fall back to + the object-convention program per item — still HIR-compiled, just with the + engine's record decode/encode around it. +- Kernels run the object program when stochasticity is disabled (the object + wrapper carries the distributions-forbidden runtime check). +- Metrics, scenarios and visualizers are separate surfaces: metrics/scenarios + keep their `new Function` compilation (they never used Babel), visualizers + (JSX) keep `@babel/standalone` in the UI package. +- Lowering is syntactic (no `ts.TypeChecker`); `.map` is disambiguated by + tracking distribution-valued bindings. +- Numeric semantics are JS doubles end-to-end; integer attributes are exact + up to ±2^53 (matching the token format). + +## Next steps + +1. **Skip `materializeEngineFrame`** in the single-run engine: the buffer + path only needs place offsets/counts and the token floats — a cheap view + would remove the remaining per-transition-per-step frame copy. +2. **Buffer-native dynamics for more shapes** (cross-token reductions), + and buffer kernels under disabled stochasticity (needs the + distributions-forbidden check on the buffer path). +3. **Metrics & scenarios on the HIR**: needs `reduce`/`filter` comprehension + nodes and a `MetricState` surface context; then their `new Function` + paths can go too. +4. **WASM backend**: the buffer ABI is already the right shape — scalars, + static offsets, no GC; lower `emit-buffer-js` output structure to WAT and + provide host shims for `Math.*`/distributions/RNG. +5. **GPU**: dynamics and Monte-Carlo runs are embarrassingly parallel; the + buffer loop maps onto a WGSL compute shader over the packed frame. +6. **DSL frontend**: see [`dsl-sketch.md`](./dsl-sketch.md) — parses to this + same HIR; TS ⇄ DSL migration is a pretty-printing problem. +7. **Artifact caching**: persist `HirArtifacts` keyed by a content hash of + the code + schema so repeated runs skip recompilation. diff --git a/libs/@hashintel/petrinaut-core/src/hir/analyze.test.ts b/libs/@hashintel/petrinaut-core/src/hir/analyze.test.ts new file mode 100644 index 00000000000..f44b9ab36f6 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/analyze.test.ts @@ -0,0 +1,223 @@ +import { describe, expect, it } from "vitest"; + +import { analyzeHir, foldHir } from "./analyze"; +import { lowerTypeScriptToHir } from "./lower-typescript"; + +import type { HirFunction, HirSurfaceKind } from "./hir"; + +function lower(code: string, surface: HirSurfaceKind): HirFunction { + const result = lowerTypeScriptToHir(code, surface); + if (!result.ok) { + throw new Error(result.diagnostics[0]?.message); + } + return result.fn; +} + +describe("analyzeHir", () => { + describe("dependencies", () => { + it("collects parameter reads", () => { + const analysis = analyzeHir( + lower( + `export default Lambda((input, parameters) => parameters.beta * parameters.alpha + parameters.beta);`, + "lambda", + ), + ); + expect(analysis.dependencies.parameters).toEqual(["alpha", "beta"]); + expect(analysis.dependencies.isDeterministic).toBe(true); + }); + + it("reports a lambda that depends on no parameters", () => { + const analysis = analyzeHir( + lower(`export default Lambda((input, parameters) => 1.5);`, "lambda"), + ); + expect(analysis.dependencies.parameters).toEqual([]); + expect(analysis.dependencies.tokenReads).toEqual([]); + }); + + it("attributes token reads to places for lambdas", () => { + const analysis = analyzeHir( + lower( + `export default Lambda((input, parameters) => input.Pool[0].x + input["Other Place"][1].y);`, + "lambda", + ), + ); + expect(analysis.dependencies.tokenReads).toEqual([ + { place: "Pool", field: "x" }, + { place: "Other Place", field: "y" }, + ]); + }); + + it("attributes token reads through .map callbacks (dynamics)", () => { + const analysis = analyzeHir( + lower( + `export default Dynamics((tokens, parameters) => tokens.map(({ x, y }) => ({ x: y * parameters.g, y: x })));`, + "dynamics", + ), + ); + expect(analysis.dependencies.tokenReads).toEqual([ + { place: "self", field: "y" }, + { place: "self", field: "x" }, + ]); + expect(analysis.dependencies.parameters).toEqual(["g"]); + }); + + it("tracks token count reads and Math.random", () => { + const analysis = analyzeHir( + lower( + `export default Lambda((input) => input.Pool.length * Math.random());`, + "lambda", + ), + ); + expect(analysis.dependencies.readsTokenCounts).toBe(true); + expect(analysis.dependencies.usesMathRandom).toBe(true); + expect(analysis.dependencies.isDeterministic).toBe(false); + }); + }); + + describe("distribution DAG", () => { + it("extracts nodes, edges and kernel output sinks", () => { + const analysis = analyzeHir( + lower( + `export default TransitionKernel((input, parameters) => { + const base = Distribution.Gaussian(0, parameters.sigma); + const scaled = base.map((value) => value * 10); + return { Out: [{ x: scaled, y: 1 }] }; +});`, + "kernel", + ), + ); + const { nodes, edges, sinks } = analysis.distributionDag; + expect(nodes).toHaveLength(2); + + const gaussian = nodes.find((node) => node.kind === "gaussian")!; + const mapped = nodes.find((node) => node.kind === "mapped")!; + expect(gaussian.bindingName).toBe("base"); + expect(gaussian.dependsOnParameters).toEqual(["sigma"]); + expect(mapped.bindingName).toBe("scaled"); + expect(edges).toEqual([{ from: gaussian.nodeId, to: mapped.nodeId }]); + expect(sinks).toEqual([ + { nodeId: mapped.nodeId, place: "Out", tokenIndex: 0, field: "x" }, + ]); + }); + + it("records constant arguments", () => { + const analysis = analyzeHir( + lower( + `export default TransitionKernel((input) => ({ + Out: [{ x: Distribution.Uniform(0, 2 * 5) }], +}));`, + "kernel", + ), + ); + expect(analysis.distributionDag.nodes[0]!.constantArgs).toEqual([0, 10]); + }); + + it("flags shared samples feeding several output slots", () => { + const analysis = analyzeHir( + lower( + `export default TransitionKernel((input) => { + const d = Distribution.Gaussian(0, 1); + return { Out: [{ x: d, y: d }] }; +});`, + "kernel", + ), + ); + const nodeId = analysis.distributionDag.nodes[0]!.nodeId; + expect(analysis.distributionDag.sinks).toHaveLength(2); + expect(analysis.distributionDag.sharedSampleNodeIds).toEqual([nodeId]); + }); + + it("marks per-iteration distributions and dynamic sinks", () => { + const analysis = analyzeHir( + lower( + `export default TransitionKernel((input) => ({ + Out: input.In.map((token) => ({ x: Distribution.Gaussian(token.x, 1) })), +}));`, + "kernel", + ), + ); + const node = analysis.distributionDag.nodes[0]!; + expect(node.perIteration).toBe(true); + expect(node.dependsOnTokens).toEqual([{ place: "In", field: "x" }]); + expect(analysis.distributionDag.sinks).toEqual([ + { + nodeId: node.nodeId, + place: "Out", + tokenIndex: "dynamic", + field: "x", + }, + ]); + expect(analysis.distributionDag.sharedSampleNodeIds).toEqual([]); + }); + + it("tracks distributions through conditionals", () => { + const analysis = analyzeHir( + lower( + `export default TransitionKernel((input, parameters) => { + const a = Distribution.Gaussian(0, 1); + const b = Distribution.Uniform(0, 1); + return { Out: [{ x: parameters.flag ? a : b }] }; +});`, + "kernel", + ), + ); + expect(analysis.distributionDag.sinks).toHaveLength(2); + }); + }); + + describe("bindings", () => { + it("counts references and finds unused bindings", () => { + const analysis = analyzeHir( + lower( + `export default Lambda((input, parameters) => { + const used = parameters.a; + const unused = parameters.b; + return used * 2; +});`, + "lambda", + ), + ); + const used = analysis.bindings.find( + (binding) => binding.name === "used", + )!; + const unused = analysis.bindings.find( + (binding) => binding.name === "unused", + )!; + expect(used.referenceCount).toBe(1); + expect(unused.referenceCount).toBe(0); + }); + }); +}); + +describe("foldHir", () => { + it("folds constant arithmetic, conditionals and Math calls", () => { + const fn = lower( + `export default Lambda((input) => 1 + 2 * 3 > 5 ? Math.sqrt(16) : 0);`, + "lambda", + ); + const folded = foldHir(fn.body); + expect(folded).toMatchObject({ kind: "numberLit", value: 4 }); + }); + + it("keeps non-constant parts intact", () => { + const fn = lower( + `export default Lambda((input, parameters) => parameters.rate * (2 + 3));`, + "lambda", + ); + const folded = foldHir(fn.body); + expect(folded).toMatchObject({ + kind: "binary", + op: "*", + left: { kind: "paramRef" }, + right: { kind: "numberLit", value: 5 }, + }); + }); + + it("never folds Math.random", () => { + const fn = lower( + `export default Lambda((input) => Math.random());`, + "lambda", + ); + expect(foldHir(fn.body).kind).toBe("mathCall"); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/analyze.ts b/libs/@hashintel/petrinaut-core/src/hir/analyze.ts new file mode 100644 index 00000000000..64706c093ce --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/analyze.ts @@ -0,0 +1,742 @@ +/** + * Static analyses over the HIR. + * + * All analyses share one abstract interpreter: because the HIR is pure, + * non-recursive and loop-free (only structural `map`), a single symbolic + * evaluation of the body is enough to answer: + * + * - which model parameters and token attributes the code depends on, + * - which probability distributions it constructs, how they derive from one + * another (the distribution DAG), and into which output slots their samples + * flow, + * - which distributions share a single draw (the runtime caches a + * distribution object's first sample, so one binding feeding two output + * fields yields identical values), + * - which `const` bindings are never used. + */ +import type { + HirDistributionKind, + HirExpr, + HirFunction, + HirNodeId, + Span, +} from "./hir"; + +/** + * The place a token read refers to: `"self"` for dynamics (the place the + * equation is attached to), otherwise the place display name. The `& {}` + * keeps the `"self"` literal from being absorbed into `string`. + */ +export type HirTokenReadPlace = "self" | (string & {}); + +/** A token attribute read. */ +export type HirTokenRead = { + place: HirTokenReadPlace; + field: string; +}; + +export type HirDependencies = { + /** Model parameters read via `parameters.` (sorted, unique). */ + parameters: string[]; + /** Token attributes read (unique by place+field). */ + tokenReads: HirTokenRead[]; + /** Whether the code reads token counts (`.length`). */ + readsTokenCounts: boolean; + /** Whether the code constructs probability distributions. */ + samplesDistributions: boolean; + /** Whether the code calls `Math.random()` (breaks reproducibility). */ + usesMathRandom: boolean; + /** True when the result is a pure function of (tokens, parameters). */ + isDeterministic: boolean; +}; + +export type DistributionDagNode = { + /** HIR node id of the `distribution` / `distributionMap` node. */ + nodeId: HirNodeId; + kind: HirDistributionKind | "mapped"; + span: Span; + /** Name of the `const` binding holding this distribution, if any. */ + bindingName?: string; + /** True when constructed inside a `.map(...)` over tokens — a fresh + * distribution (and draw) is created per token. */ + perIteration: boolean; + /** Argument values when they fold to constants, e.g. `[0, 10]`. */ + constantArgs?: number[]; + /** Parameters feeding the arguments (or the map body). */ + dependsOnParameters: string[]; + /** Token attributes feeding the arguments (or the map body). */ + dependsOnTokens: HirTokenRead[]; +}; + +/** Derivation edge: `to` is `from.map(...)`. */ +export type DistributionDagEdge = { + from: HirNodeId; + to: HirNodeId; +}; + +/** Where a distribution's sample lands in a kernel's output. */ +export type DistributionSink = { + nodeId: HirNodeId; + place: string; + /** Token position within the place's output array, or `"dynamic"` when the + * token is produced by a `.map(...)`. */ + tokenIndex: number | "dynamic"; + field: string; +}; + +export type DistributionDag = { + nodes: DistributionDagNode[]; + edges: DistributionDagEdge[]; + sinks: DistributionSink[]; + /** Nodes whose single draw feeds more than one output slot. */ + sharedSampleNodeIds: HirNodeId[]; +}; + +export type HirBindingInfo = { + name: string; + nameSpan: Span; + referenceCount: number; +}; + +export type HirAnalysis = { + dependencies: HirDependencies; + distributionDag: DistributionDag; + bindings: HirBindingInfo[]; +}; + +// --------------------------------------------------------------------------- +// Abstract values +// --------------------------------------------------------------------------- + +type AbstractValue = + /** An opaque scalar (number/boolean). */ + | { kind: "scalar" } + /** The lambda/kernel input object (`tokensByPlace`). */ + | { kind: "inputRecord" } + /** A token array belonging to a place. */ + | { kind: "tokens"; place: HirTokenReadPlace } + /** A single token of a place. */ + | { kind: "token"; place: HirTokenReadPlace } + | { kind: "record"; fields: Map } + /** `elements` is null for arrays of statically-unknown shape (map results + * carry the per-element value in `element`). */ + | { + kind: "array"; + elements: AbstractValue[] | null; + element?: AbstractValue; + } + | { kind: "dist"; nodeId: HirNodeId } + /** Either of several values (from conditionals). */ + | { kind: "union"; values: AbstractValue[] }; + +const SCALAR: AbstractValue = { kind: "scalar" }; + +type BindingRecord = { + name: string; + nameSpan: Span; + referenceCount: number; +}; + +type DepSink = { + parameters: Set; + tokenReads: Map; +}; + +function createDepSink(): DepSink { + return { parameters: new Set(), tokenReads: new Map() }; +} + +// --------------------------------------------------------------------------- +// Constant folding +// --------------------------------------------------------------------------- + +function constantValue(name: "PI" | "E" | "Infinity" | "NaN"): number { + switch (name) { + case "PI": + return Math.PI; + case "E": + return Math.E; + case "Infinity": + return Infinity; + case "NaN": + return NaN; + } +} + +function literalValue(expr: HirExpr): number | boolean | undefined { + if (expr.kind === "numberLit") { + return expr.value; + } + if (expr.kind === "boolLit") { + return expr.value; + } + if (expr.kind === "constant") { + return constantValue(expr.name); + } + return undefined; +} + +function numberNode( + original: HirExpr, + value: number, +): Extract { + return { + kind: "numberLit", + value, + raw: String(value), + id: original.id, + span: original.span, + }; +} + +function foldBinary( + op: Extract["op"], + left: number | boolean, + right: number | boolean, +): number | boolean | undefined { + switch (op) { + case "+": + return typeof left === "number" && typeof right === "number" + ? left + right + : undefined; + case "-": + return Number(left) - Number(right); + case "*": + return Number(left) * Number(right); + case "/": + return Number(left) / Number(right); + case "%": + return Number(left) % Number(right); + case "**": + return Number(left) ** Number(right); + case "<": + return Number(left) < Number(right); + case "<=": + return Number(left) <= Number(right); + case ">": + return Number(left) > Number(right); + case ">=": + return Number(left) >= Number(right); + case "==": + return left === right; + case "!=": + return left !== right; + case "&&": + return typeof left === "boolean" && typeof right === "boolean" + ? left && right + : undefined; + case "||": + return typeof left === "boolean" && typeof right === "boolean" + ? left || right + : undefined; + } +} + +/** + * Folds constant subexpressions. Folded nodes keep the id and span of the + * expression they replace, so diagnostics and analyses remain anchored. + * `Math.random()` and distributions are never folded. + */ +export function foldHir(expr: HirExpr): HirExpr { + switch (expr.kind) { + case "numberLit": + case "boolLit": + case "constant": + case "localRef": + case "paramRef": + return expr; + case "unary": { + const operand = foldHir(expr.operand); + const value = literalValue(operand); + if (value !== undefined) { + if (expr.op === "!" && typeof value === "boolean") { + return { ...expr, kind: "boolLit", value: !value } as HirExpr; + } + if (typeof value === "number") { + if (expr.op === "-") { + return numberNode(expr, -value); + } + if (expr.op === "+") { + return numberNode(expr, value); + } + } + } + return { ...expr, operand }; + } + case "binary": { + const left = foldHir(expr.left); + const right = foldHir(expr.right); + const leftValue = literalValue(left); + const rightValue = literalValue(right); + if (leftValue !== undefined && rightValue !== undefined) { + const folded = foldBinary(expr.op, leftValue, rightValue); + if (folded !== undefined) { + return typeof folded === "boolean" + ? ({ ...expr, kind: "boolLit", value: folded } as HirExpr) + : numberNode(expr, folded); + } + } + return { ...expr, left, right }; + } + case "cond": { + const condition = foldHir(expr.condition); + if (condition.kind === "boolLit") { + return condition.value + ? foldHir(expr.thenBranch) + : foldHir(expr.elseBranch); + } + return { + ...expr, + condition, + thenBranch: foldHir(expr.thenBranch), + elseBranch: foldHir(expr.elseBranch), + }; + } + case "let": { + return { + ...expr, + bindings: expr.bindings.map((binding) => ({ + ...binding, + value: foldHir(binding.value), + })), + body: foldHir(expr.body), + }; + } + case "mathCall": { + const args = expr.args.map(foldHir); + if (expr.fn !== "random") { + const values = args.map(literalValue); + if (values.every((value) => typeof value === "number")) { + const mathFn = Math[expr.fn] as (...values: number[]) => number; + return numberNode(expr, mathFn(...(values as number[]))); + } + } + return { ...expr, args }; + } + case "fieldAccess": + return { ...expr, target: foldHir(expr.target) }; + case "indexAccess": + return { + ...expr, + target: foldHir(expr.target), + index: foldHir(expr.index), + }; + case "length": + return { ...expr, target: foldHir(expr.target) }; + case "recordLit": + return { + ...expr, + entries: expr.entries.map((entry) => ({ + ...entry, + value: foldHir(entry.value), + })), + }; + case "arrayLit": + return { ...expr, elements: expr.elements.map(foldHir) }; + case "arrayMap": + return { + ...expr, + target: foldHir(expr.target), + body: foldHir(expr.body), + }; + case "distribution": + return { ...expr, args: expr.args.map(foldHir) }; + case "distributionMap": + return { ...expr, base: foldHir(expr.base), body: foldHir(expr.body) }; + } +} + +function constantArgsOf(args: HirExpr[]): number[] | undefined { + const values: number[] = []; + for (const argument of args) { + const folded = foldHir(argument); + if (folded.kind === "numberLit") { + values.push(folded.value); + } else if (folded.kind === "constant") { + values.push(constantValue(folded.name)); + } else { + return undefined; + } + } + return values; +} + +class Analyzer { + readonly globalDeps = createDepSink(); + readonly depSinkStack: DepSink[] = []; + readonly dagNodes = new Map(); + readonly dagEdges: DistributionDagEdge[] = []; + readonly bindings: BindingRecord[] = []; + readsTokenCounts = false; + usesMathRandom = false; + /** Depth of `.map(...)` iteration during evaluation. */ + private mapDepth = 0; + + constructor(private readonly fn: HirFunction) {} + + private recordParameter(name: string): void { + this.globalDeps.parameters.add(name); + for (const sink of this.depSinkStack) { + sink.parameters.add(name); + } + } + + private recordTokenRead(read: HirTokenRead): void { + const key = `${read.place}\u0000${read.field}`; + this.globalDeps.tokenReads.set(key, read); + for (const sink of this.depSinkStack) { + sink.tokenReads.set(key, read); + } + } + + private withDepSink(run: () => Result): [Result, DepSink] { + const sink = createDepSink(); + this.depSinkStack.push(sink); + try { + return [run(), sink]; + } finally { + this.depSinkStack.pop(); + } + } + + evaluate(): AbstractValue { + const env = new Map< + string, + { value: AbstractValue; binding?: BindingRecord } + >(); + const tokensParam = this.fn.params[0]; + if (tokensParam) { + env.set(tokensParam.name, { + value: + this.fn.surface === "dynamics" + ? { kind: "tokens", place: "self" } + : { kind: "inputRecord" }, + }); + } + return this.evalExpr(this.fn.body, env); + } + + private evalExpr( + expr: HirExpr, + env: Map, + ): AbstractValue { + switch (expr.kind) { + case "numberLit": + case "boolLit": + case "constant": + return SCALAR; + case "localRef": { + const entry = env.get(expr.name); + if (entry) { + if (entry.binding) { + entry.binding.referenceCount += 1; + } + return entry.value; + } + return SCALAR; + } + case "paramRef": + this.recordParameter(expr.name); + return SCALAR; + case "fieldAccess": { + const target = this.evalExpr(expr.target, env); + return this.accessField(target, expr.field); + } + case "indexAccess": { + const target = this.evalExpr(expr.target, env); + this.evalExpr(expr.index, env); + return this.accessIndex(target, expr.index); + } + case "length": { + const target = this.evalExpr(expr.target, env); + if (target.kind === "tokens" || target.kind === "inputRecord") { + this.readsTokenCounts = true; + } + return SCALAR; + } + case "unary": + this.evalExpr(expr.operand, env); + return SCALAR; + case "binary": + this.evalExpr(expr.left, env); + this.evalExpr(expr.right, env); + return SCALAR; + case "cond": { + this.evalExpr(expr.condition, env); + const thenValue = this.evalExpr(expr.thenBranch, env); + const elseValue = this.evalExpr(expr.elseBranch, env); + if (thenValue === elseValue) { + return thenValue; + } + return { kind: "union", values: [thenValue, elseValue] }; + } + case "let": { + const scoped = new Map(env); + for (const bindingExpr of expr.bindings) { + const binding: BindingRecord = { + name: bindingExpr.name, + nameSpan: bindingExpr.nameSpan, + referenceCount: 0, + }; + this.bindings.push(binding); + const value = this.evalExpr(bindingExpr.value, scoped); + this.nameDistribution(value, bindingExpr.name); + scoped.set(bindingExpr.name, { value, binding }); + } + return this.evalExpr(expr.body, scoped); + } + case "mathCall": { + if (expr.fn === "random") { + this.usesMathRandom = true; + } + for (const argument of expr.args) { + this.evalExpr(argument, env); + } + return SCALAR; + } + case "recordLit": { + const fields = new Map(); + for (const entry of expr.entries) { + fields.set(entry.key, this.evalExpr(entry.value, env)); + } + return { kind: "record", fields }; + } + case "arrayLit": + return { + kind: "array", + elements: expr.elements.map((element) => this.evalExpr(element, env)), + }; + case "arrayMap": { + const target = this.evalExpr(expr.target, env); + const scoped = new Map(env); + scoped.set(expr.param.name, { value: this.elementOf(target) }); + if (expr.indexParam) { + scoped.set(expr.indexParam.name, { value: SCALAR }); + } + this.mapDepth += 1; + const element = this.evalExpr(expr.body, scoped); + this.mapDepth -= 1; + return { kind: "array", elements: null, element }; + } + case "distribution": { + const [, argDeps] = this.withDepSink(() => { + for (const argument of expr.args) { + this.evalExpr(argument, env); + } + }); + const node: DistributionDagNode = { + nodeId: expr.id, + kind: expr.dist, + span: expr.span, + perIteration: this.mapDepth > 0, + constantArgs: constantArgsOf(expr.args), + dependsOnParameters: [...argDeps.parameters].sort(), + dependsOnTokens: [...argDeps.tokenReads.values()], + }; + this.dagNodes.set(expr.id, node); + return { kind: "dist", nodeId: expr.id }; + } + case "distributionMap": { + const base = this.evalExpr(expr.base, env); + const [, bodyDeps] = this.withDepSink(() => { + const scoped = new Map(env); + scoped.set(expr.param.name, { value: SCALAR }); + this.evalExpr(expr.body, scoped); + }); + const node: DistributionDagNode = { + nodeId: expr.id, + kind: "mapped", + span: expr.span, + perIteration: this.mapDepth > 0, + dependsOnParameters: [...bodyDeps.parameters].sort(), + dependsOnTokens: [...bodyDeps.tokenReads.values()], + }; + this.dagNodes.set(expr.id, node); + for (const baseNodeId of this.distributionNodeIds(base)) { + this.dagEdges.push({ from: baseNodeId, to: expr.id }); + } + return { kind: "dist", nodeId: expr.id }; + } + } + } + + private accessField(target: AbstractValue, field: string): AbstractValue { + switch (target.kind) { + case "inputRecord": + return { kind: "tokens", place: field }; + case "token": + this.recordTokenRead({ place: target.place, field }); + return SCALAR; + case "record": + return target.fields.get(field) ?? SCALAR; + case "union": + return { + kind: "union", + values: target.values.map((value) => this.accessField(value, field)), + }; + default: + return SCALAR; + } + } + + private accessIndex(target: AbstractValue, _index: HirExpr): AbstractValue { + switch (target.kind) { + case "tokens": + return { kind: "token", place: target.place }; + case "array": + if (target.elements) { + // Index may be dynamic; the result may be any element. + return target.elements.length === 1 + ? target.elements[0]! + : { kind: "union", values: target.elements }; + } + return target.element ?? SCALAR; + case "union": + return { + kind: "union", + values: target.values.map((value) => this.accessIndex(value, _index)), + }; + default: + return SCALAR; + } + } + + private elementOf(target: AbstractValue): AbstractValue { + switch (target.kind) { + case "tokens": + return { kind: "token", place: target.place }; + case "array": + if (target.elements) { + return target.elements.length === 1 + ? target.elements[0]! + : { kind: "union", values: target.elements }; + } + return target.element ?? SCALAR; + default: + return SCALAR; + } + } + + private nameDistribution(value: AbstractValue, name: string): void { + for (const nodeId of this.distributionNodeIds(value)) { + const node = this.dagNodes.get(nodeId); + if (node && node.bindingName === undefined) { + node.bindingName = name; + } + } + } + + private distributionNodeIds(value: AbstractValue): HirNodeId[] { + switch (value.kind) { + case "dist": + return [value.nodeId]; + case "union": + return value.values.flatMap((inner) => this.distributionNodeIds(inner)); + default: + return []; + } + } + + collectSinks(result: AbstractValue): DistributionSink[] { + if (this.fn.surface !== "kernel") { + return []; + } + const sinks: DistributionSink[] = []; + const visitToken = ( + token: AbstractValue, + place: string, + tokenIndex: number | "dynamic", + ): void => { + if (token.kind === "record") { + for (const [field, fieldValue] of token.fields) { + for (const nodeId of this.distributionNodeIds(fieldValue)) { + sinks.push({ nodeId, place, tokenIndex, field }); + } + } + } else if (token.kind === "union") { + for (const inner of token.values) { + visitToken(inner, place, tokenIndex); + } + } + }; + const visitPlaceValue = (value: AbstractValue, place: string): void => { + if (value.kind === "array") { + if (value.elements) { + for (const [index, token] of value.elements.entries()) { + visitToken(token, place, index); + } + } else if (value.element) { + visitToken(value.element, place, "dynamic"); + } + } else if (value.kind === "union") { + for (const inner of value.values) { + visitPlaceValue(inner, place); + } + } + }; + const visitOutput = (value: AbstractValue): void => { + if (value.kind === "record") { + for (const [place, placeValue] of value.fields) { + visitPlaceValue(placeValue, place); + } + } else if (value.kind === "union") { + for (const inner of value.values) { + visitOutput(inner); + } + } + }; + visitOutput(result); + return sinks; + } +} + +/** Runs all analyses over a lowered function. */ +export function analyzeHir(fn: HirFunction): HirAnalysis { + const analyzer = new Analyzer(fn); + const result = analyzer.evaluate(); + const sinks = analyzer.collectSinks(result); + + const sinkCountByNode = new Map(); + for (const sink of sinks) { + sinkCountByNode.set( + sink.nodeId, + (sinkCountByNode.get(sink.nodeId) ?? 0) + 1, + ); + } + const sharedSampleNodeIds = [...sinkCountByNode.entries()] + .filter(([nodeId, count]) => { + const node = analyzer.dagNodes.get(nodeId); + // A per-iteration distribution is fresh per token; sharing across a + // dynamic token array is not sharing a draw. Sharing within one token + // record still is, but we cannot distinguish that statically here, so + // stay conservative and only flag non-iterated nodes. + return count > 1 && node && !node.perIteration; + }) + .map(([nodeId]) => nodeId); + + const dagNodes = [...analyzer.dagNodes.values()]; + + const dependencies: HirDependencies = { + parameters: [...analyzer.globalDeps.parameters].sort(), + tokenReads: [...analyzer.globalDeps.tokenReads.values()], + readsTokenCounts: analyzer.readsTokenCounts, + samplesDistributions: dagNodes.length > 0, + usesMathRandom: analyzer.usesMathRandom, + isDeterministic: dagNodes.length === 0 && !analyzer.usesMathRandom, + }; + + return { + dependencies, + distributionDag: { + nodes: dagNodes, + edges: analyzer.dagEdges, + sinks, + sharedSampleNodeIds, + }, + bindings: analyzer.bindings.map((binding) => ({ + name: binding.name, + nameSpan: binding.nameSpan, + referenceCount: binding.referenceCount, + })), + }; +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/compile.test.ts b/libs/@hashintel/petrinaut-core/src/hir/compile.test.ts new file mode 100644 index 00000000000..69e2f3ca9be --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/compile.test.ts @@ -0,0 +1,74 @@ +/** + * Coverage gate: every example model must compile fully through the HIR — + * there is no fallback compiler anymore, so any lowering failure here means + * a shipped model cannot simulate. + */ +import { describe, expect, it } from "vitest"; + +import { + deploymentPipelineSDCPN, + probabilisticSatellitesSDCPN, + productionMachines, + sirModel, + supplyChainWithDisruption, +} from "../examples/index"; +import { buildSimulation } from "../simulation/engine/build-simulation"; +import { compileHirArtifacts } from "./compile"; + +import type { SDCPN } from "../types/sdcpn"; + +const EXAMPLES: [string, SDCPN][] = [ + ["production-with-machine-failure", productionMachines.petriNetDefinition], + ["deployment-pipeline", deploymentPipelineSDCPN.petriNetDefinition], + ["satellites-launcher", probabilisticSatellitesSDCPN.petriNetDefinition], + ["sir-model", sirModel.petriNetDefinition], + [ + "supply-chain-with-disruption", + supplyChainWithDisruption.petriNetDefinition, + ], +]; + +describe("compileHirArtifacts on example models", () => { + it.each(EXAMPLES)("compiles every item of %s", (_name, sdcpn) => { + const { failures } = compileHirArtifacts(sdcpn); + expect( + failures.map((failure) => ({ + item: `${failure.itemType}:${failure.itemId}`, + message: failure.diagnostics[0]?.message, + })), + ).toEqual([]); + }); + + it.each(EXAMPLES)( + "compiles every lambda and kernel of %s to the buffer ABI", + (_name, sdcpn) => { + const { artifacts } = compileHirArtifacts(sdcpn); + const withoutBuffer = [ + ...Object.entries(artifacts.lambdas) + .filter(([, artifact]) => !artifact.buffer) + .map(([id]) => `lambda:${id}`), + ...Object.entries(artifacts.kernels) + .filter(([, artifact]) => !artifact.buffer) + .map(([id]) => `kernel:${id}`), + ...Object.entries(artifacts.dynamics) + .filter(([, artifact]) => !artifact.buffer) + .map(([id]) => `dynamics:${id}`), + ]; + expect(withoutBuffer).toEqual([]); + }, + ); + + it.each(EXAMPLES)("builds a runnable simulation for %s", (_name, sdcpn) => { + const { artifacts } = compileHirArtifacts(sdcpn); + const simulation = buildSimulation({ + sdcpn, + initialMarking: {}, + parameterValues: {}, + seed: 1, + dt: 0.1, + maxTime: 1, + hirArtifacts: artifacts, + }); + expect(simulation.compiledTransitions.size).toBeGreaterThan(0); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/compile.ts b/libs/@hashintel/petrinaut-core/src/hir/compile.ts new file mode 100644 index 00000000000..a38225ba5ec --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/compile.ts @@ -0,0 +1,278 @@ +/** + * Compilation of user code through the HIR pipeline. + * + * `compileHirArtifacts` batch-compiles a whole SDCPN (root net and all + * subnets) into serializable artifact sources (`HirArtifacts`) — the only + * runtime compilation path. Per item it produces: + * + * - a buffer-ABI program (`emit-buffer-js.ts`) when the code typechecks + * cleanly and its shape scalarizes to direct packed-buffer reads/writes, + * - an object-convention fallback (`emit-js.ts`) whenever lowering succeeds. + * + * Items whose code cannot be lowered get no artifact and are reported in + * `failures`; `buildSimulation` refuses to run them (the LSP shows the same + * diagnostics with exact source ranges). + * + * Note: this module (transitively) imports `typescript` — keep it out of the + * simulation worker bundles; those only need `instantiate.ts`. + */ +import { + DEFAULT_PETRINAUT_EXTENSIONS, + getTransitionLogicAvailability, + sanitizeSDCPNForExtensions, + type PetrinautExtensionSettings, +} from "../extensions"; +import { emitBufferKernelJs, emitBufferLambdaJs } from "./emit-buffer-js"; +import { emitBufferDynamicsJs, emitUserFunctionJs } from "./emit-js"; +import { + instantiateHirBufferDynamics, + instantiateHirUserFn, +} from "./instantiate"; +import { lowerTypeScriptToHir } from "./lower-typescript"; +import { + buildDynamicsContext, + buildKernelContext, + buildLambdaContext, +} from "./surface-context"; +import { typecheckHir } from "./typecheck"; + +import type { SDCPN, Subnet } from "../types/sdcpn"; +import type { HirDiagnostic, HirFunction, HirSurfaceKind } from "./hir"; +import type { + HirArtifacts, + HirCompiledBufferDynamics, + HirCompiledUserFn, + HirParameterValues, +} from "./instantiate"; +import type { + HirNetScope, + HirSurfaceContext, + HirTokenElementInfo, +} from "./surface-context"; + +export type HirCompileFailure = { + itemId: string; + itemType: "differential-equation" | "transition-lambda" | "transition-kernel"; + diagnostics: HirDiagnostic[]; +}; + +export type HirCompileResult = { + artifacts: HirArtifacts; + /** Items whose code could not be lowered (no artifact emitted). */ + failures: HirCompileFailure[]; +}; + +type LoweredItem = { + fn: HirFunction; + /** True when typechecking against the context produced no errors — the + * gate for buffer-ABI emission. */ + typecheckClean: boolean; +}; + +function lowerAndCheck( + code: string, + surface: HirSurfaceKind, + context: HirSurfaceContext | null, +): LoweredItem | HirDiagnostic[] { + const lowered = lowerTypeScriptToHir(code, surface); + if (!lowered.ok) { + return lowered.diagnostics; + } + let typecheckClean = false; + if (context) { + const checked = typecheckHir(lowered.fn, context); + typecheckClean = !checked.diagnostics.some( + (diagnostic) => diagnostic.severity === "error", + ); + } + return { fn: lowered.fn, typecheckClean }; +} + +/** + * Batch-compiles all dynamics/lambda/kernel code of an SDCPN (root and + * subnets) to HIR artifacts for `buildSimulation`. Item availability mirrors + * the engine: empty lambdas and kernels without colored output places are + * skipped (the engine substitutes defaults for those). + */ +export function compileHirArtifacts( + sdcpn: SDCPN, + extensions: PetrinautExtensionSettings = DEFAULT_PETRINAUT_EXTENSIONS, +): HirCompileResult { + const sanitized = sanitizeSDCPNForExtensions(sdcpn, extensions); + const artifacts: HirArtifacts = { + version: 2, + dynamics: {}, + lambdas: {}, + kernels: {}, + }; + const failures: HirCompileFailure[] = []; + + const colorById = new Map( + [ + ...sanitized.types, + ...(sanitized.subnets ?? []).flatMap((subnet) => subnet.types), + ].map((color) => [color.id, color]), + ); + + const nets: { net: HirNetScope; subnet: Subnet | null }[] = [ + { net: sanitized, subnet: null }, + ...(sanitized.subnets ?? []).map((subnet) => ({ net: subnet, subnet })), + ]; + + for (const { net, subnet } of nets) { + const differentialEquations = subnet + ? subnet.differentialEquations + : sanitized.differentialEquations; + const transitions = subnet ? subnet.transitions : sanitized.transitions; + + for (const de of differentialEquations) { + const color = de.colorId ? colorById.get(de.colorId) : undefined; + if ( + !color || + !color.elements.some((element) => element.type === "real") + ) { + continue; + } + const context = de.colorId + ? buildDynamicsContext(sanitized, de.colorId, extensions, net) + : null; + const lowered = lowerAndCheck(de.code, "dynamics", context); + if (Array.isArray(lowered)) { + failures.push({ + itemId: de.id, + itemType: "differential-equation", + diagnostics: lowered, + }); + continue; + } + const buffer = lowered.typecheckClean + ? emitBufferDynamicsJs(lowered.fn, color.elements) + : null; + artifacts.dynamics[de.id] = { + ...(buffer !== null ? { buffer } : {}), + object: emitUserFunctionJs(lowered.fn), + }; + } + + for (const transition of transitions) { + const availability = getTransitionLogicAvailability( + transition, + sanitized, + extensions, + net, + ); + + if (availability.lambda && transition.lambdaCode.trim() !== "") { + const context = buildLambdaContext( + sanitized, + transition, + extensions, + net, + ); + const lowered = lowerAndCheck(transition.lambdaCode, "lambda", context); + if (Array.isArray(lowered)) { + failures.push({ + itemId: transition.id, + itemType: "transition-lambda", + diagnostics: lowered, + }); + } else { + const buffer = lowered.typecheckClean + ? emitBufferLambdaJs(lowered.fn, context) + : null; + artifacts.lambdas[transition.id] = { + ...(buffer !== null ? { buffer } : {}), + object: emitUserFunctionJs(lowered.fn), + }; + } + } + + if (availability.transitionKernel) { + const context = buildKernelContext( + sanitized, + transition, + extensions, + net, + ); + const lowered = lowerAndCheck( + transition.transitionKernelCode, + "kernel", + context, + ); + if (Array.isArray(lowered)) { + failures.push({ + itemId: transition.id, + itemType: "transition-kernel", + diagnostics: lowered, + }); + } else { + const buffer = lowered.typecheckClean + ? emitBufferKernelJs(lowered.fn, context) + : null; + artifacts.kernels[transition.id] = { + ...(buffer !== null ? { buffer } : {}), + object: emitUserFunctionJs(lowered.fn), + }; + } + } + } + } + + return { artifacts, failures }; +} + +// --------------------------------------------------------------------------- +// Single-item helpers (tests, tooling) +// --------------------------------------------------------------------------- + +function emitObjectSource( + code: string, + surface: HirSurfaceKind, +): string | null { + try { + const lowered = lowerTypeScriptToHir(code, surface); + if (!lowered.ok) { + return null; + } + return emitUserFunctionJs(lowered.fn); + } catch { + return null; + } +} + +/** Compiles a `Lambda(...)` module to an object-convention function, or + * `null` when the code is outside the HIR subset. */ +export function tryCompileHirLambda(code: string): HirCompiledUserFn | null { + const source = emitObjectSource(code, "lambda"); + return source === null ? null : instantiateHirUserFn(source); +} + +/** Compiles a `TransitionKernel(...)` module to an object-convention + * function, or `null` when the code is outside the HIR subset. */ +export function tryCompileHirKernel(code: string): HirCompiledUserFn | null { + const source = emitObjectSource(code, "kernel"); + return source === null ? null : instantiateHirUserFn(source); +} + +/** Compiles a `Dynamics(...)` module to a buffer-native derivative function + * with `parameterValues` pre-bound, or `null` when the body doesn't fit the + * buffer-native shape. */ +export function tryCompileHirBufferDynamics( + code: string, + elements: readonly HirTokenElementInfo[], + parameterValues: HirParameterValues, +): HirCompiledBufferDynamics | null { + try { + const lowered = lowerTypeScriptToHir(code, "dynamics"); + if (!lowered.ok) { + return null; + } + const source = emitBufferDynamicsJs(lowered.fn, elements); + if (source === null) { + return null; + } + return instantiateHirBufferDynamics(source, parameterValues); + } catch { + return null; + } +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/dsl-sketch.md b/libs/@hashintel/petrinaut-core/src/hir/dsl-sketch.md new file mode 100644 index 00000000000..32c47a43d1a --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/dsl-sketch.md @@ -0,0 +1,231 @@ +# Petrinaut DSL — design sketch + +A domain-specific language for Petrinaut user code (dynamics, lambdas, +kernels, metrics, scenario expressions), designed to **lower to the same HIR +as TypeScript** ([README](./README.md)). Users could switch a model — or a +single item — from TypeScript to the DSL without changing semantics, because +both frontends meet at the HIR. + +Status: design sketch only; nothing here is implemented. + +## Design goals + +1. **Analyzable by construction.** The TS frontend must reject code outside + the HIR subset; the DSL simply cannot express anything outside it. Every + valid DSL program lowers, so the runtime never falls back and every model + gets the distribution DAG, dependency info and optimized compilation. +2. **Expression-oriented, OCaml-flavoured.** `let … in` bindings, `if/then/else` + as an expression, no statements, no mutation. +3. **Distributions are syntax.** `~Gaussian(0, sigma)` is a literal, visibly + distinct from a number — the stochastic structure of a kernel is readable + at a glance and trivially extractable. +4. **Less ceremony than the TS surface.** No + `export default TransitionKernel((input, parameters) => …)` wrapper; the + editor pane _is_ the function body, and the model context (places, + parameters, attributes) is ambient. +5. **LSP-native.** Hand-written parser with error recovery producing a + partial tree — diagnostics, completions and hovers must work on incomplete + code while typing. + +## Surface examples + +Dynamics (per-token derivative; `'` marks a derivative binding — only real +attributes admit one): + +``` +per { x, v } -> + x' = v + v' = -params.k * x +``` + +Lambda — stochastic rate (`Pool` is an input place, `count Pool` its token +count; `inf` = always fire): + +``` +let pressure = count Pool * params.rate in +if pressure > 10 then inf else pressure +``` + +Lambda — predicate: + +``` +Pool[0].active && Pool[0].x >= params.threshold +``` + +Transition kernel (`~` introduces a distribution; `|>` maps over it; +output places are record labels checked against the transition's output arcs): + +``` +let noise = ~Gaussian(0, params.sigma) in +{ Target = [ { x = Pool[0].x + , v = noise |> v -> v * 2 + , generation = Pool[0].generation + 1 } ] } +``` + +Kernel over all input tokens: + +``` +{ Out = Pool.map(token -> { x = token.x + 1, v = token.v }) } +``` + +Metric (future surface): + +``` +sum Infected.tokens by t -> t.viral_load +``` + +Notes: + +- `params.` is the single namespace for model parameters (mirrors + `paramRef`). +- Input places are bare identifiers (quoted `` `Other Place` `` when the name + isn't identifier-shaped, including `Instance::Port` scoped names). +- `count ` lowers to `length`; `inf` to the `Infinity` constant. +- Newline-separated record fields avoid comma-noise in dynamics; both `,` and + newline are accepted separators. + +## Grammar (EBNF) + +``` +program := dynamics_body | expr +dynamics_body := "per" pattern "->" deriv_block +deriv_block := { deriv_binding } +deriv_binding := ident "'" "=" expr (NEWLINE | ";") +pattern := "{" ident { ("," | NEWLINE) ident } "}" | ident + +expr := "let" ident "=" expr "in" expr + | "if" expr "then" expr "else" expr + | pipe_expr +pipe_expr := or_expr { "|>" lambda_expr } (* distribution map *) +lambda_expr := ident "->" expr +or_expr := and_expr { "||" and_expr } +and_expr := cmp_expr { "&&" cmp_expr } +cmp_expr := add_expr [ ("<" | "<=" | ">" | ">=" | "==" | "!=") add_expr ] +add_expr := mul_expr { ("+" | "-") mul_expr } +mul_expr := unary_expr { ("*" | "/" | "%") unary_expr } +unary_expr := ("-" | "!") unary_expr | pow_expr +pow_expr := postfix_expr [ "^" unary_expr ] (* right assoc *) +postfix_expr := primary { "." ident | "[" expr "]" + | ".map" "(" lambda_expr ")" } +primary := NUMBER | "true" | "false" | "inf" | "pi" | "e" + | ident | quoted_ident + | "~" ident "(" args ")" (* distribution *) + | "count" primary + | math_fn "(" args ")" (* sin, cos, sqrt, … *) + | "(" expr ")" + | record | list +record := "{" [ field { ("," | NEWLINE) field } ] "}" +field := (ident | quoted_ident) "=" expr +list := "[" [ expr { "," expr } ] "]" +args := [ expr { "," expr } ] +``` + +Deliberate omissions: loops, recursion, user function definitions, strings, +mutation, sequencing. Comments: `(* … *)` and `# line comment`. + +## Implementation plan + +### Lexer + parser + +Hand-written: a lexer with precise token spans, and a recursive-descent parser +with a Pratt loop for binary operators (the grammar above is small enough that +the Pratt table is ~15 entries). Rationale over parser generators: + +- **Error recovery** is the whole game for an editor language. On error, emit + an `ErrorNode` carrying the span, synchronize on layout anchors (`let`, + `in`, `,`, `}`, `]`, newline in deriv blocks), and keep parsing — the + result is always a full tree with holes, never a hard failure. +- Documents are one expression, tens of lines — full reparse per keystroke is + microseconds; no incremental parsing needed (unlike tree-sitter, which + would still be a fine later swap for highlighting). + +### AST → HIR lowering + +Near 1:1 (`let`→`let`, `if`→`cond`, `|>`/`.map` on distributions → +`distributionMap`, `~D(…)`→`distribution`, `count`→`length`, records/lists → +`recordLit`/`arrayLit`, `per`-block → the `arrayMap`-over-tokens shape the +buffer-native emitter wants, with `x' = e` becoming record entry `x: e`). +`ErrorNode`s lower to a poisoned `unknown` node so downstream analyses still +run on the valid parts. Spans carry over directly — the DSL text _is_ the user +content. + +### Semantic analysis + +Reuse the HIR stack verbatim: `typecheckHir` + `analyzeHir` + `lintHirUserCode` +already operate on HIR with `SurfaceContext`, so the DSL gets every rule +(discrete derivatives, distribution-into-int, arc-weight bounds, shared +samples, …) for free. DSL-specific resolution happens pre-lowering: + +- bare identifiers resolve, in order: pattern/let/lambda bindings → input + place names → error with "did you mean" (Levenshtein over places/bindings); +- `params.x` checked against the parameter list at resolution time (same + diagnostics as `hir:unknown-parameter`). + +Type checking is bidirectional at the root: the surface fixes the expected +result type (derivative record / bool-or-rate / output record), which flows +into record literals and lists — this yields better messages than pure +inference ("this should be a token list for place `Target`" instead of a +mismatch at the leaf). + +### LSP integration + +The existing infrastructure was built language-agnostic in the right places — +transport, JSON-RPC protocol, diagnostics-store client, Monaco sync components +are all reusable unchanged. What changes is inside the worker: + +1. **Document model.** DSL items don't need virtual TS files, prefixes or + offset adjustment: the checker runs on the raw document. The + `SDCPNLanguageServer` grows a per-item language tag + (`Transition.language?: "typescript" | "petrinaut"` in the schema), and + dispatches per item to the TS service or the DSL analyzer. +2. **Diagnostics.** DSL parse/resolve/type errors are `HirDiagnostic`s → + the existing `check-hir.ts` conversion → the same + `publishDiagnostics` fan-out. No protocol change (this is already how HIR + lints ship today). +3. **Completions.** Context + partial tree driven: after `params.` list + parameters (typed); after a place expression `.` list attributes from the + color; in record position inside a kernel output, list the missing output + places / missing attributes; keywords (`let`, `if`, `then`, `else`, `in`, + `per`, `count`, `~Gaussian|Uniform|Lognormal`) elsewhere. The partial tree + plus `SurfaceContext` answers "what is expected here" — no type-service + needed. +4. **Hover/signature help.** Hover shows the resolved kind (parameter with + type/default, place with color and arc weight, binding with inferred + `HirType`); signature help covers distribution constructors and math + functions from the same tables the HIR uses. +5. **Monaco.** Register a `petrinaut` language with a Monarch tokenizer (~40 + lines); the existing `CompletionSync`/`HoverSync`/`DiagnosticsSync` + providers just need registering for the new language id alongside + `"typescript"`. + +### Migration story (TS ⇄ DSL) + +Because both frontends meet at the HIR, migration is a pretty-printer: + +- **TS → DSL**: `lowerTypeScriptToHir` then print HIR as DSL. Any item the + TS frontend can lower (the analyzable subset — in practice the default + templates and most real models) converts automatically; items outside the + subset stay TS until rewritten. A one-shot "Convert to Petrinaut language" + action per item, with a model-wide bulk action, both no-ops on failure. +- **DSL → TS**: same in reverse (print HIR as the TS idiom) — useful as an + escape hatch and for trust-building diffing. +- Storage: per-item `language` field next to the code string; default stays + `"typescript"` until the DSL is on by default. + +## Open questions + +- **Numeric semantics**: keep JS doubles (bit-compatibility with existing + runs, trivial JS backend) vs. defining int as i64 (matches token format v2's + ±2^53 discussion, matters for WASM/GPU backends). Sketch assumes doubles. +- **Comprehension power**: metrics realistically need `sum`/`filter`; is + `sum by ` special syntax (analyzable, GPU-reducible) or a + general `reduce` (more expressive, harder to fuse)? Leaning special forms. +- **Time**: none of the surfaces currently receive `t`/`dt`; if dynamics ever + become time-dependent, `t` becomes an ambient like `params`. +- **Multi-token kernels with data-dependent counts** (produce N tokens where + N is computed): today's TS surface can't either (tuple types); needs an arc + semantics decision before syntax. +- **Naming**: place references by display name mirror the TS surface but are + rename-fragile; the DSL could resolve through stable ids under the hood + (documents store ids, editor renders names) — a real advantage over the TS + frontend worth deciding early. diff --git a/libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.test.ts b/libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.test.ts new file mode 100644 index 00000000000..b564fb81a3b --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.test.ts @@ -0,0 +1,264 @@ +import { describe, expect, it } from "vitest"; + +import { emitBufferKernelJs, emitBufferLambdaJs } from "./emit-buffer-js"; +import { + instantiateHirBufferKernel, + instantiateHirBufferLambda, +} from "./instantiate"; +import { lowerTypeScriptToHir } from "./lower-typescript"; + +import type { RuntimeDistribution } from "../simulation/authoring/user-code/distribution"; +import type { HirFunction } from "./hir"; +import type { HirKernelContext, HirLambdaContext } from "./surface-context"; + +function lower(code: string, surface: "lambda" | "kernel"): HirFunction { + const result = lowerTypeScriptToHir(code, surface); + if (!result.ok) { + throw new Error(result.diagnostics[0]?.message); + } + return result.fn; +} + +// Layout: two input arcs — Pool (weight 2, elements x/alive) then +// Fuel (weight 1, element level). Slots: [Pool0, Pool1, Fuel0]. +const poolSlot = { + name: "Pool", + colorId: "c1", + elements: [ + { name: "x", type: "real" as const }, + { name: "alive", type: "boolean" as const }, + ], + tokenCount: 2, + slotStart: 0, +}; +const fuelSlot = { + name: "Fuel", + colorId: "c2", + elements: [{ name: "level", type: "real" as const }], + tokenCount: 1, + slotStart: 2, +}; + +const lambdaContext: HirLambdaContext = { + surface: "lambda", + parameters: [{ name: "rate", type: "real" }], + inputPlaces: [ + { ...poolSlot, slotStart: undefined as never }, + { ...fuelSlot, slotStart: undefined as never }, + ].map(({ slotStart: _slotStart, ...binding }) => binding), + inputSlots: [poolSlot, fuelSlot], + lambdaType: "stochastic", +}; + +const outSlot = { + name: "Out", + colorId: "c3", + elements: [ + { name: "x", type: "real" as const }, + { name: "count", type: "integer" as const }, + { name: "flag", type: "boolean" as const }, + ], + tokenCount: 1, + slotStart: 0, +}; + +const kernelContext: HirKernelContext = { + surface: "kernel", + parameters: [{ name: "sigma", type: "real" }], + inputPlaces: lambdaContext.inputPlaces, + inputSlots: lambdaContext.inputSlots, + outputPlaces: [(({ slotStart: _slotStart, ...binding }) => binding)(outSlot)], + outputSlots: [outSlot], + stochasticity: true, +}; + +// A frame: Pool tokens at bases 0 and 2, Fuel token at base 10. +// Pool[0] = { x: 1.5, alive: 1 }, Pool[1] = { x: -2, alive: 0 }, +// Fuel[0] = { level: 7 }. +const tokenValues = new Float64Array(16); +tokenValues.set([1.5, 1], 0); +tokenValues.set([-2, 0], 2); +tokenValues.set([7], 10); +const slotBases = new Int32Array([0, 2, 10]); + +describe("emitBufferLambdaJs", () => { + function compileLambda(code: string, parameters = { rate: 3 }) { + const program = emitBufferLambdaJs(lower(code, "lambda"), lambdaContext); + expect(program).not.toBeNull(); + expect(program!.inputSlotCount).toBe(3); + return instantiateHirBufferLambda(program!.source, parameters); + } + + it("reads token attributes at static offsets", () => { + const fn = compileLambda( + `export default Lambda((input, parameters) => input.Pool[0].x + input.Pool[1].x + input.Fuel[0].level);`, + ); + expect(fn(tokenValues, slotBases)).toBe(1.5 - 2 + 7); + }); + + it("decodes booleans and binds parameters", () => { + const fn = compileLambda( + `export default Lambda((input, parameters) => input.Pool[0].alive && input.Pool[0].x * parameters.rate > 4);`, + ); + expect(fn(tokenValues, slotBases)).toBe(true); + }); + + it("supports destructured bindings and guard clauses", () => { + const fn = compileLambda( + `export default Lambda((input, parameters) => { + const { x, alive } = input.Pool[0]; + if (!alive) return 0; + const { rate } = parameters; + return x * rate; +});`, + ); + expect(fn(tokenValues, slotBases)).toBe(4.5); + }); + + it("resolves .length to the static arc weight", () => { + const fn = compileLambda( + `export default Lambda((input) => input.Pool.length * 10 + input.Fuel.length);`, + ); + expect(fn(tokenValues, slotBases)).toBe(21); + }); + + it("supports array destructuring of the tuple", () => { + const fn = compileLambda( + `export default Lambda((input) => { + const [a, b] = input.Pool; + return a.x - b.x; +});`, + ); + expect(fn(tokenValues, slotBases)).toBe(3.5); + }); + + it("bails to null on dynamic token indices", () => { + const fn = lower( + `export default Lambda((input, parameters) => input.Pool[parameters.rate].x);`, + "lambda", + ); + expect(emitBufferLambdaJs(fn, lambdaContext)).toBeNull(); + }); +}); + +describe("emitBufferKernelJs", () => { + function compileKernel(code: string, parameters = { sigma: 2 }) { + const program = emitBufferKernelJs(lower(code, "kernel"), kernelContext); + expect(program).not.toBeNull(); + expect(program!.inputSlotCount).toBe(3); + expect(program!.outputFloatCount).toBe(3); + return instantiateHirBufferKernel(program!.source, parameters); + } + + function runKernel(fn: ReturnType) { + const out = new Float64Array(3); + const sinks: [number, RuntimeDistribution][] = []; + fn(tokenValues, slotBases, out, (index, dist) => sinks.push([index, dist])); + return { out, sinks }; + } + + it("writes attributes place-major with integer rounding and boolean 0/1", () => { + const fn = compileKernel( + `export default TransitionKernel((input, parameters) => ({ + Out: [{ x: input.Pool[0].x * 2, count: 2.7, flag: input.Pool[1].x < 0 }], +}));`, + ); + const { out, sinks } = runKernel(fn); + expect([...out]).toEqual([3, 3, 1]); + expect(sinks).toEqual([]); + }); + + it("defers distributions through distSink with shared identity", () => { + const fn = compileKernel( + `export default TransitionKernel((input, parameters) => { + const noise = Distribution.Gaussian(0, parameters.sigma); + const scaled = noise.map((v) => v * 10); + return { Out: [{ x: scaled, count: 1, flag: false }] }; +});`, + ); + const { out, sinks } = runKernel(fn); + expect(out[1]).toBe(1); + expect(sinks).toHaveLength(1); + const [index, dist] = sinks[0]!; + expect(index).toBe(0); + expect(dist.type).toBe("mapped"); + if (dist.type === "mapped") { + expect(dist.inner).toMatchObject({ + type: "gaussian", + mean: 0, + deviation: 2, + }); + expect(dist.fn(3)).toBe(30); + } + }); + + it("unrolls .map over input tuples", () => { + const context: HirKernelContext = { + ...kernelContext, + outputSlots: [{ ...outSlot, tokenCount: 2 }], + outputPlaces: [ + (({ slotStart: _s, ...binding }) => ({ ...binding, tokenCount: 2 }))( + outSlot, + ), + ], + }; + const program = emitBufferKernelJs( + lower( + `export default TransitionKernel((input, parameters) => ({ + Out: input.Pool.map((token, i) => ({ x: token.x + i, count: i, flag: token.alive })), +}));`, + "kernel", + ), + context, + ); + expect(program).not.toBeNull(); + const fn = instantiateHirBufferKernel(program!.source, {}); + const out = new Float64Array(6); + fn(tokenValues, slotBases, out, () => {}); + expect([...out]).toEqual([1.5, 0, 1, -1, 1, 0]); + }); + + it("forwards whole input tuples to outputs", () => { + const context: HirKernelContext = { + ...kernelContext, + outputSlots: [ + { + name: "Sink", + colorId: "c2", + elements: [{ name: "level", type: "real" as const }], + tokenCount: 1, + slotStart: 0, + }, + ], + outputPlaces: [ + { + name: "Sink", + colorId: "c2", + elements: [{ name: "level", type: "real" as const }], + tokenCount: 1, + }, + ], + }; + const program = emitBufferKernelJs( + lower( + `export default TransitionKernel((input) => ({ Sink: input.Fuel }));`, + "kernel", + ), + context, + ); + expect(program).not.toBeNull(); + const fn = instantiateHirBufferKernel(program!.source, {}); + const out = new Float64Array(1); + fn(tokenValues, slotBases, out, () => {}); + expect([...out]).toEqual([7]); + }); + + it("bails when the output token count is not statically resolvable", () => { + // Returning a conditional record — structurally dynamic. + const fn = lower( + `export default TransitionKernel((input, parameters) => parameters.sigma > 1 ? { Out: [{ x: 1, count: 0, flag: false }] } : { Out: [{ x: 2, count: 0, flag: false }] });`, + "kernel", + ); + expect(emitBufferKernelJs(fn, kernelContext)).toBeNull(); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.ts b/libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.ts new file mode 100644 index 00000000000..b84f8379b4c --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.ts @@ -0,0 +1,585 @@ +/** + * Buffer-ABI JavaScript emission for lambdas and transition kernels. + * + * Instead of the legacy object convention (decode packed tokens into + * `{ Place: [{ x, y }] }` records per combination, call, re-encode), the + * emitted functions read token attributes straight out of the engine's packed + * `Float64Array` at statically-resolved offsets: + * + * lambda: (tokenValues, slotBases) => number | boolean + * kernel: (tokenValues, slotBases, out, distSink) => void + * + * - `tokenValues` — the frame's token-value region (floats). + * - `slotBases` — one float base offset per input token slot; slots are laid + * out per colored non-inhibitor input arc, `weight` slots each, in arc + * order (see the slot layout invariant in `surface-context.ts`). The + * engine fills this array per enumerated combination — no allocation. + * - `out` — kernel staging: colored output arcs in arc order, each occupying + * `weight * elements.length` floats (integers pre-rounded, booleans 0/1). + * - `distSink(floatIndex, distribution)` — deferred sampling: the engine + * samples after the call, ordered by float index, which reproduces the + * legacy (place, token, element) sampling order and therefore the exact + * RNG stream. + * + * Emission works by symbolic evaluation: structural values (input tuples, + * tokens, records, arrays) exist only at compile time; everything that + * reaches the output is a scalar JS expression or a distribution constant. + * `.map(...)` over token tuples is unrolled (arc weights are static). + * Shapes the evaluator cannot scalarize return `null` — callers fall back to + * the object-convention emitter (`emit-js.ts`). + * + * Only emit from functions whose typecheck produced no errors: the emitter + * relies on attribute/place validity established by `typecheckHir`. + */ +import { foldHir } from "./analyze"; + +import type { HirExpr, HirFunction } from "./hir"; +import type { + HirArcSlot, + HirKernelContext, + HirLambdaContext, + HirTokenElementInfo, +} from "./surface-context"; + +export type BufferProgram = { + /** JS source of the emitted function (arrow expression). */ + source: string; + /** Expected `slotBases.length` — engine-side sanity check. */ + inputSlotCount: number; +}; + +export type BufferKernelProgram = BufferProgram & { + /** Expected `out.length` — engine-side sanity check. */ + outputFloatCount: number; +}; + +/** Maximum tokens per arc slot we are willing to unroll `.map` over. */ +const MAX_UNROLL = 16; + +const RESERVED_NAMES = [ + "tokenValues", + "slotBases", + "out", + "distSink", + "__params", + "__dist", + "Math", + "Infinity", + "NaN", +]; + +/** Internal bail signal: shape outside the buffer-ABI subset. */ +class BailError extends Error {} + +type Value = + /** A JS expression producing a number or boolean. */ + | { kind: "scalar"; code: string } + /** A JS expression producing a RuntimeDistribution (a hoisted const). */ + | { kind: "dist"; code: string } + /** The lambda/kernel first parameter (`tokensByPlace`). */ + | { kind: "inputRecord" } + /** A whole input arc slot's token tuple. */ + | { kind: "tokens"; slot: HirArcSlot } + /** One token; `baseCode` is a JS expression for its float base offset. */ + | { kind: "token"; baseCode: string; elements: HirTokenElementInfo[] } + | { kind: "record"; fields: Map } + | { kind: "array"; items: Value[] }; + +function quoteKey(key: string): string { + return JSON.stringify(key); +} + +function emitNumber(value: number, raw: string): string { + if (Number.isNaN(value)) { + return "NaN"; + } + if (Number(raw) === value) { + return raw; + } + if (value === Infinity) { + return "Infinity"; + } + if (value === -Infinity) { + return "-Infinity"; + } + return String(value); +} + +class NameAllocator { + private readonly used: Set; + + constructor(used: Iterable) { + this.used = new Set(used); + } + + allocate(preferred: string): string { + let name = preferred; + let suffix = 2; + while (this.used.has(name)) { + name = `${preferred}_${suffix}`; + suffix += 1; + } + this.used.add(name); + return name; + } +} + +class BufferEmitter { + /** Hoisted `const` statements, in evaluation order. */ + readonly lines: string[] = []; + readonly names = new NameAllocator(RESERVED_NAMES); + + constructor(private readonly inputSlots: HirArcSlot[]) {} + + /** Resolves a place display name to its arc slot (last arc wins, matching + * the runtime's object-key overwrite). */ + private slotForPlace(name: string): HirArcSlot | null { + for (let index = this.inputSlots.length - 1; index >= 0; index -= 1) { + if (this.inputSlots[index]!.name === name) { + return this.inputSlots[index]!; + } + } + return null; + } + + private tokenOfSlot(slot: HirArcSlot, tokenIndex: number): Value { + return { + kind: "token", + baseCode: `slotBases[${slot.slotStart + tokenIndex}]`, + elements: slot.elements, + }; + } + + private readAttribute( + token: Extract, + field: string, + ): Value { + const index = token.elements.findIndex((element) => element.name === field); + if (index === -1) { + throw new BailError(); + } + const read = `tokenValues[${token.baseCode} + ${index}]`; + return { + kind: "scalar", + code: + token.elements[index]!.type === "boolean" ? `(${read} !== 0)` : read, + }; + } + + private scalar(value: Value): string { + if (value.kind !== "scalar") { + throw new BailError(); + } + return value.code; + } + + eval(expr: HirExpr, env: Map): Value { + switch (expr.kind) { + case "numberLit": + return { kind: "scalar", code: emitNumber(expr.value, expr.raw) }; + case "boolLit": + return { kind: "scalar", code: expr.value ? "true" : "false" }; + case "constant": + switch (expr.name) { + case "PI": + return { kind: "scalar", code: "Math.PI" }; + case "E": + return { kind: "scalar", code: "Math.E" }; + case "Infinity": + return { kind: "scalar", code: "Infinity" }; + case "NaN": + return { kind: "scalar", code: "NaN" }; + } + break; + case "localRef": { + const value = env.get(expr.name); + if (!value) { + throw new BailError(); + } + return value; + } + case "paramRef": + return { kind: "scalar", code: `__params[${quoteKey(expr.name)}]` }; + case "fieldAccess": { + const target = this.eval(expr.target, env); + if (target.kind === "inputRecord") { + const slot = this.slotForPlace(expr.field); + if (!slot) { + throw new BailError(); + } + return { kind: "tokens", slot }; + } + if (target.kind === "token") { + return this.readAttribute(target, expr.field); + } + if (target.kind === "record") { + const field = target.fields.get(expr.field); + if (!field) { + throw new BailError(); + } + return field; + } + throw new BailError(); + } + case "indexAccess": { + const target = this.eval(expr.target, env); + const index = foldHir(expr.index); + if (index.kind !== "numberLit" || !Number.isInteger(index.value)) { + // Dynamic token indices would read unchecked memory — leave those + // to the object-convention fallback. + throw new BailError(); + } + if (target.kind === "tokens") { + if (index.value < 0 || index.value >= target.slot.tokenCount) { + throw new BailError(); + } + return this.tokenOfSlot(target.slot, index.value); + } + if (target.kind === "array") { + const item = target.items[index.value]; + if (!item) { + throw new BailError(); + } + return item; + } + throw new BailError(); + } + case "length": { + const target = this.eval(expr.target, env); + if (target.kind === "tokens") { + return { kind: "scalar", code: String(target.slot.tokenCount) }; + } + if (target.kind === "array") { + return { kind: "scalar", code: String(target.items.length) }; + } + throw new BailError(); + } + case "unary": + return { + kind: "scalar", + code: `(${expr.op}${this.scalar(this.eval(expr.operand, env))})`, + }; + case "binary": { + const op = + expr.op === "==" ? "===" : expr.op === "!=" ? "!==" : expr.op; + return { + kind: "scalar", + code: `(${this.scalar(this.eval(expr.left, env))} ${op} ${this.scalar( + this.eval(expr.right, env), + )})`, + }; + } + case "cond": { + const condition = this.scalar(this.eval(expr.condition, env)); + const thenValue = this.eval(expr.thenBranch, env); + const elseValue = this.eval(expr.elseBranch, env); + if (thenValue.kind === "scalar" && elseValue.kind === "scalar") { + return { + kind: "scalar", + code: `(${condition} ? ${thenValue.code} : ${elseValue.code})`, + }; + } + if (thenValue.kind === "dist" && elseValue.kind === "dist") { + return { + kind: "dist", + code: `(${condition} ? ${thenValue.code} : ${elseValue.code})`, + }; + } + // Structural conditionals (records/tuples) stay on the object path. + throw new BailError(); + } + case "let": { + const scoped = new Map(env); + for (const binding of expr.bindings) { + const value = this.eval(binding.value, scoped); + scoped.set(binding.name, this.hoist(binding.name, value)); + } + return this.eval(expr.body, scoped); + } + case "mathCall": { + const args = expr.args.map((argument) => + this.scalar(this.eval(argument, env)), + ); + return { kind: "scalar", code: `Math.${expr.fn}(${args.join(", ")})` }; + } + case "recordLit": { + const fields = new Map(); + for (const entry of expr.entries) { + fields.set(entry.key, this.eval(entry.value, env)); + } + return { kind: "record", fields }; + } + case "arrayLit": + return { + kind: "array", + items: expr.elements.map((element) => this.eval(element, env)), + }; + case "arrayMap": { + const target = this.eval(expr.target, env); + let items: Value[]; + if (target.kind === "tokens") { + if (target.slot.tokenCount > MAX_UNROLL) { + throw new BailError(); + } + items = Array.from({ length: target.slot.tokenCount }, (_, index) => + this.tokenOfSlot(target.slot, index), + ); + } else if (target.kind === "array") { + if (target.items.length > MAX_UNROLL) { + throw new BailError(); + } + items = target.items; + } else { + throw new BailError(); + } + return { + kind: "array", + items: items.map((item, index) => { + const scoped = new Map(env); + scoped.set(expr.param.name, item); + if (expr.indexParam) { + scoped.set(expr.indexParam.name, { + kind: "scalar", + code: String(index), + }); + } + return this.eval(expr.body, scoped); + }), + }; + } + case "distribution": { + const args = expr.args.map((argument) => + this.scalar(this.eval(argument, env)), + ); + // Always hoisted to a const: object identity is what makes several + // sinks share one draw at sampling time. + const name = this.names.allocate("__d"); + this.lines.push( + `const ${name} = __dist.${expr.dist}(${args.join(", ")});`, + ); + return { kind: "dist", code: name }; + } + case "distributionMap": { + const base = this.eval(expr.base, env); + if (base.kind !== "dist") { + throw new BailError(); + } + const paramName = this.names.allocate(expr.param.name); + const scoped = new Map(env); + scoped.set(expr.param.name, { kind: "scalar", code: paramName }); + const body = this.evalCallbackBody(expr.body, scoped); + const name = this.names.allocate("__d"); + this.lines.push( + `const ${name} = __dist.map(${base.code}, (${paramName}) => ${body});`, + ); + return { kind: "dist", code: name }; + } + } + throw new BailError(); + } + + /** + * Evaluates a distribution-map callback body to a self-contained JS body + * (its `let` bindings must stay inside the callback, not hoisted). + */ + private evalCallbackBody(expr: HirExpr, env: Map): string { + if (expr.kind !== "let") { + return `(${this.scalar(this.eval(expr, env))})`; + } + const statements: string[] = []; + const scoped = new Map(env); + for (const binding of expr.bindings) { + const value = this.eval(binding.value, scoped); + if (value.kind !== "scalar") { + throw new BailError(); + } + const name = this.names.allocate(binding.name); + statements.push(`const ${name} = ${value.code};`); + scoped.set(binding.name, { kind: "scalar", code: name }); + } + statements.push(`return ${this.scalar(this.eval(expr.body, scoped))};`); + return `{ ${statements.join(" ")} }`; + } + + /** Binds a `const`: scalars and distributions become hoisted consts (so + * they evaluate once); structural values stay symbolic. */ + private hoist(name: string, value: Value): Value { + if (value.kind === "scalar") { + const jsName = this.names.allocate(name); + this.lines.push(`const ${jsName} = ${value.code};`); + return { kind: "scalar", code: jsName }; + } + if (value.kind === "dist") { + // Distribution constructions are already hoisted consts; a `cond` of + // two dists is cheap to re-evaluate but must keep identity — hoist it. + if (!/^[A-Za-z_$][A-Za-z0-9_$]*$/.test(value.code)) { + const jsName = this.names.allocate(name); + this.lines.push(`const ${jsName} = ${value.code};`); + return { kind: "dist", code: jsName }; + } + return value; + } + return value; + } +} + +function initialEnv(fn: HirFunction): Map { + const env = new Map(); + const tokensParam = fn.params[0]; + if (tokensParam) { + env.set(tokensParam.name, { kind: "inputRecord" }); + } + return env; +} + +function slotCount(inputSlots: HirArcSlot[]): number { + return inputSlots.reduce((sum, slot) => sum + slot.tokenCount, 0); +} + +/** + * Emits a buffer-ABI lambda: `(tokenValues, slotBases) => number | boolean`. + * Reads `__params` for model parameters (bound at instantiation). Returns + * `null` when the function shape needs the object-convention fallback. + */ +export function emitBufferLambdaJs( + fn: HirFunction, + context: HirLambdaContext, +): BufferProgram | null { + try { + const emitter = new BufferEmitter(context.inputSlots); + const result = emitter.eval(foldHir(fn.body), initialEnv(fn)); + if (result.kind !== "scalar") { + return null; + } + const source = [ + `(tokenValues, slotBases) => {`, + ...emitter.lines.map((line) => ` ${line}`), + ` return ${result.code};`, + `}`, + ].join("\n"); + return { source, inputSlotCount: slotCount(context.inputSlots) }; + } catch (error) { + if (error instanceof BailError) { + return null; + } + throw error; + } +} + +/** + * Emits a buffer-ABI kernel: + * `(tokenValues, slotBases, out, distSink) => void`. + * + * Writes every output attribute into `out` (place-major, arc order): + * integers pre-rounded, booleans as 0/1, distribution values deferred via + * `distSink`. Returns `null` when the output shape is not statically + * resolvable (the object-convention fallback handles it). + */ +export function emitBufferKernelJs( + fn: HirFunction, + context: HirKernelContext, +): BufferKernelProgram | null { + try { + const emitter = new BufferEmitter(context.inputSlots); + const result = emitter.eval(foldHir(fn.body), initialEnv(fn)); + if (result.kind !== "record") { + return null; + } + + const writes: string[] = []; + let floatBase = 0; + + for (const slot of context.outputSlots) { + const dims = slot.elements.length; + const entry = result.fields.get(slot.name); + if (!entry) { + return null; + } + + // Resolve the entry to one token Value per output slot position. + let tokens: Value[]; + if (entry.kind === "array") { + tokens = entry.items; + } else if (entry.kind === "tokens") { + // Whole input tuple forwarded to an output place. + if (entry.slot.tokenCount !== slot.tokenCount) { + return null; + } + tokens = Array.from({ length: entry.slot.tokenCount }, (_, index) => ({ + kind: "token" as const, + baseCode: `slotBases[${entry.slot.slotStart + index}]`, + elements: entry.slot.elements, + })); + } else { + return null; + } + if (tokens.length !== slot.tokenCount) { + return null; + } + + for (const [tokenIndex, token] of tokens.entries()) { + for (const [elementIndex, element] of slot.elements.entries()) { + const floatIndex = floatBase + tokenIndex * dims + elementIndex; + + let value: Value; + if (token.kind === "record") { + const field = token.fields.get(element.name); + if (!field) { + return null; + } + value = field; + } else if (token.kind === "token") { + const index = token.elements.findIndex( + (candidate) => candidate.name === element.name, + ); + if (index === -1) { + return null; + } + value = { + kind: "scalar", + code: `tokenValues[${token.baseCode} + ${index}]`, + }; + } else { + return null; + } + + if (value.kind === "scalar") { + const encoded = + element.type === "integer" + ? `Math.round(${value.code})` + : element.type === "boolean" + ? `(${value.code}) ? 1 : 0` + : value.code; + writes.push(`out[${floatIndex}] = ${encoded};`); + } else if (value.kind === "dist") { + if (element.type !== "real") { + return null; + } + writes.push(`distSink(${floatIndex}, ${value.code});`); + } else { + return null; + } + } + } + + floatBase += slot.tokenCount * dims; + } + + const source = [ + `(tokenValues, slotBases, out, distSink) => {`, + ...emitter.lines.map((line) => ` ${line}`), + ...writes.map((line) => ` ${line}`), + `}`, + ].join("\n"); + return { + source, + inputSlotCount: slotCount(context.inputSlots), + outputFloatCount: floatBase, + }; + } catch (error) { + if (error instanceof BailError) { + return null; + } + throw error; + } +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/emit-js.test.ts b/libs/@hashintel/petrinaut-core/src/hir/emit-js.test.ts new file mode 100644 index 00000000000..1e3c13554f0 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/emit-js.test.ts @@ -0,0 +1,194 @@ +import { describe, expect, it } from "vitest"; + +import { + tryCompileHirBufferDynamics, + tryCompileHirKernel, + tryCompileHirLambda, +} from "./compile"; + +import type { RuntimeDistribution } from "../simulation/authoring/user-code/distribution"; + +describe("tryCompileHirLambda", () => { + it("compiles lambdas with bindings, length and conditionals", () => { + const code = `export default Lambda((input, parameters) => { + const pressure = input.Pool.length * parameters.rate; + return pressure > 10 ? Infinity : pressure; +});`; + const hirFn = tryCompileHirLambda(code)!; + expect(hirFn).not.toBeNull(); + + const tokensByPlace = { Pool: [{ x: 1 }, { x: 2 }, { x: 3 }] }; + expect(hirFn(tokensByPlace, { rate: 0.5 })).toBe(1.5); + expect(hirFn(tokensByPlace, { rate: 2 })).toBe(6); + expect(hirFn(tokensByPlace, { rate: 5 })).toBe(Infinity); + }); + + it("supports predicates over token attributes", () => { + const fn = tryCompileHirLambda( + `export default Lambda((input, parameters) => input.Pool[0].active && input.Pool[0].x >= parameters.threshold);`, + )!; + expect(fn({ Pool: [{ active: true, x: 5 }] }, { threshold: 4 })).toBe(true); + expect(fn({ Pool: [{ active: true, x: 3 }] }, { threshold: 4 })).toBe( + false, + ); + }); + + it("returns null for out-of-subset code", () => { + expect( + tryCompileHirLambda( + `export default Lambda((input) => { let x = 1; return x; });`, + ), + ).toBeNull(); + }); + + it("renames user bindings that collide with emitter internals", () => { + const fn = tryCompileHirLambda( + `export default Lambda((input, parameters) => { + const __params = parameters.rate * 2; + return __params + 1; +});`, + )!; + expect(fn({}, { rate: 3 })).toBe(7); + }); +}); + +describe("tryCompileHirKernel", () => { + it("produces runtime distribution objects compatible with the engine", () => { + const fn = tryCompileHirKernel( + `export default TransitionKernel((input, parameters) => { + const noise = Distribution.Gaussian(0, parameters.sigma); + return { Out: [{ x: noise.map((value) => value * 2), y: 1 }] }; +});`, + )!; + const output = fn({}, { sigma: 3 }) as Record< + string, + Record[] + >; + const x = output.Out![0]!.x as RuntimeDistribution; + expect(x.__brand).toBe("distribution"); + expect(x.type).toBe("mapped"); + if (x.type === "mapped") { + expect(x.inner).toMatchObject({ + type: "gaussian", + mean: 0, + deviation: 3, + }); + expect(x.fn(2)).toBe(4); + } + expect(output.Out![0]!.y).toBe(1); + }); + + it("shares one distribution object across aliased outputs", () => { + const fn = tryCompileHirKernel( + `export default TransitionKernel((input) => { + const d = Distribution.Uniform(0, 1); + return { Out: [{ x: d, y: d }] }; +});`, + )!; + const output = fn({}, {}) as Record[]>; + // Same object identity → the engine's sample cache yields one draw. + expect(output.Out![0]!.x).toBe(output.Out![0]!.y); + }); + + it("passes input tokens through .map kernels", () => { + const fn = tryCompileHirKernel( + `export default TransitionKernel((input, parameters) => ({ + Out: input.In.map((token, index) => ({ x: token.x + index, y: token.y * parameters.k })), +}));`, + )!; + const output = fn( + { + In: [ + { x: 1, y: 2 }, + { x: 10, y: 20 }, + ], + }, + { k: 3 }, + ) as Record; + expect(output.Out).toEqual([ + { x: 1, y: 6 }, + { x: 11, y: 60 }, + ]); + }); +}); + +describe("tryCompileHirBufferDynamics", () => { + const elements = [ + { name: "x", type: "real" as const }, + { name: "v", type: "real" as const }, + { name: "alive", type: "boolean" as const }, + ]; + + it("computes derivatives directly on the packed buffer", () => { + const fn = tryCompileHirBufferDynamics( + `export default Dynamics((tokens, parameters) => { + return tokens.map(({ x, v }) => { + return { x: v, v: -parameters.k * x }; + }); +});`, + elements, + { k: 2 }, + )!; + expect(fn).not.toBeNull(); + + // Two tokens: (x=1, v=10, alive=1), (x=-3, v=0, alive=0) + const state = new Float64Array([1, 10, 1, -3, 0, 0]); + const result = fn(state, 3, 2); + expect([...result]).toEqual([10, -2, 0, 0, 6, 0]); + }); + + it("decodes booleans and evaluates conditionals per token", () => { + const code = `export default Dynamics((tokens, parameters) => { + return tokens.map(({ x, v, alive }) => { + return { x: alive ? v * parameters.g : 0, v: Math.cos(x) }; + }); +});`; + const hirFn = tryCompileHirBufferDynamics(code, elements, { g: 9.81 })!; + expect(hirFn).not.toBeNull(); + + const state = new Float64Array([0.5, 2, 1, 1.5, -4, 0]); + const hirResult = hirFn(state, 3, 2); + + expect([...hirResult]).toEqual([ + 2 * 9.81, // token 0 alive: v * g + Math.cos(0.5), + 0, // discrete attribute derivative is always 0 + 0, // token 1 not alive + Math.cos(1.5), + 0, + ]); + }); + + it("supports index params, token counts and cross-token reads", () => { + const fn = tryCompileHirBufferDynamics( + `export default Dynamics((tokens) => tokens.map((token, i) => ({ + x: tokens[0].x + i * tokens.length, +})));`, + [{ name: "x", type: "real" as const }], + {}, + )!; + const state = new Float64Array([7, 100, 200]); + expect([...fn(state, 1, 3)]).toEqual([7, 10, 13]); + }); + + it("returns null when the body is not a token map", () => { + expect( + tryCompileHirBufferDynamics( + `export default Dynamics((tokens) => [{ x: 1 }]);`, + [{ name: "x", type: "real" as const }], + {}, + ), + ).toBeNull(); + }); + + it("throws on dimension mismatch like the legacy wrapper", () => { + const fn = tryCompileHirBufferDynamics( + `export default Dynamics((tokens) => tokens.map(() => ({ x: 1 })));`, + [{ name: "x", type: "real" as const }], + {}, + )!; + expect(() => fn(new Float64Array([0, 0]), 2, 1)).toThrow( + /Expected 1 dimensions/, + ); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/emit-js.ts b/libs/@hashintel/petrinaut-core/src/hir/emit-js.ts new file mode 100644 index 00000000000..1a21820b7d2 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/emit-js.ts @@ -0,0 +1,475 @@ +/** + * Compiles HIR functions to JavaScript source. + * + * Two emitters live here: + * + * - A generic expression emitter producing functions with the same signature + * as the legacy Babel-compiled user code (`(tokensByPlace, parameters) => + * result`), with distributions constructed through an injected `__dist` + * runtime instead of source-string injection. + * + * - A buffer-native dynamics emitter that compiles `tokens.map((token) => + * ({ ... }))` bodies straight to a `Float64Array` loop with no per-token + * object allocation — matching the engine's `DifferentialEquationFn` + * signature directly. + * + * Emission never throws on well-formed HIR; the buffer-native emitter returns + * `null` when the function shape doesn't fit the fast path so callers can fall + * back. + */ +import { foldHir } from "./analyze"; + +import type { HirExpr, HirFunction } from "./hir"; +import type { HirTokenElementInfo } from "./surface-context"; + +/** Names the emitted code relies on; user locals are renamed away from them. */ +const RESERVED_NAMES = [ + "__dist", + "__params", + "currentState", + "dimensions", + "numberOfTokens", + "result", + "Math", + "Infinity", + "NaN", +]; + +class NameAllocator { + private readonly used: Set; + + constructor(used: Iterable) { + this.used = new Set(used); + } + + child(): NameAllocator { + return new NameAllocator(this.used); + } + + allocate(preferred: string): string { + let name = preferred; + let suffix = 2; + while (this.used.has(name)) { + name = `${preferred}_${suffix}`; + suffix += 1; + } + this.used.add(name); + return name; + } +} + +type EmitScope = { + /** HIR local name → emitted JS name. */ + names: Map; + allocator: NameAllocator; +}; + +function childScope(scope: EmitScope): EmitScope { + return { names: new Map(scope.names), allocator: scope.allocator.child() }; +} + +function quoteKey(key: string): string { + return JSON.stringify(key); +} + +function emitNumber(value: number, raw: string): string { + // Preserve the exact source spelling when it still parses to the same + // value; otherwise round-trip through String (lossless for doubles). + if ( + Number(raw) === value || + (Number.isNaN(Number(raw)) && Number.isNaN(value)) + ) { + return Number.isNaN(Number(raw)) ? String(value) : raw; + } + if (value === Infinity) { + return "Infinity"; + } + if (value === -Infinity) { + return "-Infinity"; + } + return String(value); +} + +function emitExpr(expr: HirExpr, scope: EmitScope): string { + switch (expr.kind) { + case "numberLit": + return emitNumber(expr.value, expr.raw); + case "boolLit": + return expr.value ? "true" : "false"; + case "constant": + switch (expr.name) { + case "PI": + return "Math.PI"; + case "E": + return "Math.E"; + case "Infinity": + return "Infinity"; + case "NaN": + return "NaN"; + } + break; + case "localRef": + return scope.names.get(expr.name) ?? expr.name; + case "paramRef": + return `__params[${quoteKey(expr.name)}]`; + case "fieldAccess": + return `${emitExpr(expr.target, scope)}[${quoteKey(expr.field)}]`; + case "indexAccess": + return `${emitExpr(expr.target, scope)}[${emitExpr(expr.index, scope)}]`; + case "length": + return `${emitExpr(expr.target, scope)}.length`; + case "unary": + return `(${expr.op}${emitExpr(expr.operand, scope)})`; + case "binary": { + const op = expr.op === "==" ? "===" : expr.op === "!=" ? "!==" : expr.op; + return `(${emitExpr(expr.left, scope)} ${op} ${emitExpr(expr.right, scope)})`; + } + case "cond": + return `(${emitExpr(expr.condition, scope)} ? ${emitExpr(expr.thenBranch, scope)} : ${emitExpr(expr.elseBranch, scope)})`; + case "let": + // `let` only appears as a function/callback body; those emit blocks. + // eslint-disable-next-line no-use-before-define -- mutual recursion + return `(() => ${emitBody(expr, scope)})()`; + case "mathCall": + return `Math.${expr.fn}(${expr.args + .map((argument) => emitExpr(argument, scope)) + .join(", ")})`; + case "recordLit": + return `{ ${expr.entries + .map( + (entry) => `${quoteKey(entry.key)}: ${emitExpr(entry.value, scope)}`, + ) + .join(", ")} }`; + case "arrayLit": + return `[${expr.elements + .map((element) => emitExpr(element, scope)) + .join(", ")}]`; + case "arrayMap": { + const bodyScope = childScope(scope); + const paramName = bodyScope.allocator.allocate(expr.param.name); + bodyScope.names.set(expr.param.name, paramName); + let params = paramName; + if (expr.indexParam) { + const indexName = bodyScope.allocator.allocate(expr.indexParam.name); + bodyScope.names.set(expr.indexParam.name, indexName); + params = `${paramName}, ${indexName}`; + } + // eslint-disable-next-line no-use-before-define -- mutual recursion + return `${emitExpr(expr.target, scope)}.map((${params}) => ${emitBody(expr.body, bodyScope)})`; + } + case "distribution": + return `__dist.${expr.dist}(${expr.args + .map((argument) => emitExpr(argument, scope)) + .join(", ")})`; + case "distributionMap": { + const bodyScope = childScope(scope); + const paramName = bodyScope.allocator.allocate(expr.param.name); + bodyScope.names.set(expr.param.name, paramName); + // eslint-disable-next-line no-use-before-define -- mutual recursion + return `__dist.map(${emitExpr(expr.base, scope)}, (${paramName}) => ${emitBody(expr.body, bodyScope)})`; + } + } + throw new Error("Unreachable HIR node in emitExpr"); +} + +/** Emits a callback/function body: a block for `let`, an expression otherwise. */ +function emitBody(expr: HirExpr, scope: EmitScope): string { + if (expr.kind !== "let") { + // Parenthesize object literals so they aren't parsed as blocks. + const emitted = emitExpr(expr, scope); + return expr.kind === "recordLit" ? `(${emitted})` : emitted; + } + const bodyScope = childScope(scope); + const statements: string[] = []; + for (const binding of expr.bindings) { + const value = emitExpr(binding.value, bodyScope); + const name = bodyScope.allocator.allocate(binding.name); + bodyScope.names.set(binding.name, name); + statements.push(`const ${name} = ${value};`); + } + statements.push(`return ${emitExpr(expr.body, bodyScope)};`); + return `{ ${statements.join(" ")} }`; +} + +/** + * Emits a user-function-shaped JavaScript expression: + * `(tokens, parameters) => result`. + * + * The emitted code may reference `__dist` (distribution runtime) and + * `__params` (the parameters object); callers bind both when instantiating. + * The declared `parameters` parameter is accepted for signature compatibility + * but reads go through `__params` — instantiators alias them. + */ +export function emitUserFunctionJs(fn: HirFunction): string { + const folded = foldHir(fn.body); + const allocator = new NameAllocator(RESERVED_NAMES); + const scope: EmitScope = { names: new Map(), allocator }; + + const paramNames = fn.params.map((parameter) => + allocator.allocate(parameter.name), + ); + for (const [index, parameter] of fn.params.entries()) { + scope.names.set(parameter.name, paramNames[index]!); + } + + // The second user parameter *is* the parameters object: alias __params to + // it so emitted `__params[...]` reads resolve to the call argument. + const signature = paramNames.join(", "); + const statements: string[] = []; + if (fn.params.length > 1) { + statements.push(` const __params = ${paramNames[1]!};`); + } + + if (folded.kind === "let") { + const bodyScope = childScope(scope); + for (const binding of folded.bindings) { + const value = emitExpr(binding.value, bodyScope); + const name = bodyScope.allocator.allocate(binding.name); + bodyScope.names.set(binding.name, name); + statements.push(` const ${name} = ${value};`); + } + statements.push(` return ${emitExpr(folded.body, bodyScope)};`); + return [`(${signature}) => {`, ...statements, `}`].join("\n"); + } + + statements.push(` return ${emitExpr(folded, scope)};`); + return [`(${signature}) => {`, ...statements, `}`].join("\n"); +} + +// --------------------------------------------------------------------------- +// Buffer-native dynamics +// --------------------------------------------------------------------------- + +/** Internal bail signal: the function doesn't fit the fast path. */ +class BailError extends Error {} + +type BufferEmitContext = { + tokensName: string; + mapParamName: string; + indexParamName: string | null; + elementIndex: Map; + elementType: Map; + scope: EmitScope; +}; + +function emitBufferExpr(expr: HirExpr, context: BufferEmitContext): string { + switch (expr.kind) { + case "numberLit": + return emitNumber(expr.value, expr.raw); + case "boolLit": + return expr.value ? "true" : "false"; + case "constant": + return emitExpr(expr, context.scope); + case "localRef": { + if (expr.name === context.indexParamName) { + return "__i"; + } + if ( + expr.name === context.mapParamName || + expr.name === context.tokensName + ) { + // A token record / the token array escaping as a value (e.g. being + // returned or stored) cannot be scalarized. + throw new BailError(); + } + return context.scope.names.get(expr.name) ?? expr.name; + } + case "paramRef": + return `__params[${quoteKey(expr.name)}]`; + case "fieldAccess": { + // eslint-disable-next-line no-use-before-define -- mutual recursion + const read = emitBufferTokenRead(expr, context); + if (read !== null) { + return read; + } + throw new BailError(); + } + case "indexAccess": + throw new BailError(); + case "length": { + if ( + expr.target.kind === "localRef" && + expr.target.name === context.tokensName + ) { + return "numberOfTokens"; + } + throw new BailError(); + } + case "unary": + return `(${expr.op}${emitBufferExpr(expr.operand, context)})`; + case "binary": { + const op = expr.op === "==" ? "===" : expr.op === "!=" ? "!==" : expr.op; + return `(${emitBufferExpr(expr.left, context)} ${op} ${emitBufferExpr(expr.right, context)})`; + } + case "cond": + return `(${emitBufferExpr(expr.condition, context)} ? ${emitBufferExpr(expr.thenBranch, context)} : ${emitBufferExpr(expr.elseBranch, context)})`; + case "mathCall": + return `Math.${expr.fn}(${expr.args + .map((argument) => emitBufferExpr(argument, context)) + .join(", ")})`; + case "let": + case "recordLit": + case "arrayLit": + case "arrayMap": + case "distribution": + case "distributionMap": + throw new BailError(); + } +} + +/** + * Emits a read of `token.field` / `tokens[j].field` against the packed + * buffer, or `null` when the access is not a token read. + */ +function emitBufferTokenRead( + expr: Extract, + context: BufferEmitContext, +): string | null { + const emitRead = (base: string, field: string): string => { + const index = context.elementIndex.get(field); + if (index === undefined) { + throw new BailError(); + } + const read = `currentState[${base} + ${index}]`; + return context.elementType.get(field) === "boolean" + ? `(${read} !== 0)` + : read; + }; + + // token.field — the map parameter. + if ( + expr.target.kind === "localRef" && + expr.target.name === context.mapParamName + ) { + return emitRead("__base", expr.field); + } + // tokens[j].field — cross-token access. + if ( + expr.target.kind === "indexAccess" && + expr.target.target.kind === "localRef" && + expr.target.target.name === context.tokensName + ) { + const index = emitBufferExpr(expr.target.index, context); + return emitRead(`(${index}) * dimensions`, expr.field); + } + return null; +} + +/** + * Compiles a dynamics function to a buffer-native derivative kernel: + * + * (__params) => (currentState, dimensions, numberOfTokens) => Float64Array + * + * Returns `null` when the body doesn't fit the + * `tokens.map((token, index?) => ({ ...derivatives }))` shape (callers fall + * back to the object-API path). Discrete attributes always get derivative 0, + * matching the engine's behaviour. + */ +export function emitBufferDynamicsJs( + fn: HirFunction, + elements: readonly HirTokenElementInfo[], +): string | null { + const tokensParam = fn.params[0]; + if (!tokensParam) { + return null; + } + + let folded = foldHir(fn.body); + // Top-level bindings before the map (e.g. `const mu = parameters.g;`) are + // token-independent — hoist them before the loop. + let outerBindings: Extract["bindings"] = []; + if (folded.kind === "let") { + outerBindings = folded.bindings; + folded = folded.body; + } + if ( + folded.kind !== "arrayMap" || + folded.target.kind !== "localRef" || + folded.target.name !== tokensParam.name + ) { + return null; + } + + const mapBody = folded.body; + const bindings = mapBody.kind === "let" ? mapBody.bindings : []; + const record = mapBody.kind === "let" ? mapBody.body : mapBody; + if (record.kind !== "recordLit") { + return null; + } + + const context: BufferEmitContext = { + tokensName: tokensParam.name, + mapParamName: folded.param.name, + indexParamName: folded.indexParam?.name ?? null, + elementIndex: new Map( + elements.map((element, index) => [element.name, index]), + ), + elementType: new Map( + elements.map((element) => [element.name, element.type]), + ), + scope: { + names: new Map(), + allocator: new NameAllocator([...RESERVED_NAMES, "__base", "__i"]), + }, + }; + + try { + const statements: string[] = []; + // Hoisted before the loop: may read parameters/constants/token counts + // but never token attributes (there is no current token yet). + const hoisted: string[] = []; + const outerContext: BufferEmitContext = { + ...context, + mapParamName: "", + indexParamName: null, + }; + for (const binding of outerBindings) { + const value = emitBufferExpr(binding.value, outerContext); + const name = context.scope.allocator.allocate(binding.name); + context.scope.names.set(binding.name, name); + hoisted.push(` const ${name} = ${value};`); + } + + for (const binding of bindings) { + const value = emitBufferExpr(binding.value, context); + const name = context.scope.allocator.allocate(binding.name); + context.scope.names.set(binding.name, name); + statements.push(` const ${name} = ${value};`); + } + + for (const entry of record.entries) { + const index = context.elementIndex.get(entry.key); + const type = context.elementType.get(entry.key); + // Unknown keys are ignored and discrete attributes are forced to a zero + // derivative by the engine — skipping them reproduces that exactly + // (expressions are pure, so skipping loses no effects). + if (index === undefined || type !== "real") { + continue; + } + statements.push( + ` result[__base + ${index}] = ${emitBufferExpr(entry.value, context)};`, + ); + } + + return [ + `(__params) => (currentState, dimensions, numberOfTokens) => {`, + ` "use strict";`, + ` if (dimensions !== ${elements.length}) {`, + ` throw new Error("Expected ${elements.length} dimensions, got " + dimensions);`, + ` }`, + ...hoisted, + ` const result = new Float64Array(numberOfTokens * dimensions);`, + ` for (let __i = 0; __i < numberOfTokens; __i++) {`, + ` const __base = __i * dimensions;`, + ...statements, + ` }`, + ` return result;`, + `}`, + ].join("\n"); + } catch (error) { + if (error instanceof BailError) { + return null; + } + throw error; + } +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/hir.ts b/libs/@hashintel/petrinaut-core/src/hir/hir.ts new file mode 100644 index 00000000000..71fd844a564 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/hir.ts @@ -0,0 +1,385 @@ +/** + * HIR — high-level intermediate representation for Petrinaut user code. + * + * The HIR is a typed, source-spanned, JSON-serializable expression tree that + * sits between the surface language (today TypeScript, later the Petrinaut + * DSL) and the execution backends (JavaScript today, WASM/GPU later). + * + * Design invariants: + * - Every node carries a stable `id` (unique within one `HirFunction`) and a + * `span` into the *user-visible* source text, so analyses can be keyed off + * nodes without mutating them and diagnostics always land on the right + * range in the editor. + * - The tree is pure and expression-oriented (OCaml-like): `let` bindings, + * conditionals and `map` comprehensions instead of statements and loops. + * There is no mutation and no recursion, which makes symbolic evaluation + * (distribution DAGs, dependency sets, constant folding) decidable. + * - Probability distributions are first-class node kinds, not opaque calls, + * so the sampling structure of a transition kernel is directly visible. + */ + +/** Half-open span into the user-visible source text, in UTF-16 code units. */ +export type Span = { + start: number; + length: number; +}; + +/** Node identifier, unique within a single `HirFunction`. */ +export type HirNodeId = number; + +/** + * The user-code surfaces the HIR can currently represent. + * + * `dynamics` — differential equations (`Dynamics(...)`) + * `lambda` — transition firing predicates / stochastic rates (`Lambda(...)`) + * `kernel` — transition kernels (`TransitionKernel(...)`) + * + * Metrics and scenario expressions are expressible with the same node set but + * do not have lowering entry points yet (see `README.md`). + */ +export type HirSurfaceKind = "dynamics" | "lambda" | "kernel"; + +/** Scalar and structural types inferred over HIR nodes. */ +export type HirType = + | { kind: "real" } + | { kind: "int" } + | { kind: "bool" } + | { + kind: "record"; + fields: { name: string; type: HirType }[]; + } + | { + kind: "array"; + element: HirType; + /** Statically-known length (e.g. arc-weight token tuples). */ + length?: number; + } + /** A probability distribution over reals. */ + | { kind: "distribution" } + | { kind: "unknown" }; + +export const HIR_TYPE_REAL: HirType = { kind: "real" }; +export const HIR_TYPE_INT: HirType = { kind: "int" }; +export const HIR_TYPE_BOOL: HirType = { kind: "bool" }; +export const HIR_TYPE_DISTRIBUTION: HirType = { kind: "distribution" }; +export const HIR_TYPE_UNKNOWN: HirType = { kind: "unknown" }; + +export type HirBinaryOp = + | "+" + | "-" + | "*" + | "/" + | "%" + | "**" + | "<" + | "<=" + | ">" + | ">=" + | "==" + | "!=" + | "&&" + | "||"; + +export type HirUnaryOp = "-" | "+" | "!"; + +/** Math builtins the HIR understands (subset of ECMAScript `Math`). */ +export const HIR_MATH_FNS = [ + "abs", + "acos", + "asin", + "atan", + "atan2", + "cbrt", + "ceil", + "cos", + "cosh", + "exp", + "floor", + "hypot", + "log", + "log10", + "log2", + "max", + "min", + "pow", + "random", + "round", + "sign", + "sin", + "sinh", + "sqrt", + "tan", + "tanh", + "trunc", +] as const; + +export type HirMathFn = (typeof HIR_MATH_FNS)[number]; + +export type HirConstantName = "PI" | "E" | "Infinity" | "NaN"; + +export type HirDistributionKind = "gaussian" | "uniform" | "lognormal"; + +type HirNodeBase = { + id: HirNodeId; + span: Span; +}; + +/** Numeric literal. `raw` preserves the exact source text (e.g. `"1e-9"`). */ +export type HirNumberLit = HirNodeBase & { + kind: "numberLit"; + value: number; + raw: string; +}; + +export type HirBoolLit = HirNodeBase & { + kind: "boolLit"; + value: boolean; +}; + +/** Named numeric constant (`Math.PI`, `Infinity`, ...). */ +export type HirConstant = HirNodeBase & { + kind: "constant"; + name: HirConstantName; +}; + +/** Reference to a local binding: a function parameter, `let`/`const` binding, + * or `map` callback parameter. */ +export type HirLocalRef = HirNodeBase & { + kind: "localRef"; + name: string; +}; + +/** Model parameter access — lowered from `parameters.`. */ +export type HirParamRef = HirNodeBase & { + kind: "paramRef"; + name: string; +}; + +/** Record field access: `target.field` / `target["field"]`. */ +export type HirFieldAccess = HirNodeBase & { + kind: "fieldAccess"; + target: HirExpr; + field: string; + /** Span of just the field name, for precise diagnostics. */ + fieldSpan: Span; +}; + +/** Array element access: `target[index]`. */ +export type HirIndexAccess = HirNodeBase & { + kind: "indexAccess"; + target: HirExpr; + index: HirExpr; +}; + +/** Array length access: `target.length`. */ +export type HirLength = HirNodeBase & { + kind: "length"; + target: HirExpr; +}; + +export type HirUnary = HirNodeBase & { + kind: "unary"; + op: HirUnaryOp; + operand: HirExpr; +}; + +export type HirBinary = HirNodeBase & { + kind: "binary"; + op: HirBinaryOp; + left: HirExpr; + right: HirExpr; +}; + +/** Conditional expression — lowered from ternaries. */ +export type HirCond = HirNodeBase & { + kind: "cond"; + condition: HirExpr; + thenBranch: HirExpr; + elseBranch: HirExpr; +}; + +export type HirLetBinding = { + name: string; + nameSpan: Span; + value: HirExpr; +}; + +/** Sequential (non-recursive) bindings scoping a body expression — lowered + * from `const` statements followed by `return`. */ +export type HirLet = HirNodeBase & { + kind: "let"; + bindings: HirLetBinding[]; + body: HirExpr; +}; + +export type HirMathCall = HirNodeBase & { + kind: "mathCall"; + fn: HirMathFn; + args: HirExpr[]; +}; + +export type HirRecordEntry = { + key: string; + keySpan: Span; + value: HirExpr; +}; + +export type HirRecordLit = HirNodeBase & { + kind: "recordLit"; + entries: HirRecordEntry[]; +}; + +export type HirArrayLit = HirNodeBase & { + kind: "arrayLit"; + elements: HirExpr[]; +}; + +/** + * Array comprehension — lowered from `collection.map((param, indexParam) => + * body)`. Object-destructured callback parameters are desugared during + * lowering into `fieldAccess(localRef(param), field)` nodes. + */ +export type HirArrayMap = HirNodeBase & { + kind: "arrayMap"; + target: HirExpr; + param: { name: string; span: Span }; + indexParam?: { name: string; span: Span }; + body: HirExpr; +}; + +/** Distribution construction: `Distribution.Gaussian(mean, deviation)`, ... */ +export type HirDistribution = HirNodeBase & { + kind: "distribution"; + dist: HirDistributionKind; + args: HirExpr[]; +}; + +/** Derived distribution: `base.map((param) => body)`. */ +export type HirDistributionMap = HirNodeBase & { + kind: "distributionMap"; + base: HirExpr; + param: { name: string; span: Span }; + body: HirExpr; +}; + +export type HirExpr = + | HirNumberLit + | HirBoolLit + | HirConstant + | HirLocalRef + | HirParamRef + | HirFieldAccess + | HirIndexAccess + | HirLength + | HirUnary + | HirBinary + | HirCond + | HirLet + | HirMathCall + | HirRecordLit + | HirArrayLit + | HirArrayMap + | HirDistribution + | HirDistributionMap; + +/** + * A lowered user function. `params[0]` is the tokens/input parameter, + * `params[1]` (when present) is the parameters object — its property + * accesses are lowered to `paramRef` nodes, so it never appears as a + * `localRef` in the body. + */ +export type HirFunction = { + hirVersion: 1; + surface: HirSurfaceKind; + params: { name: string; span: Span }[]; + body: HirExpr; + /** Span of the whole user function expression in the source text. */ + span: Span; +}; + +export type HirDiagnosticSeverity = "error" | "warning" | "info" | "hint"; + +/** + * Diagnostic produced by lowering or by HIR analyses. Spans are always + * relative to the user-visible source text (not any generated wrapper), so + * they can be surfaced in the editor without further adjustment. + */ +export type HirDiagnostic = { + /** Stable string code, e.g. `"hir:unsupported-syntax"`. */ + code: string; + message: string; + severity: HirDiagnosticSeverity; + span: Span; +}; + +/** Returns the direct children of a node, in source order. */ +export function hirChildren(expr: HirExpr): HirExpr[] { + switch (expr.kind) { + case "numberLit": + case "boolLit": + case "constant": + case "localRef": + case "paramRef": + return []; + case "fieldAccess": + return [expr.target]; + case "indexAccess": + return [expr.target, expr.index]; + case "length": + return [expr.target]; + case "unary": + return [expr.operand]; + case "binary": + return [expr.left, expr.right]; + case "cond": + return [expr.condition, expr.thenBranch, expr.elseBranch]; + case "let": + return [...expr.bindings.map((binding) => binding.value), expr.body]; + case "mathCall": + return expr.args; + case "recordLit": + return expr.entries.map((entry) => entry.value); + case "arrayLit": + return expr.elements; + case "arrayMap": + return [expr.target, expr.body]; + case "distribution": + return expr.args; + case "distributionMap": + return [expr.base, expr.body]; + } +} + +/** Depth-first pre-order walk over an expression tree. */ +export function walkHir(expr: HirExpr, visit: (node: HirExpr) => void): void { + visit(expr); + for (const child of hirChildren(expr)) { + walkHir(child, visit); + } +} + +/** Formats a type for use in diagnostics. */ +export function formatHirType(type: HirType): string { + switch (type.kind) { + case "real": + return "real"; + case "int": + return "integer"; + case "bool": + return "boolean"; + case "distribution": + return "Distribution"; + case "record": + return `{ ${type.fields + .map((field) => `${field.name}: ${formatHirType(field.type)}`) + .join("; ")} }`; + case "array": + return type.length === undefined + ? `${formatHirType(type.element)}[]` + : `[${Array.from({ length: type.length }) + .fill(formatHirType(type.element)) + .join(", ")}]`; + case "unknown": + return "unknown"; + } +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/instantiate.ts b/libs/@hashintel/petrinaut-core/src/hir/instantiate.ts new file mode 100644 index 00000000000..321fdf94bb5 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/instantiate.ts @@ -0,0 +1,189 @@ +/** + * Instantiation of HIR-emitted JavaScript sources. + * + * Deliberately free of any compiler dependency (no `typescript` import): the + * simulation engine and its workers instantiate precompiled artifact strings + * without bundling the TS→HIR frontend. Compilation lives in `compile.ts`; + * artifact sources are emitted by `emit-js.ts` (object convention) and + * `emit-buffer-js.ts` (buffer ABI). + */ +import type { RuntimeDistribution } from "../simulation/authoring/user-code/distribution"; + +export type HirParameterValues = Record; +type TokensByPlace = Record[]>; + +/** Legacy object-convention user function (fallback path). */ +export type HirCompiledUserFn = ( + tokensByPlace: TokensByPlace, + parameters: HirParameterValues, +) => unknown; + +/** Buffer-native dynamics (engine `DifferentialEquationFn` shape). */ +export type HirCompiledBufferDynamics = ( + currentState: Float64Array, + dimensions: number, + numberOfTokens: number, +) => Float64Array; + +/** + * Buffer-ABI lambda: reads token attributes at `slotBases[slot] + attr` + * directly from the frame's token-value floats. Parameters are pre-bound. + */ +export type HirCompiledBufferLambda = ( + tokenValues: Float64Array, + slotBases: Int32Array, +) => number | boolean; + +/** + * Buffer-ABI kernel: writes output attributes into `out` (place-major, arc + * order) and defers distribution-valued attributes through `distSink`. + * Parameters are pre-bound. + */ +export type HirCompiledBufferKernel = ( + tokenValues: Float64Array, + slotBases: Int32Array, + out: Float64Array, + distSink: (floatIndex: number, distribution: RuntimeDistribution) => void, +) => void; + +/** A compiled program for one user-code item. `buffer` is preferred by the + * engine; `object` is the fallback for shapes the buffer emitter cannot + * scalarize. At least one is present for every compilable item. */ +export type HirLambdaArtifact = { + buffer?: { + source: string; + /** Expected `slotBases.length` — engine-side sanity check. */ + inputSlotCount: number; + }; + object?: string; +}; + +export type HirKernelArtifact = { + buffer?: { + source: string; + inputSlotCount: number; + /** Expected staging size — engine-side sanity check. */ + outputFloatCount: number; + }; + object?: string; +}; + +export type HirDynamicsArtifact = { + /** Buffer-native derivative factory (see `emitBufferDynamicsJs`). */ + buffer?: string; + /** Object-convention `(tokens, parameters) => derivatives[]` fallback. */ + object?: string; +}; + +/** + * Precompiled HIR artifacts for one SDCPN, keyed by item id (differential + * equation id / transition id, pre-flattening — the engine resolves flattened + * `path::id` ids back to their source id). Produced by `compileHirArtifacts`. + * + * Artifacts must be produced from the same SDCPN snapshot that is simulated — + * they are the only compilation path; items without an artifact fail to + * build with a diagnosis to fix the code. + */ +export type HirArtifacts = { + version: 2; + dynamics: Record; + lambdas: Record; + kernels: Record; +}; + +/** + * Distribution constructors injected into emitted code as `__dist`. + * Produces the same branded objects as the previous runtime, minus the + * `.map` method — emitted code calls `__dist.map(...)` instead. + */ +export const hirDistributionRuntime = { + gaussian: (mean: number, deviation: number): RuntimeDistribution => ({ + __brand: "distribution", + type: "gaussian", + mean, + deviation, + }), + uniform: (min: number, max: number): RuntimeDistribution => ({ + __brand: "distribution", + type: "uniform", + min, + max, + }), + lognormal: (mu: number, sigma: number): RuntimeDistribution => ({ + __brand: "distribution", + type: "lognormal", + mu, + sigma, + }), + map: ( + inner: RuntimeDistribution, + fn: (value: number) => number, + ): RuntimeDistribution => ({ + __brand: "distribution", + type: "mapped", + inner, + fn, + }), +}; + +function instantiate(source: string): unknown { + // eslint-disable-next-line no-new-func, @typescript-eslint/no-implied-eval, @typescript-eslint/no-unsafe-call + return new Function("__dist", `"use strict"; return (${source});`)( + hirDistributionRuntime, + ); +} + +function instantiateWithParams( + source: string, + parameterValues: HirParameterValues, +): unknown { + // eslint-disable-next-line no-new-func, @typescript-eslint/no-implied-eval, @typescript-eslint/no-unsafe-call + return new Function( + "__dist", + "__params", + `"use strict"; return (${source});`, + )(hirDistributionRuntime, parameterValues); +} + +/** Instantiates an emitted object-convention source (`emitUserFunctionJs`). */ +export function instantiateHirUserFn(source: string): HirCompiledUserFn { + return instantiate(source) as HirCompiledUserFn; +} + +/** + * Instantiates an emitted buffer-native dynamics factory source + * (`emitBufferDynamicsJs`), binding the run's parameter values. + */ +export function instantiateHirBufferDynamics( + source: string, + parameterValues: HirParameterValues, +): HirCompiledBufferDynamics { + const factory = instantiate(source) as ( + parameters: HirParameterValues, + ) => HirCompiledBufferDynamics; + return factory(parameterValues); +} + +/** Instantiates a buffer-ABI lambda source (`emitBufferLambdaJs`), binding + * the run's parameter values. */ +export function instantiateHirBufferLambda( + source: string, + parameterValues: HirParameterValues, +): HirCompiledBufferLambda { + return instantiateWithParams( + source, + parameterValues, + ) as HirCompiledBufferLambda; +} + +/** Instantiates a buffer-ABI kernel source (`emitBufferKernelJs`), binding + * the run's parameter values. */ +export function instantiateHirBufferKernel( + source: string, + parameterValues: HirParameterValues, +): HirCompiledBufferKernel { + return instantiateWithParams( + source, + parameterValues, + ) as HirCompiledBufferKernel; +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/lint.test.ts b/libs/@hashintel/petrinaut-core/src/hir/lint.test.ts new file mode 100644 index 00000000000..018f3027db6 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/lint.test.ts @@ -0,0 +1,319 @@ +import { describe, expect, it } from "vitest"; + +import { lintHirUserCode } from "./lint"; + +import type { + HirDynamicsContext, + HirKernelContext, + HirLambdaContext, +} from "./surface-context"; + +const dynamicsContext: HirDynamicsContext = { + surface: "dynamics", + parameters: [{ name: "g", type: "real" }], + elements: [ + { name: "x", type: "real" }, + { name: "v", type: "real" }, + { name: "count", type: "integer" }, + { name: "alive", type: "boolean" }, + ], +}; + +const poolBinding = { + name: "Pool", + colorId: "color-1", + elements: [ + { name: "x", type: "real" as const }, + { name: "alive", type: "boolean" as const }, + ], + tokenCount: 2, +}; + +const outBinding = { + name: "Out", + colorId: "color-2", + elements: [ + { name: "x", type: "real" as const }, + { name: "count", type: "integer" as const }, + ], + tokenCount: 1, +}; + +const lambdaContext: HirLambdaContext = { + surface: "lambda", + parameters: [{ name: "rate", type: "real" }], + inputPlaces: [poolBinding], + inputSlots: [{ ...poolBinding, slotStart: 0 }], + lambdaType: "stochastic", +}; + +const kernelContext: HirKernelContext = { + surface: "kernel", + parameters: [{ name: "sigma", type: "real" }], + inputPlaces: lambdaContext.inputPlaces, + inputSlots: lambdaContext.inputSlots, + outputPlaces: [outBinding], + outputSlots: [{ ...outBinding, slotStart: 0 }], + stochasticity: true, +}; + +function codes(code: string, context: Parameters[1]) { + return lintHirUserCode(code, context).diagnostics.map((diagnostic) => [ + diagnostic.code, + diagnostic.severity, + ]); +} + +describe("lintHirUserCode", () => { + it("returns no diagnostics for the default templates", () => { + expect( + codes( + `export default Dynamics((tokens, parameters) => { + return tokens.map(({ x, v }) => { + return { x: 1, v: 1 }; + }); +});`, + dynamicsContext, + ), + ).toEqual([]); + expect( + codes( + `export default Lambda((input, parameters) => 1.0);`, + lambdaContext, + ), + ).toEqual([]); + expect( + codes( + `export default TransitionKernel((input, parameters) => ({ + Out: [{ x: 0, count: 0 }], +}));`, + kernelContext, + ), + ).toEqual([]); + }); + + it("reports out-of-subset code as an error by default", () => { + const result = lintHirUserCode( + `export default Lambda((input) => { let x = 1; return x; });`, + lambdaContext, + ); + expect(result.fn).toBeUndefined(); + expect(result.diagnostics).toHaveLength(1); + expect(result.diagnostics[0]!.severity).toBe("error"); + expect(result.diagnostics[0]!.message).toContain( + "restricted TypeScript subset", + ); + }); + + it("downgrades out-of-subset code to the configured severity", () => { + const result = lintHirUserCode( + `export default Lambda((input) => { let x = 1; return x; });`, + lambdaContext, + { subsetSeverity: "info" }, + ); + expect(result.diagnostics[0]!.severity).toBe("info"); + }); + + it("keeps real errors (unknown identifier) as errors when lowering fails", () => { + const result = lintHirUserCode( + `export default Lambda((input) => nonsense + 1);`, + lambdaContext, + ); + expect(result.diagnostics[0]!.code).toBe("hir:unknown-identifier"); + expect(result.diagnostics[0]!.severity).toBe("error"); + }); + + describe("typecheck rules", () => { + it("flags derivatives on discrete attributes with a friendly message", () => { + const code = `export default Dynamics((tokens, parameters) => { + return tokens.map(({ x }) => { + return { x: 1, count: 1 }; + }); +});`; + const result = lintHirUserCode(code, dynamicsContext); + const diagnostic = result.diagnostics.find( + (candidate) => candidate.code === "hir:discrete-derivative", + )!; + expect(diagnostic.message).toContain("count"); + expect(diagnostic.message).toContain("integer"); + expect( + code.slice( + diagnostic.span.start, + diagnostic.span.start + diagnostic.span.length, + ), + ).toBe("count"); + }); + + it("flags unknown token attributes at the field span", () => { + const code = `export default Lambda((input, parameters) => input.Pool[0].missing > 0);`; + const result = lintHirUserCode(code, lambdaContext); + const diagnostic = result.diagnostics.find( + (candidate) => candidate.code === "hir:unknown-field", + )!; + expect( + code.slice( + diagnostic.span.start, + diagnostic.span.start + diagnostic.span.length, + ), + ).toBe("missing"); + }); + + it("flags unknown parameters and places", () => { + expect( + codes( + `export default Lambda((input, parameters) => parameters.nope);`, + lambdaContext, + ), + ).toContainEqual(["hir:unknown-parameter", "error"]); + expect( + codes( + `export default Lambda((input, parameters) => input.Nowhere[0].x);`, + lambdaContext, + ), + ).toContainEqual(["hir:unknown-field", "error"]); + }); + + it("flags out-of-bounds token indices (arc weight)", () => { + expect( + codes( + `export default Lambda((input, parameters) => input.Pool[2].x);`, + lambdaContext, + ), + ).toContainEqual(["hir:index-out-of-bounds", "error"]); + }); + + it("flags Distribution into a discrete attribute (H-6519 rule)", () => { + expect( + codes( + `export default TransitionKernel((input, parameters) => ({ + Out: [{ x: 0, count: Distribution.Gaussian(0, 1) }], +}));`, + kernelContext, + ), + ).toContainEqual(["hir:distribution-discrete-attribute", "error"]); + }); + + it("flags distributions outside kernels", () => { + expect( + codes( + `export default Lambda((input, parameters) => { + const d = Distribution.Gaussian(0, 1); + return 1.0; +});`, + lambdaContext, + ), + ).toContainEqual(["hir:distribution-outside-kernel", "error"]); + }); + + it("flags unknown output places and token count mismatches", () => { + const results = codes( + `export default TransitionKernel((input, parameters) => ({ + Elsewhere: [{ x: 0 }], + Out: [{ x: 0, count: 0 }, { x: 1, count: 0 }], +}));`, + kernelContext, + ); + expect(results).toContainEqual(["hir:unknown-output-place", "error"]); + expect(results).toContainEqual(["hir:output-token-count", "error"]); + }); + + it("flags missing attributes in kernel outputs", () => { + expect( + codes( + `export default TransitionKernel((input, parameters) => ({ + Out: [{ x: 0 }], +}));`, + kernelContext, + ), + ).toContainEqual(["hir:missing-attribute", "error"]); + }); + + it("flags predicate/stochastic return type mismatches", () => { + expect( + codes(`export default Lambda((input, parameters) => true);`, { + ...lambdaContext, + lambdaType: "stochastic", + }), + ).toContainEqual(["hir:lambda-return", "error"]); + expect( + codes(`export default Lambda((input, parameters) => 1);`, { + ...lambdaContext, + lambdaType: "predicate", + }), + ).toContainEqual(["hir:lambda-return", "error"]); + }); + }); + + describe("semantic rules", () => { + it("warns on Math.random", () => { + expect( + codes( + `export default Lambda((input, parameters) => Math.random());`, + lambdaContext, + ), + ).toContainEqual(["hir:math-random", "warning"]); + }); + + it("warns when a transition can never fire", () => { + expect( + codes( + `export default Lambda((input, parameters) => 2 - 2);`, + lambdaContext, + ), + ).toContainEqual(["hir:transition-never-fires", "warning"]); + expect( + codes(`export default Lambda((input, parameters) => false);`, { + ...lambdaContext, + lambdaType: "predicate", + }), + ).toContainEqual(["hir:transition-never-fires", "warning"]); + }); + + it("informs about shared distribution samples", () => { + expect( + codes( + `export default TransitionKernel((input, parameters) => { + const d = Distribution.Gaussian(0, 1); + return { Out: [{ x: d, count: 0 }] }; +});`, + kernelContext, + ), + ).toEqual([]); + expect( + codes( + `export default TransitionKernel((input, parameters) => { + const d = Distribution.Gaussian(0, 1); + return { Out: [{ x: d, count: 0 }], Out2: [{ x: d }] }; +});`, + { + ...kernelContext, + outputPlaces: [ + ...kernelContext.outputPlaces, + { + name: "Out2", + colorId: "color-3", + elements: [{ name: "x", type: "real" }], + tokenCount: 1, + }, + ], + }, + ), + ).toContainEqual(["hir:shared-sample", "info"]); + }); + + it("hints at unused bindings but ignores underscore names", () => { + const results = codes( + `export default Lambda((input, parameters) => { + const unused = parameters.rate; + const _ignored = parameters.rate; + return 1.0; +});`, + lambdaContext, + ); + expect(results).toContainEqual(["hir:unused-binding", "hint"]); + expect( + results.filter(([code]) => code === "hir:unused-binding"), + ).toHaveLength(1); + }); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/lint.ts b/libs/@hashintel/petrinaut-core/src/hir/lint.ts new file mode 100644 index 00000000000..a948b8dc085 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/lint.ts @@ -0,0 +1,178 @@ +/** + * HIR-based semantic linting. + * + * Combines lowering, type checking and analyses into a single diagnostics + * pass over one piece of user code. All spans are relative to the + * user-visible source text, so results can be surfaced in the editor without + * adjustment. + * + * Lowering failures short-circuit: when the code is outside the analyzable + * subset the result carries a single positioned diagnostic (downgraded to a + * severity chosen by the caller) and no artifacts. + */ +import { analyzeHir, foldHir } from "./analyze"; +import { walkHir } from "./hir"; +import { lowerTypeScriptToHir } from "./lower-typescript"; +import { typecheckHir } from "./typecheck"; + +import type { HirAnalysis } from "./analyze"; +import type { HirDiagnostic, HirFunction, Span } from "./hir"; +import type { HirSurfaceContext } from "./surface-context"; +import type { HirTypecheckResult } from "./typecheck"; + +export type HirLintResult = { + diagnostics: HirDiagnostic[]; + /** Present when lowering succeeded. */ + fn?: HirFunction; + analysis?: HirAnalysis; + typecheck?: HirTypecheckResult; +}; + +export type HirLintOptions = { + /** + * Severity for out-of-subset diagnostics. The HIR pipeline is the only + * compiler — code outside the subset cannot run — so these default to + * errors; tooling that only analyzes may downgrade them. + * @default "error" + */ + subsetSeverity?: HirDiagnostic["severity"]; +}; + +/** Lint codes that describe out-of-subset code rather than definite bugs. */ +const SUBSET_CODES = new Set([ + "hir:unsupported-statement", + "hir:unsupported-syntax", + "hir:unsupported-operator", + "hir:unsupported-call", + "hir:mutable-binding", + "hir:destructured-binding", + "hir:destructured-parameter", + "hir:if-statement", + "hir:loop-statement", + "hir:early-return", + "hir:spread", + "hir:computed-key", + "hir:string-value", + "hir:map-callback", + "hir:unknown-math-function", + "hir:math-reference", +]); + +const SUBSET_NOTE = + " Petrinaut compiles a restricted TypeScript subset — rewrite using `const` bindings, conditionals (`?:` or guard `if`s), `.map(...)` and `Math.*`."; + +function collectMathRandomSpans(fn: HirFunction): Span[] { + const spans: Span[] = []; + walkHir(fn.body, (expr) => { + if (expr.kind === "mathCall" && expr.fn === "random") { + spans.push(expr.span); + } + }); + return spans; +} + +/** + * Lints one piece of user code (a full `export default Ctor(...)` module). + * + * Callers should skip this when the TypeScript checker already reports errors + * for the same code — HIR lints assume syntactically and type-valid input. + */ +export function lintHirUserCode( + code: string, + context: HirSurfaceContext, + options: HirLintOptions = {}, +): HirLintResult { + const subsetSeverity = options.subsetSeverity ?? "error"; + const lowered = lowerTypeScriptToHir(code, context.surface); + + if (!lowered.ok) { + return { + diagnostics: lowered.diagnostics.map((diagnostic) => + SUBSET_CODES.has(diagnostic.code) + ? { + ...diagnostic, + severity: subsetSeverity, + message: diagnostic.message + SUBSET_NOTE, + } + : diagnostic, + ), + }; + } + + const { fn } = lowered; + const typecheck = typecheckHir(fn, context); + const analysis = analyzeHir(fn); + const diagnostics: HirDiagnostic[] = [...typecheck.diagnostics]; + + // --- Reproducibility ----------------------------------------------------- + if (analysis.dependencies.usesMathRandom) { + for (const node of collectMathRandomSpans(fn)) { + diagnostics.push({ + code: "hir:math-random", + severity: "warning", + message: + "`Math.random()` is not seeded — simulation runs will not be reproducible." + + (context.surface === "kernel" + ? " Use `Distribution.Uniform(min, max)` instead." + : ""), + span: node, + }); + } + } + + // --- Dead transitions ---------------------------------------------------- + if (context.surface === "lambda") { + const folded = foldHir(fn.body); + if ( + (context.lambdaType === "stochastic" && + folded.kind === "numberLit" && + folded.value === 0) || + (context.lambdaType === "predicate" && + folded.kind === "boolLit" && + !folded.value) + ) { + diagnostics.push({ + code: "hir:transition-never-fires", + severity: "warning", + message: + context.lambdaType === "stochastic" + ? "This rate is always 0 — the transition will never fire." + : "This predicate is always false — the transition will never fire.", + span: folded.span, + }); + } + } + + // --- Shared distribution draws ------------------------------------------- + for (const nodeId of analysis.distributionDag.sharedSampleNodeIds) { + const node = analysis.distributionDag.nodes.find( + (candidate) => candidate.nodeId === nodeId, + ); + if (!node) { + continue; + } + const sinkCount = analysis.distributionDag.sinks.filter( + (sink) => sink.nodeId === nodeId, + ).length; + diagnostics.push({ + code: "hir:shared-sample", + severity: "info", + message: `This distribution feeds ${sinkCount} output attributes — they will all receive the same sampled value. Construct separate distributions for independent samples.`, + span: node.span, + }); + } + + // --- Unused bindings ------------------------------------------------------- + for (const binding of analysis.bindings) { + if (binding.referenceCount === 0 && !binding.name.startsWith("_")) { + diagnostics.push({ + code: "hir:unused-binding", + severity: "hint", + message: `\`${binding.name}\` is never used.`, + span: binding.nameSpan, + }); + } + } + + return { diagnostics, fn, analysis, typecheck }; +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/lower-typescript.test.ts b/libs/@hashintel/petrinaut-core/src/hir/lower-typescript.test.ts new file mode 100644 index 00000000000..e720a971d7f --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/lower-typescript.test.ts @@ -0,0 +1,530 @@ +import { describe, expect, it } from "vitest"; + +import { lowerTypeScriptToHir } from "./lower-typescript"; + +import type { HirExpr } from "./hir"; + +function lowerOk(code: string, surface: "dynamics" | "lambda" | "kernel") { + const result = lowerTypeScriptToHir(code, surface); + if (!result.ok) { + throw new Error( + `Expected lowering to succeed, got: ${result.diagnostics + .map((diagnostic) => diagnostic.message) + .join("; ")}`, + ); + } + return result.fn; +} + +function lowerErr(code: string, surface: "dynamics" | "lambda" | "kernel") { + const result = lowerTypeScriptToHir(code, surface); + if (result.ok) { + throw new Error("Expected lowering to fail"); + } + return result.diagnostics; +} + +describe("lowerTypeScriptToHir", () => { + describe("module shape", () => { + it("lowers the default dynamics template", () => { + const fn = lowerOk( + `export default Dynamics((tokens, parameters) => { + return tokens.map(({ x, y }) => { + return { x: 1, y: 1 }; + }); +});`, + "dynamics", + ); + expect(fn.surface).toBe("dynamics"); + expect(fn.params.map((parameter) => parameter.name)).toEqual([ + "tokens", + "parameters", + ]); + expect(fn.body.kind).toBe("arrayMap"); + }); + + it("rejects a missing default export with a full-file span", () => { + const [diagnostic] = lowerErr(`const x = 1;`, "lambda"); + expect(diagnostic!.code).toBe("hir:unsupported-statement"); + }); + + it("rejects the wrong constructor", () => { + const [diagnostic] = lowerErr( + `export default Lambda((input) => true);`, + "kernel", + ); + expect(diagnostic!.code).toBe("hir:missing-constructor"); + }); + + it("short-circuits on syntax errors with parse diagnostics", () => { + const diagnostics = lowerErr( + `export default Lambda((input) => { return 1 + ; });`, + "lambda", + ); + expect(diagnostics[0]!.code).toBe("hir:parse-error"); + expect(diagnostics[0]!.span.start).toBeGreaterThan(0); + }); + }); + + describe("expressions", () => { + it("lowers parameters access to paramRef", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => parameters.rate * 2);`, + "lambda", + ); + expect(fn.body).toMatchObject({ + kind: "binary", + op: "*", + left: { kind: "paramRef", name: "rate" }, + right: { kind: "numberLit", value: 2 }, + }); + }); + + it("keeps exact numeric literals", () => { + const fn = lowerOk(`export default Lambda((input) => 1e-9);`, "lambda"); + expect(fn.body).toMatchObject({ kind: "numberLit", raw: "1e-9" }); + }); + + it("lowers token field access chains", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => input.Pool[0].x > 1);`, + "lambda", + ); + expect(fn.body).toMatchObject({ + kind: "binary", + op: ">", + left: { + kind: "fieldAccess", + field: "x", + target: { + kind: "indexAccess", + target: { kind: "fieldAccess", field: "Pool" }, + }, + }, + }); + }); + + it("lowers blocks with const bindings to let", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => { + const rate = parameters.base * 2; + return rate + 1; +});`, + "lambda", + ); + expect(fn.body.kind).toBe("let"); + const letNode = fn.body as Extract; + expect(letNode.bindings[0]!.name).toBe("rate"); + expect(letNode.body.kind).toBe("binary"); + }); + + it("desugars destructured map params to field accesses", () => { + const fn = lowerOk( + `export default Dynamics((tokens, parameters) => tokens.map(({ x }) => ({ x: x * 2 })));`, + "dynamics", + ); + const mapNode = fn.body as Extract; + const record = mapNode.body as Extract; + expect(record.entries[0]!.value).toMatchObject({ + kind: "binary", + left: { + kind: "fieldAccess", + field: "x", + target: { kind: "localRef", name: "__element" }, + }, + }); + }); + + it("supports .length and index parameters", () => { + const fn = lowerOk( + `export default Dynamics((tokens) => tokens.map((token, index) => ({ x: index / tokens.length })));`, + "dynamics", + ); + const mapNode = fn.body as Extract; + expect(mapNode.indexParam?.name).toBe("index"); + const record = mapNode.body as Extract; + expect(record.entries[0]!.value).toMatchObject({ + kind: "binary", + op: "/", + right: { kind: "length" }, + }); + }); + + it("unwraps type assertions and parentheses", () => { + const fn = lowerOk( + `export default Lambda((input) => ((1 as number)) + 2);`, + "lambda", + ); + expect(fn.body).toMatchObject({ + kind: "binary", + left: { kind: "numberLit", value: 1 }, + }); + }); + }); + + describe("distributions", () => { + it("lowers Distribution constructors", () => { + const fn = lowerOk( + `export default TransitionKernel((input, parameters) => ({ + Out: [{ x: Distribution.Gaussian(0, parameters.sigma) }], +}));`, + "kernel", + ); + const record = fn.body as Extract; + const tokens = record.entries[0]!.value as Extract< + HirExpr, + { kind: "arrayLit" } + >; + const token = tokens.elements[0] as Extract< + HirExpr, + { kind: "recordLit" } + >; + expect(token.entries[0]!.value).toMatchObject({ + kind: "distribution", + dist: "gaussian", + args: [{ kind: "numberLit" }, { kind: "paramRef", name: "sigma" }], + }); + }); + + it("distinguishes distribution .map from array .map", () => { + const fn = lowerOk( + `export default TransitionKernel((input) => { + const base = Distribution.Uniform(0, 1); + const scaled = base.map((value) => value * 10); + return { Out: [{ x: scaled }] }; +});`, + "kernel", + ); + const letNode = fn.body as Extract; + expect(letNode.bindings[1]!.value.kind).toBe("distributionMap"); + }); + + it("expands Math function references in distribution .map", () => { + const fn = lowerOk( + `export default TransitionKernel((input) => ({ + Out: [{ x: Distribution.Gaussian(0, 10).map(Math.cos) }], +}));`, + "kernel", + ); + const record = fn.body as Extract; + const tokens = record.entries[0]!.value as Extract< + HirExpr, + { kind: "arrayLit" } + >; + const token = tokens.elements[0] as Extract< + HirExpr, + { kind: "recordLit" } + >; + expect(token.entries[0]!.value).toMatchObject({ + kind: "distributionMap", + body: { kind: "mathCall", fn: "cos" }, + }); + }); + }); + + describe("destructuring", () => { + it("lowers const { a, b } = parameters to parameter reads", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => { + const { rate, threshold } = parameters; + return rate * threshold; +});`, + "lambda", + ); + const letNode = fn.body as Extract; + expect(letNode.bindings).toHaveLength(2); + expect(letNode.bindings[0]).toMatchObject({ + name: "rate", + value: { kind: "paramRef", name: "rate" }, + }); + expect(letNode.bindings[1]).toMatchObject({ + name: "threshold", + value: { kind: "paramRef", name: "threshold" }, + }); + }); + + it("supports renames when destructuring parameters", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => { + const { rate: r } = parameters; + return r; +});`, + "lambda", + ); + const letNode = fn.body as Extract; + expect(letNode.bindings[0]).toMatchObject({ + name: "r", + value: { kind: "paramRef", name: "rate" }, + }); + }); + + it("lowers object destructuring from token expressions", () => { + const fn = lowerOk( + `export default Lambda((tokens, parameters) => { + const { x, y } = tokens.Space[0]; + return x + y; +});`, + "lambda", + ); + const letNode = fn.body as Extract; + // temp binding + two field reads + expect(letNode.bindings).toHaveLength(3); + expect(letNode.bindings[1]).toMatchObject({ + name: "x", + value: { kind: "fieldAccess", field: "x" }, + }); + expect(letNode.bindings[2]).toMatchObject({ + name: "y", + value: { kind: "fieldAccess", field: "y" }, + }); + }); + + it("lowers array destructuring to index reads", () => { + const fn = lowerOk( + `export default Lambda((tokens, parameters) => { + const [a, b] = tokens.Space; + return a.x - b.x; +});`, + "lambda", + ); + const letNode = fn.body as Extract; + expect(letNode.bindings[1]).toMatchObject({ + name: "a", + value: { kind: "indexAccess", index: { kind: "numberLit", value: 0 } }, + }); + expect(letNode.bindings[2]).toMatchObject({ + name: "b", + value: { kind: "indexAccess", index: { kind: "numberLit", value: 1 } }, + }); + }); + + it("does not create a temp binding when destructuring a local", () => { + const fn = lowerOk( + `export default Lambda((tokens, parameters) => { + const token = tokens.Space[0]; + const { x } = token; + return x; +});`, + "lambda", + ); + const letNode = fn.body as Extract; + expect(letNode.bindings.map((binding) => binding.name)).toEqual([ + "token", + "x", + ]); + }); + + it("supports destructured function parameters", () => { + const fn = lowerOk( + `export default Lambda(({ Pool }, { rate }) => Pool[0].x * rate);`, + "lambda", + ); + expect(fn.params.map((parameter) => parameter.name)).toEqual([ + "__input", + "__parameters", + ]); + expect(fn.body).toMatchObject({ + kind: "binary", + left: { + kind: "fieldAccess", + field: "x", + target: { + kind: "indexAccess", + target: { kind: "fieldAccess", field: "Pool" }, + }, + }, + right: { kind: "paramRef", name: "rate" }, + }); + }); + + it("rejects rest and defaults in destructuring", () => { + expect( + lowerErr( + `export default Lambda((input, parameters) => { + const { a = 1 } = parameters; + return a; +});`, + "lambda", + )[0]!.code, + ).toBe("hir:destructured-binding"); + expect( + lowerErr( + `export default Lambda((input, parameters) => { + const { ...rest } = parameters; + return 1; +});`, + "lambda", + )[0]!.code, + ).toBe("hir:destructured-binding"); + }); + }); + + describe("guard clauses and early returns", () => { + it("lowers a guard if + early return to a conditional", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => { + const order = input.OpenOrders[0]; + if (order.age < order.promised_lead_time) return 0; + return parameters.rate * (order.priority > 0.5 ? 1.8 : 1); +});`, + "lambda", + ); + const letNode = fn.body as Extract; + expect(letNode.bindings[0]!.name).toBe("order"); + expect(letNode.body).toMatchObject({ + kind: "cond", + condition: { kind: "binary", op: "<" }, + thenBranch: { kind: "numberLit", value: 0 }, + elseBranch: { kind: "binary", op: "*" }, + }); + }); + + it("lowers if/else where both branches return", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => { + if (parameters.enabled) { + return parameters.rate; + } else { + return 0; + } +});`, + "lambda", + ); + expect(fn.body).toMatchObject({ + kind: "cond", + thenBranch: { kind: "paramRef", name: "rate" }, + elseBranch: { kind: "numberLit", value: 0 }, + }); + }); + + it("supports chained guards with bindings in between", () => { + const fn = lowerOk( + `export default Lambda((input, parameters) => { + if (input.Pool.length == 0) return 0; + const first = input.Pool[0]; + if (first.x < 0) return 0.5; + return first.x; +});`, + "lambda", + ); + expect(fn.body).toMatchObject({ + kind: "cond", + elseBranch: { + kind: "let", + body: { kind: "cond" }, + }, + }); + }); + + it("rejects unreachable code after return", () => { + const [diagnostic] = lowerErr( + `export default Lambda((input) => { + return 1; + return 2; +});`, + "lambda", + ); + expect(diagnostic!.code).toBe("hir:unreachable-code"); + }); + + it("rejects if branches that do not return", () => { + const [diagnostic] = lowerErr( + `export default Lambda((input, parameters) => { + if (parameters.enabled) { + const x = 1; + } + return 1; +});`, + "lambda", + ); + expect(diagnostic!.code).toBe("hir:missing-return"); + }); + }); + + describe("out-of-subset rejections with positions", () => { + it("rejects let bindings", () => { + const code = `export default Lambda((input) => { + let x = 1; + return x; +});`; + const [diagnostic] = lowerErr(code, "lambda"); + expect(diagnostic!.code).toBe("hir:mutable-binding"); + expect(code.slice(diagnostic!.span.start).startsWith("let x = 1")).toBe( + true, + ); + }); + + it("rejects unknown identifiers at the right span", () => { + const code = `export default Lambda((input) => wat + 1);`; + const [diagnostic] = lowerErr(code, "lambda"); + expect(diagnostic!.code).toBe("hir:unknown-identifier"); + expect( + code.slice( + diagnostic!.span.start, + diagnostic!.span.start + diagnostic!.span.length, + ), + ).toBe("wat"); + }); + + it("rejects arbitrary function calls", () => { + const [diagnostic] = lowerErr( + `export default Lambda((input) => fetch("x"));`, + "lambda", + ); + expect(diagnostic!.code).toBe("hir:unsupported-call"); + }); + + it("rejects object spread", () => { + const [diagnostic] = lowerErr( + `export default TransitionKernel((input) => ({ + Out: input.In.map((token) => ({ ...token })), +}));`, + "kernel", + ); + expect(diagnostic!.code).toBe("hir:spread"); + }); + + it("rejects bare use of the parameters object", () => { + const [diagnostic] = lowerErr( + `export default Lambda((input, parameters) => parameters);`, + "lambda", + ); + expect(diagnostic!.code).toBe("hir:bare-parameters-object"); + }); + }); + + describe("spans", () => { + it("every node span maps back into the source text", () => { + const code = `export default Lambda((input, parameters) => { + const a = parameters.alpha; + return a > 0 ? a * 2 : Math.abs(a); +});`; + const fn = lowerOk(code, "lambda"); + const stack = [fn.body]; + while (stack.length > 0) { + const node = stack.pop()!; + expect(node.span.start).toBeGreaterThanOrEqual(0); + expect(node.span.start + node.span.length).toBeLessThanOrEqual( + code.length, + ); + switch (node.kind) { + case "binary": + stack.push(node.left, node.right); + break; + case "let": + stack.push(...node.bindings.map((binding) => binding.value)); + stack.push(node.body); + break; + case "cond": + stack.push(node.condition, node.thenBranch, node.elseBranch); + break; + case "mathCall": + stack.push(...node.args); + break; + default: + break; + } + } + }); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/hir/lower-typescript.ts b/libs/@hashintel/petrinaut-core/src/hir/lower-typescript.ts new file mode 100644 index 00000000000..ea61a6e5fd2 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/lower-typescript.ts @@ -0,0 +1,1253 @@ +/** + * Lowers user-authored TypeScript modules into the HIR. + * + * The lowering accepts the analyzable subset of TypeScript that Petrinaut + * user code is written in: a module of the form + * + * export default ((tokens, parameters) => ) + * + * where the body is a pure expression tree — `const` bindings (with object + * and array destructuring), guard-clause `if`/early returns, ternaries, + * arithmetic/logic, `Math.*` calls, token/parameter access, `.map(...)` + * comprehensions and `Distribution.*` constructors. + * + * Anything outside that subset short-circuits lowering and produces a single + * `HirDiagnostic` whose span points at the offending syntax in the + * user-visible source text. The HIR pipeline is the only compiler, so these + * surface as errors in the LSP and the affected item cannot simulate until + * fixed. + */ +import ts from "typescript"; + +import { HIR_MATH_FNS } from "./hir"; + +import type { + HirBinaryOp, + HirDiagnostic, + HirDistributionKind, + HirExpr, + HirFunction, + HirLetBinding, + HirMathFn, + HirSurfaceKind, + Span, +} from "./hir"; + +export type LowerTypeScriptResult = + | { ok: true; fn: HirFunction; diagnostics: HirDiagnostic[] } + | { ok: false; diagnostics: HirDiagnostic[] }; + +const CONSTRUCTOR_NAMES: Record = { + dynamics: "Dynamics", + lambda: "Lambda", + kernel: "TransitionKernel", +}; + +const DISTRIBUTION_FACTORIES: Record = { + Gaussian: "gaussian", + Uniform: "uniform", + Lognormal: "lognormal", +}; + +const MATH_CONSTANTS = new Set(["PI", "E"]); + +const BINARY_OPS: Partial> = { + [ts.SyntaxKind.PlusToken]: "+", + [ts.SyntaxKind.MinusToken]: "-", + [ts.SyntaxKind.AsteriskToken]: "*", + [ts.SyntaxKind.SlashToken]: "/", + [ts.SyntaxKind.PercentToken]: "%", + [ts.SyntaxKind.AsteriskAsteriskToken]: "**", + [ts.SyntaxKind.LessThanToken]: "<", + [ts.SyntaxKind.LessThanEqualsToken]: "<=", + [ts.SyntaxKind.GreaterThanToken]: ">", + [ts.SyntaxKind.GreaterThanEqualsToken]: ">=", + [ts.SyntaxKind.EqualsEqualsToken]: "==", + [ts.SyntaxKind.EqualsEqualsEqualsToken]: "==", + [ts.SyntaxKind.ExclamationEqualsToken]: "!=", + [ts.SyntaxKind.ExclamationEqualsEqualsToken]: "!=", + [ts.SyntaxKind.AmpersandAmpersandToken]: "&&", + [ts.SyntaxKind.BarBarToken]: "||", +}; + +/** `Omit` distributed over a union, preserving each variant's shape. */ +type DistributiveOmit = Type extends unknown + ? Omit + : never; + +/** Internal signal used to short-circuit lowering with one diagnostic. */ +class LowerError extends Error { + diagnostic: HirDiagnostic; + + constructor(diagnostic: HirDiagnostic) { + super(diagnostic.message); + this.diagnostic = diagnostic; + } +} + +type LowerScope = { + /** In-scope local names: fn params, `const` bindings, map params. */ + locals: Set; + /** Locals whose bound value is distribution-valued (drives `.map` + * disambiguation between arrays and distributions). */ + distributionLocals: Set; + /** Identifiers introduced by object-destructured map params, mapped to the + * synthetic parameter they read from. */ + destructuredFields: Map; + /** Identifiers introduced by destructuring the top-level parameters + * function parameter, mapped to the model parameter they read. */ + parameterAliases: Map; + /** Name of the `parameters` function parameter, if declared. */ + parametersName: string | null; +}; + +function childScope(scope: LowerScope): LowerScope { + return { + locals: new Set(scope.locals), + distributionLocals: new Set(scope.distributionLocals), + destructuredFields: new Map(scope.destructuredFields), + parameterAliases: new Map(scope.parameterAliases), + parametersName: scope.parametersName, + }; +} + +class Lowering { + private nextId = 0; + + constructor( + private readonly sourceFile: ts.SourceFile, + private readonly surface: HirSurfaceKind, + ) {} + + spanOf(node: ts.Node): Span { + const start = node.getStart(this.sourceFile); + return { start, length: node.getWidth(this.sourceFile) }; + } + + fail(node: ts.Node, code: string, message: string): never { + throw new LowerError({ + code, + message, + severity: "error", + span: this.spanOf(node), + }); + } + + private make>( + node: ts.Node, + expr: Init, + ): Extract { + return { + ...expr, + id: this.nextId++, + span: this.spanOf(node), + } as unknown as Extract; + } + + lowerModule(): HirFunction { + const constructorName = CONSTRUCTOR_NAMES[this.surface]; + let exportAssignment: ts.ExportAssignment | undefined; + + for (const statement of this.sourceFile.statements) { + if (ts.isExportAssignment(statement) && !statement.isExportEquals) { + if (exportAssignment) { + this.fail( + statement, + "hir:multiple-default-exports", + "Only one default export is allowed.", + ); + } + exportAssignment = statement; + } else if ( + ts.isImportDeclaration(statement) && + statement.importClause?.isTypeOnly + ) { + // Type-only imports are erased; ignore them. + } else { + this.fail( + statement, + "hir:unsupported-statement", + `Only \`export default ${constructorName}(...)\` is supported at the top level.`, + ); + } + } + + if (!exportAssignment) { + throw new LowerError({ + code: "hir:missing-default-export", + message: `Expected \`export default ${constructorName}(...)\`.`, + severity: "error", + span: { start: 0, length: Math.max(this.sourceFile.text.length, 1) }, + }); + } + + const call = exportAssignment.expression; + if ( + !ts.isCallExpression(call) || + !ts.isIdentifier(call.expression) || + call.expression.text !== constructorName + ) { + this.fail( + call, + "hir:missing-constructor", + `The default export must be a \`${constructorName}(...)\` call.`, + ); + } + if (call.arguments.length !== 1) { + this.fail( + call, + "hir:constructor-arity", + `\`${constructorName}\` expects exactly one function argument.`, + ); + } + + const fnArg = call.arguments[0]!; + if (!ts.isArrowFunction(fnArg) && !ts.isFunctionExpression(fnArg)) { + this.fail( + fnArg, + "hir:missing-function", + `\`${constructorName}\` expects a function argument, e.g. \`(tokens, parameters) => ...\`.`, + ); + } + + if (fnArg.parameters.length > 2) { + this.fail( + fnArg.parameters[2]!, + "hir:too-many-parameters", + "User functions take at most two parameters (tokens and parameters).", + ); + } + + const params: HirFunction["params"] = []; + const scope: LowerScope = { + locals: new Set(), + distributionLocals: new Set(), + destructuredFields: new Map(), + parameterAliases: new Map(), + parametersName: null, + }; + + for (const [index, parameter] of fnArg.parameters.entries()) { + if (ts.isIdentifier(parameter.name)) { + const name = parameter.name.text; + params.push({ name, span: this.spanOf(parameter.name) }); + if (index === 0) { + scope.locals.add(name); + } else { + scope.parametersName = name; + } + continue; + } + if (ts.isObjectBindingPattern(parameter.name)) { + // Destructured parameter: synthesize a name and route the bound + // identifiers to the underlying tokens object / parameter reads. + const syntheticName = index === 0 ? "__input" : "__parameters"; + params.push({ + name: syntheticName, + span: this.spanOf(parameter.name), + }); + for (const element of parameter.name.elements) { + const bound = this.bindingElementParts(element); + if (bound.name !== bound.sourceName) { + this.fail( + element, + "hir:destructured-parameter", + "Renames are not supported when destructuring function parameters.", + ); + } + if (index === 0) { + scope.destructuredFields.set(bound.name, syntheticName); + } else { + scope.parameterAliases.set(bound.name, bound.sourceName); + } + } + if (index === 0) { + scope.locals.add(syntheticName); + } + continue; + } + this.fail( + parameter.name, + "hir:destructured-parameter", + "Function parameters must be plain names or object destructuring like `({ Pool }, { rate })`.", + ); + } + + const body = ts.isBlock(fnArg.body) + ? this.lowerBlock(fnArg.body, scope) + : this.lowerExpr(fnArg.body, scope); + + return { + hirVersion: 1, + surface: this.surface, + params, + body, + span: this.spanOf(fnArg), + }; + } + + /** Lowers a `{ const...; if guards...; return expr; }` block. */ + private lowerBlock(block: ts.Block, outerScope: LowerScope): HirExpr { + return this.lowerStatements( + block.statements, + childScope(outerScope), + block, + ); + } + + /** + * Lowers a statement list to an expression. Supported statements: + * `const` bindings (with object/array destructuring), guard-clause + * `if (cond) return a;` (the remaining statements become the else branch), + * terminating `if/else` where both branches return, and a final `return`. + */ + private lowerStatements( + statements: readonly ts.Statement[], + scope: LowerScope, + anchor: ts.Node, + ): HirExpr { + const bindings: HirLetBinding[] = []; + + const wrapBindings = (body: HirExpr): HirExpr => { + if (bindings.length === 0) { + return body; + } + return { + kind: "let", + bindings, + body, + id: this.nextId++, + span: this.spanOf(anchor), + }; + }; + + const failUnreachable = (statement: ts.Statement): never => + this.fail( + statement, + "hir:unreachable-code", + "This code is unreachable — the function has already returned.", + ); + + for (const [index, statement] of statements.entries()) { + if (ts.isVariableStatement(statement)) { + // eslint-disable-next-line no-bitwise -- ts.NodeFlags is a bitfield + if (!(statement.declarationList.flags & ts.NodeFlags.Const)) { + this.fail( + statement, + "hir:mutable-binding", + "`let` and `var` are not supported — use `const` (bindings are immutable).", + ); + } + for (const declaration of statement.declarationList.declarations) { + bindings.push(...this.lowerDeclaration(declaration, scope)); + } + } else if (ts.isReturnStatement(statement)) { + if (index !== statements.length - 1) { + failUnreachable(statements[index + 1]!); + } + if (!statement.expression) { + this.fail( + statement, + "hir:empty-return", + "The function must return a value.", + ); + } + return wrapBindings(this.lowerExpr(statement.expression, scope)); + } else if (ts.isIfStatement(statement)) { + const condition = this.lowerExpr(statement.expression, scope); + const thenBranch = this.lowerBranch(statement.thenStatement, scope); + + if (statement.elseStatement) { + // Both branches terminate; nothing may follow. + if (index !== statements.length - 1) { + failUnreachable(statements[index + 1]!); + } + const elseBranch = this.lowerBranch(statement.elseStatement, scope); + return wrapBindings({ + kind: "cond", + condition, + thenBranch, + elseBranch, + id: this.nextId++, + span: this.spanOf(statement), + }); + } + + // Guard clause: the remaining statements are the else branch. + const rest = statements.slice(index + 1); + if (rest.length === 0) { + this.fail( + statement, + "hir:missing-return", + "A guard `if` must be followed by more statements ending in `return`.", + ); + } + const elseBranch = this.lowerStatements( + rest, + childScope(scope), + rest[0]!, + ); + return wrapBindings({ + kind: "cond", + condition, + thenBranch, + elseBranch, + id: this.nextId++, + span: this.spanOf(statement), + }); + } else if ( + ts.isForStatement(statement) || + ts.isForOfStatement(statement) || + ts.isForInStatement(statement) || + ts.isWhileStatement(statement) || + ts.isDoStatement(statement) + ) { + this.fail( + statement, + "hir:loop-statement", + "Loops are not supported — use `.map(...)` over token arrays.", + ); + } else { + this.fail( + statement, + "hir:unsupported-statement", + `Unsupported statement (${ts.SyntaxKind[statement.kind]}) — only \`const\` bindings, guard \`if\`s and a final \`return\` are supported.`, + ); + } + } + + this.fail( + anchor, + "hir:missing-return", + "The function body must end with a `return` statement.", + ); + } + + /** Lowers an `if` branch, which must terminate with a `return`. */ + private lowerBranch(statement: ts.Statement, scope: LowerScope): HirExpr { + if (ts.isReturnStatement(statement)) { + if (!statement.expression) { + this.fail( + statement, + "hir:empty-return", + "The function must return a value.", + ); + } + return this.lowerExpr(statement.expression, scope); + } + if (ts.isBlock(statement)) { + return this.lowerStatements( + statement.statements, + childScope(scope), + statement, + ); + } + this.fail( + statement, + "hir:if-statement", + "`if` branches must end with a `return` statement.", + ); + } + + /** + * Lowers one `const` declaration to bindings, expanding object and array + * destructuring patterns. `const { a, b } = parameters` binds directly to + * parameter reads. + */ + private lowerDeclaration( + declaration: ts.VariableDeclaration, + scope: LowerScope, + ): HirLetBinding[] { + if (!declaration.initializer) { + this.fail( + declaration, + "hir:missing-initializer", + "`const` bindings must have an initializer.", + ); + } + const initializer = declaration.initializer; + const pattern = declaration.name; + + const registerBinding = (name: string, value: HirExpr): void => { + scope.locals.add(name); + scope.destructuredFields.delete(name); + scope.parameterAliases.delete(name); + if (this.isDistributionValued(value, scope)) { + scope.distributionLocals.add(name); + } else { + scope.distributionLocals.delete(name); + } + }; + + // Simple binding: const name = expr + if (ts.isIdentifier(pattern)) { + const value = this.lowerExpr(initializer, scope); + registerBinding(pattern.text, value); + return [{ name: pattern.text, nameSpan: this.spanOf(pattern), value }]; + } + + // Destructuring from the parameters object binds parameter reads + // directly: const { a, b: alias } = parameters + if ( + ts.isObjectBindingPattern(pattern) && + ts.isIdentifier(initializer) && + initializer.text === scope.parametersName && + !scope.locals.has(initializer.text) + ) { + const bindings: HirLetBinding[] = []; + for (const element of pattern.elements) { + const bound = this.bindingElementParts(element); + const value = this.make(element, { + kind: "paramRef", + name: bound.sourceName, + }); + registerBinding(bound.name, value); + bindings.push({ + name: bound.name, + nameSpan: bound.nameSpan, + value, + }); + } + return bindings; + } + + // General destructuring: bind the source once, then one binding per + // element reading from it. + const source = this.lowerExpr(initializer, scope); + const bindings: HirLetBinding[] = []; + let sourceRefName: string; + if (source.kind === "localRef") { + sourceRefName = source.name; + } else { + sourceRefName = `__destructured_${this.nextId}`; + scope.locals.add(sourceRefName); + bindings.push({ + name: sourceRefName, + nameSpan: this.spanOf(pattern), + value: source, + }); + } + + const sourceRef = (node: ts.Node): HirExpr => + this.make(node, { kind: "localRef", name: sourceRefName }); + + if (ts.isObjectBindingPattern(pattern)) { + for (const element of pattern.elements) { + const bound = this.bindingElementParts(element); + const value = this.make(element, { + kind: "fieldAccess", + target: sourceRef(element), + field: bound.sourceName, + fieldSpan: bound.sourceSpan, + }); + registerBinding(bound.name, value); + bindings.push({ name: bound.name, nameSpan: bound.nameSpan, value }); + } + return bindings; + } + + if (ts.isArrayBindingPattern(pattern)) { + for (const [elementIndex, element] of pattern.elements.entries()) { + if (ts.isOmittedExpression(element)) { + continue; + } + if ( + element.dotDotDotToken || + element.initializer || + !ts.isIdentifier(element.name) + ) { + this.fail( + element, + "hir:destructured-binding", + "Only simple array destructuring like `const [a, b] = ...` is supported.", + ); + } + const index = this.make(element, { + kind: "numberLit", + value: elementIndex, + raw: String(elementIndex), + }); + const value = this.make(element, { + kind: "indexAccess", + target: sourceRef(element), + index, + }); + registerBinding(element.name.text, value); + bindings.push({ + name: element.name.text, + nameSpan: this.spanOf(element.name), + value, + }); + } + return bindings; + } + + this.fail( + pattern, + "hir:destructured-binding", + "Unsupported binding pattern.", + ); + } + + /** Extracts (boundName, sourceProperty) from an object binding element, + * supporting renames (`{ a: alias }`). */ + private bindingElementParts(element: ts.BindingElement): { + name: string; + nameSpan: Span; + sourceName: string; + sourceSpan: Span; + } { + if ( + element.dotDotDotToken || + element.initializer || + !ts.isIdentifier(element.name) + ) { + this.fail( + element, + "hir:destructured-binding", + "Only simple destructuring like `{ a, b }` or `{ a: alias }` is supported (no defaults, rest or nesting).", + ); + } + if (element.propertyName !== undefined) { + if ( + !ts.isIdentifier(element.propertyName) && + !ts.isStringLiteralLike(element.propertyName) + ) { + this.fail( + element.propertyName, + "hir:destructured-binding", + "Computed keys are not supported in destructuring.", + ); + } + return { + name: element.name.text, + nameSpan: this.spanOf(element.name), + sourceName: element.propertyName.text, + sourceSpan: this.spanOf(element.propertyName), + }; + } + return { + name: element.name.text, + nameSpan: this.spanOf(element.name), + sourceName: element.name.text, + sourceSpan: this.spanOf(element.name), + }; + } + + private lowerExpr(node: ts.Expression, scope: LowerScope): HirExpr { + // Unwrap constructs that are transparent at runtime. + if (ts.isParenthesizedExpression(node)) { + return this.lowerExpr(node.expression, scope); + } + if ( + ts.isAsExpression(node) || + ts.isSatisfiesExpression(node) || + ts.isNonNullExpression(node) + ) { + return this.lowerExpr(node.expression, scope); + } + + if (ts.isNumericLiteral(node)) { + return this.make(node, { + kind: "numberLit", + value: Number(node.text), + raw: node.getText(this.sourceFile), + }); + } + if (node.kind === ts.SyntaxKind.TrueKeyword) { + return this.make(node, { kind: "boolLit", value: true }); + } + if (node.kind === ts.SyntaxKind.FalseKeyword) { + return this.make(node, { kind: "boolLit", value: false }); + } + if (ts.isIdentifier(node)) { + return this.lowerIdentifier(node, scope); + } + if (ts.isPrefixUnaryExpression(node)) { + return this.lowerUnary(node, scope); + } + if (ts.isBinaryExpression(node)) { + const op = BINARY_OPS[node.operatorToken.kind]; + if (!op) { + this.fail( + node.operatorToken, + "hir:unsupported-operator", + `Unsupported operator \`${node.operatorToken.getText(this.sourceFile)}\`.`, + ); + } + return this.make(node, { + kind: "binary", + op, + left: this.lowerExpr(node.left, scope), + right: this.lowerExpr(node.right, scope), + }); + } + if (ts.isConditionalExpression(node)) { + return this.make(node, { + kind: "cond", + condition: this.lowerExpr(node.condition, scope), + thenBranch: this.lowerExpr(node.whenTrue, scope), + elseBranch: this.lowerExpr(node.whenFalse, scope), + }); + } + if (ts.isPropertyAccessExpression(node)) { + return this.lowerPropertyAccess(node, scope); + } + if (ts.isElementAccessExpression(node)) { + const argument = node.argumentExpression; + if (ts.isStringLiteralLike(argument)) { + return this.make(node, { + kind: "fieldAccess", + target: this.lowerExpr(node.expression, scope), + field: argument.text, + fieldSpan: this.spanOf(argument), + }); + } + return this.make(node, { + kind: "indexAccess", + target: this.lowerExpr(node.expression, scope), + index: this.lowerExpr(argument, scope), + }); + } + if (ts.isCallExpression(node)) { + return this.lowerCall(node, scope); + } + if (ts.isObjectLiteralExpression(node)) { + return this.lowerObjectLiteral(node, scope); + } + if (ts.isArrayLiteralExpression(node)) { + const elements: HirExpr[] = []; + for (const element of node.elements) { + if (ts.isSpreadElement(element)) { + this.fail( + element, + "hir:spread", + "Spread elements are not supported yet — list tokens explicitly.", + ); + } + elements.push(this.lowerExpr(element, scope)); + } + return this.make(node, { kind: "arrayLit", elements }); + } + if (ts.isStringLiteralLike(node)) { + this.fail( + node, + "hir:string-value", + "String values are not supported — token attributes are numbers and booleans.", + ); + } + + this.fail( + node, + "hir:unsupported-syntax", + `Unsupported syntax (${ts.SyntaxKind[node.kind]}).`, + ); + } + + private lowerIdentifier(node: ts.Identifier, scope: LowerScope): HirExpr { + const name = node.text; + if (name === "Infinity") { + return this.make(node, { kind: "constant", name: "Infinity" }); + } + if (name === "NaN") { + return this.make(node, { kind: "constant", name: "NaN" }); + } + const destructuredBase = scope.destructuredFields.get(name); + if (destructuredBase !== undefined) { + const target = this.make(node, { + kind: "localRef", + name: destructuredBase, + }); + return this.make(node, { + kind: "fieldAccess", + target, + field: name, + fieldSpan: this.spanOf(node), + }); + } + const aliasedParameter = scope.parameterAliases.get(name); + if (aliasedParameter !== undefined) { + return this.make(node, { kind: "paramRef", name: aliasedParameter }); + } + if (scope.locals.has(name)) { + return this.make(node, { kind: "localRef", name }); + } + if (name === scope.parametersName) { + this.fail( + node, + "hir:bare-parameters-object", + `The parameters object can only be used via property access, e.g. \`${name}.rate\`.`, + ); + } + this.fail( + node, + "hir:unknown-identifier", + `Unknown identifier \`${name}\`.`, + ); + } + + private lowerUnary( + node: ts.PrefixUnaryExpression, + scope: LowerScope, + ): HirExpr { + const op = + node.operator === ts.SyntaxKind.MinusToken + ? "-" + : node.operator === ts.SyntaxKind.PlusToken + ? "+" + : node.operator === ts.SyntaxKind.ExclamationToken + ? "!" + : null; + if (!op) { + this.fail( + node, + "hir:unsupported-operator", + "Unsupported unary operator.", + ); + } + // Fold `-1` style negated literals so raw text is preserved. + if (op === "-" && ts.isNumericLiteral(node.operand)) { + return this.make(node, { + kind: "numberLit", + value: -Number(node.operand.text), + raw: node.getText(this.sourceFile), + }); + } + return this.make(node, { + kind: "unary", + op, + operand: this.lowerExpr(node.operand, scope), + }); + } + + private lowerPropertyAccess( + node: ts.PropertyAccessExpression, + scope: LowerScope, + ): HirExpr { + const property = node.name.text; + + if (ts.isIdentifier(node.expression)) { + const objectName = node.expression.text; + + if (objectName === "Math") { + if (MATH_CONSTANTS.has(property)) { + return this.make(node, { + kind: "constant", + name: property as "PI" | "E", + }); + } + this.fail( + node, + "hir:math-reference", + `\`Math.${property}\` can only be used as a function call.`, + ); + } + + if (objectName === "Number") { + if (property === "POSITIVE_INFINITY") { + return this.make(node, { kind: "constant", name: "Infinity" }); + } + if (property === "NaN") { + return this.make(node, { kind: "constant", name: "NaN" }); + } + this.fail( + node, + "hir:unsupported-syntax", + `\`Number.${property}\` is not supported.`, + ); + } + + if ( + objectName === scope.parametersName && + !scope.locals.has(objectName) + ) { + return this.make(node, { kind: "paramRef", name: property }); + } + } + + if (property === "length") { + return this.make(node, { + kind: "length", + target: this.lowerExpr(node.expression, scope), + }); + } + + return this.make(node, { + kind: "fieldAccess", + target: this.lowerExpr(node.expression, scope), + field: property, + fieldSpan: this.spanOf(node.name), + }); + } + + private lowerCall(node: ts.CallExpression, scope: LowerScope): HirExpr { + const callee = node.expression; + + if (ts.isPropertyAccessExpression(callee)) { + const method = callee.name.text; + + // Math.fn(...) + if ( + ts.isIdentifier(callee.expression) && + callee.expression.text === "Math" + ) { + if (!(HIR_MATH_FNS as readonly string[]).includes(method)) { + this.fail( + callee.name, + "hir:unknown-math-function", + `\`Math.${method}\` is not supported.`, + ); + } + return this.make(node, { + kind: "mathCall", + fn: method as HirMathFn, + args: node.arguments.map((argument) => + this.lowerExpr(argument, scope), + ), + }); + } + + // Distribution.Gaussian(...) / Uniform / Lognormal + if ( + ts.isIdentifier(callee.expression) && + callee.expression.text === "Distribution" + ) { + const dist = DISTRIBUTION_FACTORIES[method]; + if (!dist) { + this.fail( + callee.name, + "hir:unknown-distribution", + `Unknown distribution \`Distribution.${method}\` — expected Gaussian, Uniform or Lognormal.`, + ); + } + return this.make(node, { + kind: "distribution", + dist, + args: node.arguments.map((argument) => + this.lowerExpr(argument, scope), + ), + }); + } + + // .map(...) + if (method === "map") { + return this.lowerMapCall(node, callee, scope); + } + } + + this.fail( + node, + "hir:unsupported-call", + "Only `Math.*`, `Distribution.*` and `.map(...)` calls are supported.", + ); + } + + private lowerMapCall( + node: ts.CallExpression, + callee: ts.PropertyAccessExpression, + scope: LowerScope, + ): HirExpr { + if (node.arguments.length !== 1) { + this.fail( + node, + "hir:map-arity", + "`.map(...)` expects exactly one callback argument.", + ); + } + const target = this.lowerExpr(callee.expression, scope); + const callback = node.arguments[0]!; + + if (this.isDistributionValued(target, scope)) { + return this.lowerDistributionMap(node, target, callback, scope); + } + + if (!ts.isArrowFunction(callback) && !ts.isFunctionExpression(callback)) { + this.fail( + callback, + "hir:map-callback", + "`.map(...)` expects an inline function, e.g. `.map((token) => ...)`.", + ); + } + if (callback.parameters.length > 2) { + this.fail( + callback.parameters[2]!, + "hir:map-arity", + "`.map(...)` callbacks take at most (element, index).", + ); + } + + const bodyScope = childScope(scope); + let param: { name: string; span: Span }; + + const firstParam = callback.parameters[0]; + if (!firstParam) { + param = { name: "__element", span: this.spanOf(callback) }; + } else if (ts.isIdentifier(firstParam.name)) { + param = { + name: firstParam.name.text, + span: this.spanOf(firstParam.name), + }; + bodyScope.locals.add(param.name); + bodyScope.distributionLocals.delete(param.name); + bodyScope.destructuredFields.delete(param.name); + bodyScope.parameterAliases.delete(param.name); + } else if (ts.isObjectBindingPattern(firstParam.name)) { + param = { name: "__element", span: this.spanOf(firstParam.name) }; + for (const element of firstParam.name.elements) { + if ( + element.dotDotDotToken || + element.propertyName !== undefined || + !ts.isIdentifier(element.name) + ) { + this.fail( + element, + "hir:destructured-binding", + "Only simple destructuring like `({ x, y })` is supported in `.map(...)` callbacks.", + ); + } + bodyScope.destructuredFields.set(element.name.text, param.name); + bodyScope.locals.delete(element.name.text); + bodyScope.distributionLocals.delete(element.name.text); + } + bodyScope.locals.add(param.name); + } else { + this.fail( + firstParam.name, + "hir:destructured-binding", + "Array destructuring is not supported in `.map(...)` callbacks.", + ); + } + + let indexParam: { name: string; span: Span } | undefined; + const secondParam = callback.parameters[1]; + if (secondParam) { + if (!ts.isIdentifier(secondParam.name)) { + this.fail( + secondParam.name, + "hir:destructured-binding", + "The `.map(...)` index parameter must be a plain name.", + ); + } + indexParam = { + name: secondParam.name.text, + span: this.spanOf(secondParam.name), + }; + bodyScope.locals.add(indexParam.name); + bodyScope.distributionLocals.delete(indexParam.name); + bodyScope.destructuredFields.delete(indexParam.name); + bodyScope.parameterAliases.delete(indexParam.name); + } + + const body = ts.isBlock(callback.body) + ? this.lowerBlock(callback.body, bodyScope) + : this.lowerExpr(callback.body, bodyScope); + + return this.make(node, { + kind: "arrayMap", + target, + param, + indexParam, + body, + }); + } + + private lowerDistributionMap( + node: ts.CallExpression, + base: HirExpr, + callback: ts.Expression, + scope: LowerScope, + ): HirExpr { + // `dist.map(Math.cos)` — expand the function reference to a callback. + if ( + ts.isPropertyAccessExpression(callback) && + ts.isIdentifier(callback.expression) && + callback.expression.text === "Math" + ) { + const method = callback.name.text; + if (!(HIR_MATH_FNS as readonly string[]).includes(method)) { + this.fail( + callback.name, + "hir:unknown-math-function", + `\`Math.${method}\` is not supported.`, + ); + } + const paramName = "__sample"; + const argRef = this.make(callback, { + kind: "localRef", + name: paramName, + }); + const body = this.make(callback, { + kind: "mathCall", + fn: method as HirMathFn, + args: [argRef], + }); + return this.make(node, { + kind: "distributionMap", + base, + param: { name: paramName, span: this.spanOf(callback) }, + body, + }); + } + + if (!ts.isArrowFunction(callback) && !ts.isFunctionExpression(callback)) { + this.fail( + callback, + "hir:map-callback", + "Distribution `.map(...)` expects an inline function or a `Math.*` reference.", + ); + } + if (callback.parameters.length !== 1) { + this.fail( + callback, + "hir:map-arity", + "Distribution `.map(...)` callbacks take exactly one parameter (the sampled value).", + ); + } + const parameter = callback.parameters[0]!; + if (!ts.isIdentifier(parameter.name)) { + this.fail( + parameter.name, + "hir:destructured-binding", + "Distribution `.map(...)` parameters must be plain names.", + ); + } + + const bodyScope = childScope(scope); + const paramName = parameter.name.text; + bodyScope.locals.add(paramName); + bodyScope.distributionLocals.delete(paramName); + bodyScope.destructuredFields.delete(paramName); + bodyScope.parameterAliases.delete(paramName); + + const body = ts.isBlock(callback.body) + ? this.lowerBlock(callback.body, bodyScope) + : this.lowerExpr(callback.body, bodyScope); + + return this.make(node, { + kind: "distributionMap", + base, + param: { name: paramName, span: this.spanOf(parameter.name) }, + body, + }); + } + + private lowerObjectLiteral( + node: ts.ObjectLiteralExpression, + scope: LowerScope, + ): HirExpr { + const entries: { key: string; keySpan: Span; value: HirExpr }[] = []; + for (const property of node.properties) { + if (ts.isPropertyAssignment(property)) { + const name = property.name; + if (ts.isIdentifier(name) || ts.isStringLiteralLike(name)) { + entries.push({ + key: name.text, + keySpan: this.spanOf(name), + value: this.lowerExpr(property.initializer, scope), + }); + } else { + this.fail( + name, + "hir:computed-key", + "Computed object keys are not supported.", + ); + } + } else if (ts.isShorthandPropertyAssignment(property)) { + entries.push({ + key: property.name.text, + keySpan: this.spanOf(property.name), + value: this.lowerIdentifier(property.name, scope), + }); + } else if (ts.isSpreadAssignment(property)) { + this.fail( + property, + "hir:spread", + "Object spread is not supported yet — list attributes explicitly.", + ); + } else { + this.fail( + property, + "hir:unsupported-syntax", + "Only `key: value` entries are supported in object literals.", + ); + } + } + return this.make(node, { kind: "recordLit", entries }); + } + + /** + * Conservative distribution-ness test, used to disambiguate `.map(...)` + * between array comprehensions and derived distributions. + */ + private isDistributionValued(expr: HirExpr, scope: LowerScope): boolean { + switch (expr.kind) { + case "distribution": + case "distributionMap": + return true; + case "localRef": + return scope.distributionLocals.has(expr.name); + case "cond": + return ( + this.isDistributionValued(expr.thenBranch, scope) || + this.isDistributionValued(expr.elseBranch, scope) + ); + case "let": + return this.isDistributionValued(expr.body, scope); + default: + return false; + } + } +} + +/** + * Lowers a user-authored TypeScript module to an `HirFunction`. + * + * `code` must be the user-visible source text (not the LSP-wrapped virtual + * file); all spans in the result are relative to it. Returns `ok: false` with + * a positioned diagnostic when the code is syntactically invalid or falls + * outside the analyzable subset. + */ +export function lowerTypeScriptToHir( + code: string, + surface: HirSurfaceKind, +): LowerTypeScriptResult { + const sourceFile = ts.createSourceFile( + "user-code.ts", + code, + ts.ScriptTarget.ES2020, + /* setParentNodes */ true, + ); + + // Short-circuit on parse errors so downstream diagnostics don't pile on top + // of syntactically broken code. `parseDiagnostics` is not part of the + // public API but has been stable across TypeScript versions. + const parseDiagnostics = ( + sourceFile as ts.SourceFile & { + parseDiagnostics?: ts.DiagnosticWithLocation[]; + } + ).parseDiagnostics; + if (parseDiagnostics && parseDiagnostics.length > 0) { + return { + ok: false, + diagnostics: parseDiagnostics.slice(0, 3).map((diagnostic) => ({ + code: "hir:parse-error", + message: ts.flattenDiagnosticMessageText(diagnostic.messageText, "\n"), + severity: "error" as const, + span: { + start: diagnostic.start, + length: Math.max(diagnostic.length, 1), + }, + })), + }; + } + + try { + const fn = new Lowering(sourceFile, surface).lowerModule(); + return { ok: true, fn, diagnostics: [] }; + } catch (error) { + if (error instanceof LowerError) { + return { ok: false, diagnostics: [error.diagnostic] }; + } + throw error; + } +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts b/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts new file mode 100644 index 00000000000..effe996e72f --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/surface-context.ts @@ -0,0 +1,339 @@ +/** + * Surface contexts describe the model-derived environment a piece of user + * code executes in: which parameters exist, which places deliver tokens (and + * with which attributes), and what shape the result must have. + * + * They are consumed by the HIR type checker, lint rules and the buffer-ABI + * emitters. Lowering and generic JS emission work without a context so they + * stay usable where the SDCPN is not available (e.g. isolated tests). + * + * ## Slot layout invariant (buffer ABI) + * + * `inputSlots` lists one entry per **colored, non-inhibitor input arc whose + * color has at least one element**, in arc order — this must match the + * engine's `inputPlacesWithTokenValues` filter (`dimensions > 0 && arcType + * !== "inhibitor"`, see `compute-possible-transition.ts`). Each entry + * occupies `tokenCount` (arc weight) consecutive token slots. `outputSlots` + * lists colored output arcs in arc order (including zero-element colors, + * which occupy no floats but must still be produced by kernels). + * + * When several arcs reference places with the same display name, the runtime + * `tokensByPlace` object keeps the LAST arc's tokens (object key overwrite) — + * both the merged `inputPlaces` bindings and name-based slot resolution + * follow that rule. + */ +import { getArcEndpoint } from "../arc-endpoints"; +import { + createArcPlaceResolver, + DEFAULT_PETRINAUT_EXTENSIONS, + getEffectiveTransitionLambdaType, + getTransitionLogicAvailability, + type PetrinautExtensionSettings, +} from "../extensions"; + +import type { + Color, + ColorElementType, + ComponentInstance, + InputArc, + OutputArc, + Parameter, + Place, + SDCPN, + Transition, +} from "../types/sdcpn"; + +/** The net a transition/differential equation belongs to (root or subnet). */ +export type HirNetScope = { + places: Place[]; + parameters: Parameter[]; + componentInstances?: ComponentInstance[]; +}; + +export type HirParameterInfo = { + /** The identifier used in user code (`parameters.`). */ + name: string; + type: ColorElementType; +}; + +export type HirTokenElementInfo = { + name: string; + type: ColorElementType; +}; + +/** A place, as seen from a transition through its arcs (merged by name). */ +export type HirPlaceBinding = { + /** Display name used as the object key in user code (scoped + * `Instance::Port` for component ports). */ + name: string; + colorId: string; + elements: HirTokenElementInfo[]; + /** Number of tokens delivered/produced (arc weight; last arc wins for + * duplicate names, matching runtime object-key overwrite). */ + tokenCount: number; +}; + +/** + * One arc's worth of token slots in the buffer ABI (see module doc). + * `slotStart` is the index of the arc's first token slot. + */ +export type HirArcSlot = { + name: string; + colorId: string; + elements: HirTokenElementInfo[]; + tokenCount: number; + slotStart: number; +}; + +export type HirDynamicsContext = { + surface: "dynamics"; + parameters: HirParameterInfo[]; + /** Attributes of the color the differential equation is attached to. */ + elements: HirTokenElementInfo[]; +}; + +export type HirLambdaContext = { + surface: "lambda"; + parameters: HirParameterInfo[]; + inputPlaces: HirPlaceBinding[]; + /** Buffer-ABI input slot layout (see module doc). */ + inputSlots: HirArcSlot[]; + lambdaType: "predicate" | "stochastic"; +}; + +export type HirKernelContext = { + surface: "kernel"; + parameters: HirParameterInfo[]; + inputPlaces: HirPlaceBinding[]; + /** Buffer-ABI input slot layout (see module doc). */ + inputSlots: HirArcSlot[]; + outputPlaces: HirPlaceBinding[]; + /** Buffer-ABI output slot layout: colored output arcs in arc order. */ + outputSlots: HirArcSlot[]; + /** Whether the stochasticity extension is enabled (Distribution allowed). */ + stochasticity: boolean; +}; + +export type HirSurfaceContext = + | HirDynamicsContext + | HirLambdaContext + | HirKernelContext; + +const SCOPE_SEPARATOR = "::"; + +function toParameterInfos(parameters: Parameter[]): HirParameterInfo[] { + return parameters.map((parameter) => ({ + name: parameter.variableName, + type: parameter.type, + })); +} + +/** + * Display name for the place an arc points at — must match the property keys + * generated for the LSP `Input`/`Output` types and the runtime + * `tokensByPlace` objects (`Instance::Port` scoping for component ports). + */ +function getArcPlaceDisplayName( + arc: InputArc | OutputArc, + sdcpn: SDCPN, + net: HirNetScope, +): string { + const endpoint = getArcEndpoint(arc); + + if (endpoint.kind === "place") { + const place = net.places.find((pl) => pl.id === endpoint.placeId); + return place?.name ?? endpoint.placeId; + } + + const instance = net.componentInstances?.find( + (inst) => inst.id === endpoint.componentInstanceId, + ); + const subnet = instance + ? sdcpn.subnets?.find((sn) => sn.id === instance.subnetId) + : undefined; + const portPlace = subnet?.places.find( + (pl) => pl.id === endpoint.portPlaceId && pl.isPort, + ); + + if (instance && portPlace) { + return `${instance.name}${SCOPE_SEPARATOR}${portPlace.name}`; + } + + return endpoint.componentInstanceId + SCOPE_SEPARATOR + endpoint.portPlaceId; +} + +type ArcCollection = { + /** Merged by display name; last arc wins (runtime key-overwrite parity). */ + bindings: HirPlaceBinding[]; + /** Per-arc slots in arc order (engine filter applied). */ + slots: HirArcSlot[]; +}; + +function collectArcs( + arcs: (InputArc | OutputArc)[], + sdcpn: SDCPN, + net: HirNetScope, + extensions: PetrinautExtensionSettings, + colorById: Map, + options: { includeZeroElementColors: boolean }, +): ArcCollection { + const resolveArcPlace = createArcPlaceResolver(sdcpn, net, { + componentPortsEnabled: extensions.subnets, + }); + const bindings = new Map(); + const slots: HirArcSlot[] = []; + let slotStart = 0; + + for (const arc of arcs) { + if ("type" in arc && arc.type === "inhibitor") { + continue; + } + const place = resolveArcPlace(arc); + if (!extensions.colors || !place?.colorId) { + continue; + } + const color = colorById.get(place.colorId); + if (!color) { + continue; + } + if (color.elements.length === 0 && !options.includeZeroElementColors) { + continue; + } + + const name = getArcPlaceDisplayName(arc, sdcpn, net); + const elements = color.elements.map((element) => ({ + name: element.name, + type: element.type, + })); + + // Last arc with a given name wins, matching the runtime's object-key + // overwrite when building `tokensByPlace`. + bindings.set(name, { + name, + colorId: color.id, + elements, + tokenCount: arc.weight, + }); + + slots.push({ + name, + colorId: color.id, + elements, + tokenCount: arc.weight, + slotStart, + }); + slotStart += arc.weight; + } + + return { bindings: [...bindings.values()], slots }; +} + +function collectColors( + sdcpn: SDCPN, + extensions: PetrinautExtensionSettings, +): Map { + const allColors = extensions.colors + ? [ + ...sdcpn.types, + ...(sdcpn.subnets ?? []).flatMap((subnet) => subnet.types), + ] + : []; + return new Map(allColors.map((color) => [color.id, color])); +} + +/** + * Builds the context for a differential equation attached to `color`. + * Pass `net` when the equation belongs to a subnet (its parameters apply). + */ +export function buildDynamicsContext( + sdcpn: SDCPN, + colorId: string, + extensions: PetrinautExtensionSettings = DEFAULT_PETRINAUT_EXTENSIONS, + net: HirNetScope = sdcpn, +): HirDynamicsContext | null { + const color = collectColors(sdcpn, extensions).get(colorId); + if (!color) { + return null; + } + return { + surface: "dynamics", + parameters: toParameterInfos(extensions.parameters ? net.parameters : []), + elements: color.elements.map((element) => ({ + name: element.name, + type: element.type, + })), + }; +} + +/** + * Builds the context for a transition's lambda. Pass `net` when the + * transition belongs to a subnet. + */ +export function buildLambdaContext( + sdcpn: SDCPN, + transition: Transition, + extensions: PetrinautExtensionSettings = DEFAULT_PETRINAUT_EXTENSIONS, + net: HirNetScope = sdcpn, +): HirLambdaContext { + const availability = getTransitionLogicAvailability( + transition, + sdcpn, + extensions, + net, + ); + const { bindings, slots } = collectArcs( + transition.inputArcs, + sdcpn, + net, + extensions, + collectColors(sdcpn, extensions), + { includeZeroElementColors: false }, + ); + return { + surface: "lambda", + parameters: toParameterInfos(extensions.parameters ? net.parameters : []), + inputPlaces: bindings, + inputSlots: slots, + lambdaType: getEffectiveTransitionLambdaType(transition, availability), + }; +} + +/** + * Builds the context for a transition's kernel. Pass `net` when the + * transition belongs to a subnet. + */ +export function buildKernelContext( + sdcpn: SDCPN, + transition: Transition, + extensions: PetrinautExtensionSettings = DEFAULT_PETRINAUT_EXTENSIONS, + net: HirNetScope = sdcpn, +): HirKernelContext { + const colorById = collectColors(sdcpn, extensions); + const inputs = collectArcs( + transition.inputArcs, + sdcpn, + net, + extensions, + colorById, + { includeZeroElementColors: false }, + ); + const outputs = collectArcs( + transition.outputArcs, + sdcpn, + net, + extensions, + colorById, + // Zero-element colored outputs still require a (token-count-sized) entry + // in the kernel result, they just carry no attribute floats. + { includeZeroElementColors: true }, + ); + return { + surface: "kernel", + parameters: toParameterInfos(extensions.parameters ? net.parameters : []), + inputPlaces: inputs.bindings, + inputSlots: inputs.slots, + outputPlaces: outputs.bindings, + outputSlots: outputs.slots, + stochasticity: extensions.stochasticity, + }; +} diff --git a/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts b/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts new file mode 100644 index 00000000000..f89610ef260 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/hir/typecheck.ts @@ -0,0 +1,797 @@ +/** + * Type inference and checking over the HIR. + * + * Types flow bottom-up from literals and the surface context (token attribute + * types, parameter types). The checker is deliberately pragmatic: `unknown` + * propagates silently, and diagnostics are only produced where the checker is + * confident — TypeScript remains the authoritative type oracle for TS-authored + * code, so these diagnostics focus on friendlier, domain-specific messages + * (unknown attribute, derivative on a discrete attribute, Distribution into a + * discrete attribute, ...) and on standalone use where no TS checker runs + * (the future DSL, runtime compilation). + */ +import { + formatHirType, + HIR_TYPE_BOOL, + HIR_TYPE_DISTRIBUTION, + HIR_TYPE_INT, + HIR_TYPE_REAL, + HIR_TYPE_UNKNOWN, +} from "./hir"; + +import type { + HirDiagnostic, + HirExpr, + HirFunction, + HirMathFn, + HirNodeId, + HirType, + Span, +} from "./hir"; +import type { HirSurfaceContext, HirTokenElementInfo } from "./surface-context"; + +export type HirTypecheckResult = { + /** Inferred type per HIR node id. */ + types: Map; + returnType: HirType; + diagnostics: HirDiagnostic[]; +}; + +const MATH_FN_ARITY: Record = { + abs: { min: 1, max: 1 }, + acos: { min: 1, max: 1 }, + asin: { min: 1, max: 1 }, + atan: { min: 1, max: 1 }, + atan2: { min: 2, max: 2 }, + cbrt: { min: 1, max: 1 }, + ceil: { min: 1, max: 1 }, + cos: { min: 1, max: 1 }, + cosh: { min: 1, max: 1 }, + exp: { min: 1, max: 1 }, + floor: { min: 1, max: 1 }, + hypot: { min: 1, max: Infinity }, + log: { min: 1, max: 1 }, + log10: { min: 1, max: 1 }, + log2: { min: 1, max: 1 }, + max: { min: 1, max: Infinity }, + min: { min: 1, max: Infinity }, + pow: { min: 2, max: 2 }, + random: { min: 0, max: 0 }, + round: { min: 1, max: 1 }, + sign: { min: 1, max: 1 }, + sin: { min: 1, max: 1 }, + sinh: { min: 1, max: 1 }, + sqrt: { min: 1, max: 1 }, + tan: { min: 1, max: 1 }, + tanh: { min: 1, max: 1 }, + trunc: { min: 1, max: 1 }, +}; + +/** Math functions that always produce integers. */ +const INT_MATH_FNS = new Set([ + "ceil", + "floor", + "round", + "sign", + "trunc", +]); + +function elementType(element: HirTokenElementInfo): HirType { + switch (element.type) { + case "real": + return HIR_TYPE_REAL; + case "integer": + return HIR_TYPE_INT; + case "boolean": + return HIR_TYPE_BOOL; + } +} + +function tokenRecordType(elements: HirTokenElementInfo[]): HirType { + return { + kind: "record", + fields: elements.map((element) => ({ + name: element.name, + type: elementType(element), + })), + }; +} + +function isNumeric(type: HirType): boolean { + return type.kind === "real" || type.kind === "int" || type.kind === "unknown"; +} + +function isBoolish(type: HirType): boolean { + return type.kind === "bool" || type.kind === "unknown"; +} + +/** Least upper bound, collapsing to `unknown` on structural mismatch. */ +function joinTypes(left: HirType, right: HirType): HirType { + if (left.kind === "unknown" || right.kind === "unknown") { + return HIR_TYPE_UNKNOWN; + } + if (left.kind === right.kind) { + if (left.kind === "array" && right.kind === "array") { + return { + kind: "array", + element: joinTypes(left.element, right.element), + length: left.length === right.length ? left.length : undefined, + }; + } + if (left.kind === "record" && right.kind === "record") { + const rightByName = new Map( + right.fields.map((field) => [field.name, field.type]), + ); + if ( + left.fields.length !== right.fields.length || + !left.fields.every((field) => rightByName.has(field.name)) + ) { + return HIR_TYPE_UNKNOWN; + } + return { + kind: "record", + fields: left.fields.map((field) => ({ + name: field.name, + type: joinTypes(field.type, rightByName.get(field.name)!), + })), + }; + } + return left; + } + if ( + (left.kind === "int" && right.kind === "real") || + (left.kind === "real" && right.kind === "int") + ) { + return HIR_TYPE_REAL; + } + return HIR_TYPE_UNKNOWN; +} + +class Typechecker { + readonly types = new Map(); + readonly diagnostics: HirDiagnostic[] = []; + + constructor(private readonly context: HirSurfaceContext) {} + + report( + span: Span, + code: string, + message: string, + severity: HirDiagnostic["severity"] = "error", + ): void { + this.diagnostics.push({ code, message, severity, span }); + } + + check(fn: HirFunction): HirType { + const env = new Map(); + const tokensParam = fn.params[0]; + if (tokensParam) { + env.set(tokensParam.name, this.tokensParamType()); + } + const returnType = this.infer(fn.body, env); + this.checkReturnType(fn, returnType); + return returnType; + } + + private tokensParamType(): HirType { + const context = this.context; + switch (context.surface) { + case "dynamics": + return { kind: "array", element: tokenRecordType(context.elements) }; + case "lambda": + case "kernel": + return { + kind: "record", + fields: context.inputPlaces.map((place) => ({ + name: place.name, + type: { + kind: "array", + element: tokenRecordType(place.elements), + length: place.tokenCount, + } satisfies HirType, + })), + }; + } + } + + private infer(expr: HirExpr, env: Map): HirType { + const type = this.inferUncached(expr, env); + this.types.set(expr.id, type); + return type; + } + + private inferUncached(expr: HirExpr, env: Map): HirType { + switch (expr.kind) { + case "numberLit": + return Number.isInteger(expr.value) ? HIR_TYPE_INT : HIR_TYPE_REAL; + case "boolLit": + return HIR_TYPE_BOOL; + case "constant": + return HIR_TYPE_REAL; + case "localRef": + return env.get(expr.name) ?? HIR_TYPE_UNKNOWN; + case "paramRef": { + const parameter = this.context.parameters.find( + (candidate) => candidate.name === expr.name, + ); + if (!parameter) { + this.report( + expr.span, + "hir:unknown-parameter", + `Unknown parameter \`${expr.name}\`.`, + ); + return HIR_TYPE_UNKNOWN; + } + return elementType({ name: parameter.name, type: parameter.type }); + } + case "fieldAccess": { + const targetType = this.infer(expr.target, env); + if (targetType.kind === "record") { + const field = targetType.fields.find( + (candidate) => candidate.name === expr.field, + ); + if (!field) { + this.report( + expr.fieldSpan, + "hir:unknown-field", + `\`${expr.field}\` does not exist here — expected one of: ${targetType.fields + .map((candidate) => candidate.name) + .join(", ")}.`, + ); + return HIR_TYPE_UNKNOWN; + } + return field.type; + } + if (targetType.kind !== "unknown") { + this.report( + expr.fieldSpan, + "hir:invalid-field-access", + `Cannot access field \`${expr.field}\` on a ${formatHirType(targetType)} value.`, + ); + } + return HIR_TYPE_UNKNOWN; + } + case "indexAccess": { + const targetType = this.infer(expr.target, env); + const indexType = this.infer(expr.index, env); + if (!isNumeric(indexType)) { + this.report( + expr.index.span, + "hir:invalid-index", + "Array indices must be numbers.", + ); + } + if (targetType.kind === "array") { + if ( + targetType.length !== undefined && + expr.index.kind === "numberLit" && + expr.index.value >= targetType.length + ) { + this.report( + expr.span, + "hir:index-out-of-bounds", + `Index ${expr.index.raw} is out of bounds — only ${targetType.length} token(s) are available here.`, + ); + } + return targetType.element; + } + if (targetType.kind !== "unknown") { + this.report( + expr.span, + "hir:invalid-index", + `Cannot index into a ${formatHirType(targetType)} value.`, + ); + } + return HIR_TYPE_UNKNOWN; + } + case "length": { + const targetType = this.infer(expr.target, env); + if (targetType.kind !== "array" && targetType.kind !== "unknown") { + this.report( + expr.span, + "hir:invalid-length", + `\`.length\` is only available on arrays, not ${formatHirType(targetType)}.`, + ); + } + return HIR_TYPE_INT; + } + case "unary": { + const operandType = this.infer(expr.operand, env); + if (expr.op === "!") { + if (!isBoolish(operandType)) { + this.report( + expr.span, + "hir:type-mismatch", + `\`!\` expects a boolean, got ${formatHirType(operandType)}.`, + ); + } + return HIR_TYPE_BOOL; + } + if (!isNumeric(operandType)) { + this.report( + expr.span, + "hir:type-mismatch", + `Unary \`${expr.op}\` expects a number, got ${formatHirType(operandType)}.`, + ); + return HIR_TYPE_UNKNOWN; + } + return operandType; + } + case "binary": { + const leftType = this.infer(expr.left, env); + const rightType = this.infer(expr.right, env); + switch (expr.op) { + case "&&": + case "||": + if (!isBoolish(leftType) || !isBoolish(rightType)) { + this.report( + expr.span, + "hir:type-mismatch", + `\`${expr.op}\` expects booleans.`, + ); + } + return HIR_TYPE_BOOL; + case "==": + case "!=": + return HIR_TYPE_BOOL; + case "<": + case "<=": + case ">": + case ">=": + if (!isNumeric(leftType) || !isNumeric(rightType)) { + this.report( + expr.span, + "hir:type-mismatch", + `\`${expr.op}\` expects numbers.`, + ); + } + return HIR_TYPE_BOOL; + case "/": + case "**": + this.expectNumeric(expr, leftType, rightType); + return HIR_TYPE_REAL; + default: + this.expectNumeric(expr, leftType, rightType); + return leftType.kind === "int" && rightType.kind === "int" + ? HIR_TYPE_INT + : leftType.kind === "unknown" || rightType.kind === "unknown" + ? HIR_TYPE_UNKNOWN + : HIR_TYPE_REAL; + } + } + case "cond": { + const conditionType = this.infer(expr.condition, env); + if (!isBoolish(conditionType)) { + this.report( + expr.condition.span, + "hir:type-mismatch", + `Conditions must be booleans, got ${formatHirType(conditionType)}.`, + ); + } + return joinTypes( + this.infer(expr.thenBranch, env), + this.infer(expr.elseBranch, env), + ); + } + case "let": { + const scoped = new Map(env); + for (const binding of expr.bindings) { + scoped.set(binding.name, this.infer(binding.value, scoped)); + } + return this.infer(expr.body, scoped); + } + case "mathCall": { + const arity = MATH_FN_ARITY[expr.fn]; + if (expr.args.length < arity.min || expr.args.length > arity.max) { + this.report( + expr.span, + "hir:math-arity", + arity.min === arity.max + ? `\`Math.${expr.fn}\` expects ${arity.min} argument(s).` + : `\`Math.${expr.fn}\` expects at least ${arity.min} argument(s).`, + ); + } + let allInt = true; + for (const argument of expr.args) { + const argumentType = this.infer(argument, env); + if (!isNumeric(argumentType)) { + this.report( + argument.span, + "hir:type-mismatch", + `\`Math.${expr.fn}\` expects numbers, got ${formatHirType(argumentType)}.`, + ); + } + if (argumentType.kind !== "int") { + allInt = false; + } + } + if (INT_MATH_FNS.has(expr.fn)) { + return HIR_TYPE_INT; + } + if ( + (expr.fn === "min" || expr.fn === "max" || expr.fn === "abs") && + allInt + ) { + return HIR_TYPE_INT; + } + return HIR_TYPE_REAL; + } + case "recordLit": { + const seen = new Set(); + const fields: { name: string; type: HirType }[] = []; + for (const entry of expr.entries) { + if (seen.has(entry.key)) { + this.report( + entry.keySpan, + "hir:duplicate-key", + `Duplicate key \`${entry.key}\`.`, + ); + } + seen.add(entry.key); + fields.push({ name: entry.key, type: this.infer(entry.value, env) }); + } + return { kind: "record", fields }; + } + case "arrayLit": { + let element: HirType = HIR_TYPE_UNKNOWN; + for (const [index, item] of expr.elements.entries()) { + const itemType = this.infer(item, env); + element = index === 0 ? itemType : joinTypes(element, itemType); + } + return { kind: "array", element, length: expr.elements.length }; + } + case "arrayMap": { + const targetType = this.infer(expr.target, env); + let elementType_: HirType = HIR_TYPE_UNKNOWN; + let length: number | undefined; + if (targetType.kind === "array") { + elementType_ = targetType.element; + length = targetType.length; + } else if (targetType.kind !== "unknown") { + this.report( + expr.target.span, + "hir:invalid-map", + `\`.map(...)\` is only available on arrays, not ${formatHirType(targetType)}.`, + ); + } + const scoped = new Map(env); + scoped.set(expr.param.name, elementType_); + if (expr.indexParam) { + scoped.set(expr.indexParam.name, HIR_TYPE_INT); + } + return { + kind: "array", + element: this.infer(expr.body, scoped), + length, + }; + } + case "distribution": { + this.checkDistributionAllowed(expr.span); + if (expr.args.length !== 2) { + this.report( + expr.span, + "hir:distribution-arity", + "Distributions take exactly two arguments.", + ); + } + for (const argument of expr.args) { + const argumentType = this.infer(argument, env); + if (!isNumeric(argumentType)) { + this.report( + argument.span, + "hir:type-mismatch", + `Distribution arguments must be numbers, got ${formatHirType(argumentType)}.`, + ); + } + } + return HIR_TYPE_DISTRIBUTION; + } + case "distributionMap": { + const baseType = this.infer(expr.base, env); + if (baseType.kind !== "distribution" && baseType.kind !== "unknown") { + this.report( + expr.base.span, + "hir:invalid-map", + `Distribution \`.map(...)\` requires a distribution, got ${formatHirType(baseType)}.`, + ); + } + const scoped = new Map(env); + scoped.set(expr.param.name, HIR_TYPE_REAL); + const bodyType = this.infer(expr.body, scoped); + if (!isNumeric(bodyType)) { + this.report( + expr.body.span, + "hir:type-mismatch", + `Distribution \`.map(...)\` must return a number, got ${formatHirType(bodyType)}.`, + ); + } + return HIR_TYPE_DISTRIBUTION; + } + } + } + + private expectNumeric( + expr: Extract, + leftType: HirType, + rightType: HirType, + ): void { + if (!isNumeric(leftType) || !isNumeric(rightType)) { + this.report( + expr.span, + "hir:type-mismatch", + `\`${expr.op}\` expects numbers.`, + ); + } + } + + private checkDistributionAllowed(span: Span): void { + if (this.context.surface !== "kernel") { + this.report( + span, + "hir:distribution-outside-kernel", + "Distributions can only be produced in transition kernel outputs — they are sampled when the transition fires.", + ); + } else if (!this.context.stochasticity) { + this.report( + span, + "hir:distribution-disabled", + "Distributions require the stochasticity extension, which is disabled for this net.", + ); + } + } + + /** Surface-specific checks on the function's result type. */ + private checkReturnType(fn: HirFunction, returnType: HirType): void { + const context = this.context; + const bodySpan = fn.body.kind === "let" ? fn.body.body.span : fn.body.span; + + switch (context.surface) { + case "dynamics": { + if (returnType.kind !== "array") { + if (returnType.kind !== "unknown") { + this.report( + bodySpan, + "hir:dynamics-return", + `Dynamics must return an array of derivative records (one per token), got ${formatHirType(returnType)}.`, + ); + } + return; + } + if (returnType.element.kind !== "record") { + return; + } + const elementByName = new Map( + context.elements.map((element) => [element.name, element]), + ); + for (const field of returnType.element.fields) { + const element = elementByName.get(field.name); + if (!element) { + this.report( + this.derivativeKeySpan(fn.body, field.name) ?? bodySpan, + "hir:unknown-attribute", + `\`${field.name}\` is not an attribute of this type.`, + ); + } else if (element.type !== "real") { + this.report( + this.derivativeKeySpan(fn.body, field.name) ?? bodySpan, + "hir:discrete-derivative", + `\`${field.name}\` is a discrete (${element.type}) attribute — only real attributes can have derivatives.`, + ); + } else if (!isNumeric(field.type)) { + this.report( + this.derivativeKeySpan(fn.body, field.name) ?? bodySpan, + "hir:type-mismatch", + `The derivative of \`${field.name}\` must be a number, got ${formatHirType(field.type)}.`, + ); + } + } + return; + } + case "lambda": { + if (context.lambdaType === "predicate") { + if (!isBoolish(returnType)) { + this.report( + bodySpan, + "hir:lambda-return", + `Predicate lambdas must return a boolean, got ${formatHirType(returnType)}.`, + ); + } + } else if (!isNumeric(returnType)) { + this.report( + bodySpan, + "hir:lambda-return", + `Stochastic lambdas must return a number (the firing rate), got ${formatHirType(returnType)}.`, + ); + } + return; + } + case "kernel": { + if (returnType.kind !== "record") { + if (returnType.kind !== "unknown") { + this.report( + bodySpan, + "hir:kernel-return", + `Transition kernels must return an object mapping output places to token arrays, got ${formatHirType(returnType)}.`, + ); + } + return; + } + this.checkKernelOutput(fn, returnType); + return; + } + } + } + + private checkKernelOutput( + fn: HirFunction, + returnType: Extract, + ): void { + const context = this.context; + if (context.surface !== "kernel") { + return; + } + const placeByName = new Map( + context.outputPlaces.map((place) => [place.name, place]), + ); + const outputRecord = this.resultRecordLit(fn.body); + + for (const place of context.outputPlaces) { + if (!returnType.fields.some((field) => field.name === place.name)) { + this.report( + outputRecord?.span ?? fn.body.span, + "hir:missing-output-place", + `The kernel must return tokens for output place \`${place.name}\`.`, + ); + } + } + + for (const field of returnType.fields) { + const place = placeByName.get(field.name); + const keySpan = + outputRecord?.entries.find((entry) => entry.key === field.name) + ?.keySpan ?? fn.body.span; + if (!place) { + this.report( + keySpan, + "hir:unknown-output-place", + `\`${field.name}\` is not an output place of this transition${ + context.outputPlaces.length > 0 + ? ` — expected: ${context.outputPlaces + .map((candidate) => candidate.name) + .join(", ")}` + : "" + }.`, + ); + continue; + } + if (field.type.kind !== "array") { + continue; + } + if ( + field.type.length !== undefined && + field.type.length !== place.tokenCount + ) { + this.report( + keySpan, + "hir:output-token-count", + `\`${field.name}\` receives ${place.tokenCount} token(s) per firing (arc weight), but ${field.type.length} were returned.`, + ); + } + if (field.type.element.kind !== "record") { + continue; + } + const elementByName = new Map( + place.elements.map((element) => [element.name, element]), + ); + for (const tokenField of field.type.element.fields) { + const element = elementByName.get(tokenField.name); + if (!element) { + this.report( + keySpan, + "hir:unknown-attribute", + `\`${tokenField.name}\` is not an attribute of tokens in \`${field.name}\`.`, + ); + continue; + } + if ( + tokenField.type.kind === "distribution" && + element.type !== "real" + ) { + this.report( + keySpan, + "hir:distribution-discrete-attribute", + `\`${tokenField.name}\` is a discrete (${element.type}) attribute — Distributions can only produce real attributes.`, + ); + } + if (tokenField.type.kind === "bool" && element.type !== "boolean") { + this.report( + keySpan, + "hir:type-mismatch", + `\`${tokenField.name}\` is a ${element.type} attribute but received a boolean.`, + ); + } + if ( + (tokenField.type.kind === "real" || tokenField.type.kind === "int") && + element.type === "boolean" + ) { + this.report( + keySpan, + "hir:type-mismatch", + `\`${tokenField.name}\` is a boolean attribute but received a number.`, + ); + } + } + for (const element of place.elements) { + if ( + !field.type.element.fields.some( + (tokenField) => tokenField.name === element.name, + ) + ) { + this.report( + keySpan, + "hir:missing-attribute", + `Tokens for \`${field.name}\` are missing the \`${element.name}\` attribute.`, + ); + } + } + } + } + + /** The record literal the function ultimately returns, if syntactically + * evident (unwrapping `let` bodies). */ + private resultRecordLit( + expr: HirExpr, + ): Extract | null { + if (expr.kind === "recordLit") { + return expr; + } + if (expr.kind === "let") { + return this.resultRecordLit(expr.body); + } + return null; + } + + /** Span of the `name:` key inside the derivative record literal returned by + * a dynamics body, when it is syntactically evident. */ + private derivativeKeySpan(expr: HirExpr, key: string): Span | null { + if (expr.kind === "let") { + return this.derivativeKeySpan(expr.body, key); + } + if (expr.kind === "arrayMap") { + return this.derivativeKeySpan(expr.body, key); + } + if (expr.kind === "arrayLit") { + for (const element of expr.elements) { + const span = this.derivativeKeySpan(element, key); + if (span) { + return span; + } + } + return null; + } + if (expr.kind === "recordLit") { + return expr.entries.find((entry) => entry.key === key)?.keySpan ?? null; + } + if (expr.kind === "cond") { + return ( + this.derivativeKeySpan(expr.thenBranch, key) ?? + this.derivativeKeySpan(expr.elseBranch, key) + ); + } + return null; + } +} + +/** Runs type inference and surface checks over a lowered function. */ +export function typecheckHir( + fn: HirFunction, + context: HirSurfaceContext, +): HirTypecheckResult { + const checker = new Typechecker(context); + const returnType = checker.check(fn); + return { + types: checker.types, + returnType, + diagnostics: checker.diagnostics, + }; +} diff --git a/libs/@hashintel/petrinaut-core/src/index.ts b/libs/@hashintel/petrinaut-core/src/index.ts index 6ad41881619..f9ad3ba7e63 100644 --- a/libs/@hashintel/petrinaut-core/src/index.ts +++ b/libs/@hashintel/petrinaut-core/src/index.ts @@ -243,6 +243,15 @@ export type { TextDocumentIdentifier, } from "./lsp"; +// --- HIR (type-only from the main entry; the compiler itself stays in the +// LSP worker, runtime instantiation in ./hir-runtime) --- +export type { + HirArtifacts, + HirCompileFailure, + HirCompileResult, + HirDiagnostic, +} from "./hir"; + // --- Playback --- export { createPlayback, diff --git a/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts b/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts index c382ca936cb..b939b384a29 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/language-client.ts @@ -1,3 +1,5 @@ +import { DiagnosticSeverity } from "vscode-languageserver-types"; + import { createWorkerLspTransport, type LspTransport, @@ -5,6 +7,8 @@ import { } from "./transport"; import type { PetrinautExtensionSettings } from "../extensions"; +// Type-only: must not pull the compiler (`typescript`) into client bundles. +import type { HirCompileResult } from "../hir"; import type { ReadableStore } from "../store"; import type { SDCPN } from "../types/sdcpn"; import type { @@ -25,6 +29,12 @@ import type { export type DiagnosticsSnapshot = { byUri: Map; total: number; + /** + * Error-severity diagnostics only. Use this (not `total`) to decide whether + * the net is simulatable — warnings/hints (e.g. HIR semantic lints) don't + * block running. + */ + errorCount: number; }; /** @@ -74,6 +84,16 @@ export interface LanguageClient { uri: DocumentUri, position: Position, ): Promise; + /** + * Compiles the SDCPN's user code to HIR artifacts (in the worker, where the + * TypeScript frontend lives). Pass the result as `hirArtifacts` when + * starting simulations/experiments — the engine has no compiler of its own. + */ + requestHirArtifacts( + this: void, + sdcpn: SDCPN, + extensions?: PetrinautExtensionSettings, + ): Promise; /** * Tear down the transport. Pending requests reject with "Worker terminated". @@ -89,6 +109,7 @@ export type CreateLanguageClientConfig = const EMPTY_DIAGNOSTICS: DiagnosticsSnapshot = { byUri: new Map(), total: 0, + errorCount: 0, }; function createReadableStore(initial: T): ReadableStore & { @@ -119,14 +140,20 @@ function buildSnapshot( ): DiagnosticsSnapshot { const byUri = new Map(); let total = 0; + let errorCount = 0; for (const param of allParams) { if (param.diagnostics.length === 0) { continue; } byUri.set(param.uri, param.diagnostics); total += param.diagnostics.length; + for (const diagnostic of param.diagnostics) { + if (diagnostic.severity === DiagnosticSeverity.Error) { + errorCount += 1; + } + } } - return { byUri, total }; + return { byUri, total, errorCount }; } /** @@ -287,6 +314,12 @@ export function createLanguageClient( position, }); }, + requestHirArtifacts(sdcpn, extensions) { + return sendRequest("sdcpn/compileHirArtifacts", { + sdcpn, + extensions, + }); + }, dispose() { if (disposed) { diff --git a/libs/@hashintel/petrinaut-core/src/lsp/lib/check-hir.ts b/libs/@hashintel/petrinaut-core/src/lsp/lib/check-hir.ts new file mode 100644 index 00000000000..4d035d68d5d --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/lsp/lib/check-hir.ts @@ -0,0 +1,128 @@ +/** + * Bridges HIR semantic lints into the TypeScript-diagnostic-shaped pipeline + * used by the checker and the LSP worker. + * + * HIR lint spans are already relative to the user-visible source text, so no + * prefix adjustment is needed — the produced diagnostics can be serialized + * with the same `serializeDiagnostic` path as adjusted TS diagnostics. + */ +import ts from "typescript"; + +import { lintHirUserCode } from "../../hir"; + +import type { HirDiagnostic, HirSurfaceContext } from "../../hir"; + +/** + * Stable numeric codes for HIR diagnostics (ts.Diagnostic.code is numeric). + * Codes ≥ 99000 are reserved for Petrinaut's own checks. Append only — these + * appear in editor UI and may be referenced in docs. + */ +export const HIR_DIAGNOSTIC_CODES: Record = { + "hir:parse-error": 99001, + "hir:missing-default-export": 99002, + "hir:multiple-default-exports": 99003, + "hir:missing-constructor": 99004, + "hir:constructor-arity": 99005, + "hir:missing-function": 99006, + "hir:too-many-parameters": 99007, + "hir:destructured-parameter": 99008, + "hir:unsupported-statement": 99009, + "hir:mutable-binding": 99010, + "hir:destructured-binding": 99011, + "hir:missing-initializer": 99012, + "hir:early-return": 99013, + "hir:empty-return": 99014, + "hir:if-statement": 99015, + "hir:loop-statement": 99016, + "hir:missing-return": 99017, + "hir:unsupported-operator": 99018, + "hir:unsupported-syntax": 99019, + "hir:unsupported-call": 99020, + "hir:unknown-identifier": 99021, + "hir:bare-parameters-object": 99022, + "hir:math-reference": 99023, + "hir:unknown-math-function": 99024, + "hir:unknown-distribution": 99025, + "hir:map-arity": 99026, + "hir:map-callback": 99027, + "hir:computed-key": 99028, + "hir:spread": 99029, + "hir:string-value": 99030, + "hir:unknown-parameter": 99040, + "hir:unknown-field": 99041, + "hir:invalid-field-access": 99042, + "hir:invalid-index": 99043, + "hir:index-out-of-bounds": 99044, + "hir:invalid-length": 99045, + "hir:type-mismatch": 99046, + "hir:duplicate-key": 99047, + "hir:invalid-map": 99048, + "hir:math-arity": 99049, + "hir:distribution-arity": 99050, + "hir:distribution-outside-kernel": 99051, + "hir:distribution-disabled": 99052, + "hir:dynamics-return": 99053, + "hir:unknown-attribute": 99054, + "hir:discrete-derivative": 99055, + "hir:lambda-return": 99056, + "hir:kernel-return": 99057, + "hir:unknown-output-place": 99058, + "hir:output-token-count": 99059, + "hir:distribution-discrete-attribute": 99060, + "hir:missing-attribute": 99061, + "hir:missing-output-place": 99062, + "hir:unreachable-code": 99063, + "hir:math-random": 99070, + "hir:transition-never-fires": 99071, + "hir:shared-sample": 99072, + "hir:unused-binding": 99073, +}; + +const FALLBACK_CODE = 99000; + +function toTsCategory( + severity: HirDiagnostic["severity"], +): ts.DiagnosticCategory { + switch (severity) { + case "error": + return ts.DiagnosticCategory.Error; + case "warning": + return ts.DiagnosticCategory.Warning; + case "info": + return ts.DiagnosticCategory.Message; + case "hint": + return ts.DiagnosticCategory.Suggestion; + } +} + +function toTsDiagnostic(diagnostic: HirDiagnostic): ts.Diagnostic { + return { + file: undefined, + start: diagnostic.span.start, + length: diagnostic.span.length, + messageText: diagnostic.message, + category: toTsCategory(diagnostic.severity), + code: HIR_DIAGNOSTIC_CODES[diagnostic.code] ?? FALLBACK_CODE, + source: "hir", + }; +} + +/** + * Runs the HIR semantic linter over one item's user code, returning + * TS-diagnostic-shaped results with offsets relative to the user content. + * + * Callers should only invoke this when TypeScript reported no errors for the + * item — HIR lints assume type-valid input, and stacking both would be noise. + */ +export function getHirDiagnosticsForItem( + code: string, + context: HirSurfaceContext, +): ts.Diagnostic[] { + if (code.trim() === "") { + return []; + } + // Out-of-subset code is an error: the HIR pipeline is the only compiler, + // so such code cannot simulate. + const result = lintHirUserCode(code, context, { subsetSeverity: "error" }); + return result.diagnostics.map(toTsDiagnostic); +} diff --git a/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.test.ts b/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.test.ts index 4a3ab30571d..508eecd38d8 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.test.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.test.ts @@ -82,9 +82,14 @@ describe("checkSDCPN", () => { // WHEN const result = check(sdcpn); - // THEN + // THEN — valid; only the HIR unused-binding hint remains (`value`) expect(result.isValid).toBe(true); - expect(result.itemDiagnostics).toHaveLength(0); + const nonHintDiagnostics = result.itemDiagnostics.flatMap((item) => + item.diagnostics.filter( + (diag) => diag.category !== ts.DiagnosticCategory.Suggestion, + ), + ); + expect(nonHintDiagnostics).toHaveLength(0); }); it("returns valid for code accessing defined parameters", () => { @@ -115,9 +120,14 @@ describe("checkSDCPN", () => { // WHEN const result = check(sdcpn); - // THEN + // THEN — valid; only HIR unused-binding hints remain (`a`, `e`) expect(result.isValid).toBe(true); - expect(result.itemDiagnostics).toHaveLength(0); + const nonHintDiagnostics = result.itemDiagnostics.flatMap((item) => + item.diagnostics.filter( + (diag) => diag.category !== ts.DiagnosticCategory.Suggestion, + ), + ); + expect(nonHintDiagnostics).toHaveLength(0); }); it("returns invalid when accessing undefined token property", () => { @@ -1355,4 +1365,135 @@ describe("checkSDCPN", () => { expect(result.itemDiagnostics).toHaveLength(0); }); }); + + describe("HIR semantic lints", () => { + it("surfaces Math.random as a warning with source 'hir' at the right span", () => { + // GIVEN — TS-valid code with a reproducibility problem + const lambdaCode = `export default Lambda((input, parameters) => { + return input.Source[0].value * Math.random(); +});`; + const sdcpn = createSDCPN({ + types: [{ id: "color1", elements: [{ name: "value", type: "real" }] }], + places: [ + { id: "place1", name: "Source", colorId: "color1" }, + { id: "place2", name: "Target", colorId: "color1" }, + ], + transitions: [ + { + id: "t1", + lambdaType: "stochastic", + inputArcs: [{ placeId: "place1", weight: 1, type: "standard" }], + outputArcs: [{ placeId: "place2", weight: 1 }], + lambdaCode, + transitionKernelCode: `export default TransitionKernel((input, parameters) => { + return { Target: [input.Source[0]] }; + });`, + }, + ], + }); + + // WHEN + const result = check(sdcpn); + + // THEN — a warning (does not invalidate the net) at Math.random() + expect(result.isValid).toBe(true); + const lambdaItem = result.itemDiagnostics.find( + (item) => item.itemType === "transition-lambda", + ); + expect(lambdaItem).toBeDefined(); + const hirDiagnostic = lambdaItem!.diagnostics.find( + (diag) => diag.source === "hir", + ); + expect(hirDiagnostic).toBeDefined(); + expect(hirDiagnostic!.category).toBe(ts.DiagnosticCategory.Warning); + expect( + lambdaCode.slice( + hirDiagnostic!.start!, + hirDiagnostic!.start! + hirDiagnostic!.length!, + ), + ).toBe("Math.random()"); + }); + + it("skips HIR lints when TypeScript already reports errors", () => { + // GIVEN — code with a TS error (unknown property) + const sdcpn = createSDCPN({ + types: [{ id: "color1", elements: [{ name: "value", type: "real" }] }], + places: [ + { id: "place1", name: "Source", colorId: "color1" }, + { id: "place2", name: "Target", colorId: "color1" }, + ], + transitions: [ + { + id: "t1", + lambdaType: "predicate", + inputArcs: [{ placeId: "place1", weight: 1, type: "standard" }], + outputArcs: [{ placeId: "place2", weight: 1 }], + lambdaCode: `export default Lambda((input, parameters) => { + return input.Source[0].missing > Math.random(); + });`, + transitionKernelCode: `export default TransitionKernel((input, parameters) => { + return { Target: [input.Source[0]] }; + });`, + }, + ], + }); + + // WHEN + const result = check(sdcpn); + + // THEN — only TS diagnostics; no doubled-up HIR results + expect(result.isValid).toBe(false); + const lambdaItem = result.itemDiagnostics.find( + (item) => item.itemType === "transition-lambda", + )!; + expect( + lambdaItem.diagnostics.every((diag) => diag.source !== "hir"), + ).toBe(true); + }); + + it("marks out-of-subset code as an error (it cannot be compiled)", () => { + // GIVEN — valid TS using a loop (outside the HIR subset) + const sdcpn = createSDCPN({ + types: [{ id: "color1", elements: [{ name: "value", type: "real" }] }], + places: [ + { id: "place1", name: "Source", colorId: "color1" }, + { id: "place2", name: "Target", colorId: "color1" }, + ], + transitions: [ + { + id: "t1", + lambdaType: "predicate", + inputArcs: [{ placeId: "place1", weight: 1, type: "standard" }], + outputArcs: [{ placeId: "place2", weight: 1 }], + lambdaCode: `export default Lambda((input, parameters) => { + let total = 0; + for (const token of input.Source) { + total += token.value; + } + return total > 0; + });`, + transitionKernelCode: `export default TransitionKernel((input, parameters) => { + return { Target: [input.Source[0]] }; + });`, + }, + ], + }); + + // WHEN + const result = check(sdcpn); + + // THEN — the HIR pipeline is the only compiler, so this blocks running + expect(result.isValid).toBe(false); + const lambdaItem = result.itemDiagnostics.find( + (item) => item.itemType === "transition-lambda", + )!; + const hirDiagnostic = lambdaItem.diagnostics.find( + (diag) => diag.source === "hir", + )!; + expect(hirDiagnostic.category).toBe(ts.DiagnosticCategory.Error); + expect( + ts.flattenDiagnosticMessageText(hirDiagnostic.messageText, "\n"), + ).toContain("restricted TypeScript subset"); + }); + }); }); diff --git a/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.ts b/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.ts index 6fb9458e375..839c740b262 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/lib/checker.ts @@ -3,8 +3,15 @@ import { getTransitionLogicAvailability, type PetrinautExtensionSettings, } from "../../extensions"; +import { + buildDynamicsContext, + buildKernelContext, + buildLambdaContext, +} from "../../hir"; +import { getHirDiagnosticsForItem } from "./check-hir"; import { getItemFilePath } from "./file-paths"; +import type { HirSurfaceContext } from "../../hir"; import type { SDCPN } from "../../types/sdcpn"; import type { SDCPNLanguageServer } from "./create-sdcpn-language-service"; import type ts from "typescript"; @@ -26,15 +33,48 @@ export type SDCPNDiagnostic = { }; export type SDCPNCheckResult = { - /** Whether the SDCPN is valid (no errors) */ + /** Whether the SDCPN is valid (no error-severity diagnostics). */ isValid: boolean; /** All diagnostics grouped by item */ itemDiagnostics: SDCPNDiagnostic[]; }; +/** TS `DiagnosticCategory.Error` (kept numeric — `typescript` is imported + * type-only here). */ +const TS_CATEGORY_ERROR = 1; + +/** + * Collects TS diagnostics for one code file and, when TypeScript found no + * errors, appends HIR semantic lints (friendlier domain rules, analyzability + * notes). HIR spans are user-content-relative, matching the adjusted TS + * diagnostics. + */ +function collectItemDiagnostics( + server: SDCPNLanguageServer, + filePath: string, + hirContext: HirSurfaceContext | null, +): ts.Diagnostic[] { + const semanticDiagnostics = server.getSemanticDiagnostics(filePath); + const syntacticDiagnostics = server.getSyntacticDiagnostics(filePath); + const allDiagnostics = [...syntacticDiagnostics, ...semanticDiagnostics]; + + const hasTsError = allDiagnostics.some( + (diagnostic) => diagnostic.category === TS_CATEGORY_ERROR, + ); + if (!hasTsError && hirContext) { + const userContent = server.getUserContent(filePath); + if (userContent !== undefined) { + allDiagnostics.push(...getHirDiagnosticsForItem(userContent, hirContext)); + } + } + + return allDiagnostics; +} + /** - * Checks the validity of an SDCPN by running TypeScript validation - * on all user-provided code (transitions and differential equations). + * Checks the validity of an SDCPN by running TypeScript validation and HIR + * semantic lints on all user-provided code (transitions and differential + * equations). */ export function checkSDCPN( sdcpn: SDCPN, @@ -50,9 +90,11 @@ export function checkSDCPN( const filePath = getItemFilePath("differential-equation-code", { id: de.id, }); - const semanticDiagnostics = server.getSemanticDiagnostics(filePath); - const syntacticDiagnostics = server.getSyntacticDiagnostics(filePath); - const allDiagnostics = [...syntacticDiagnostics, ...semanticDiagnostics]; + const allDiagnostics = collectItemDiagnostics( + server, + filePath, + de.colorId ? buildDynamicsContext(sdcpn, de.colorId, extensions) : null, + ); if (allDiagnostics.length > 0) { itemDiagnostics.push({ @@ -77,14 +119,11 @@ export function checkSDCPN( const lambdaFilePath = getItemFilePath("transition-lambda-code", { transitionId: transition.id, }); - const lambdaSemanticDiagnostics = - server.getSemanticDiagnostics(lambdaFilePath); - const lambdaSyntacticDiagnostics = - server.getSyntacticDiagnostics(lambdaFilePath); - const lambdaDiagnostics = [ - ...lambdaSyntacticDiagnostics, - ...lambdaSemanticDiagnostics, - ]; + const lambdaDiagnostics = collectItemDiagnostics( + server, + lambdaFilePath, + buildLambdaContext(sdcpn, transition, extensions), + ); if (lambdaDiagnostics.length > 0) { itemDiagnostics.push({ @@ -100,14 +139,11 @@ export function checkSDCPN( const kernelFilePath = getItemFilePath("transition-kernel-code", { transitionId: transition.id, }); - const kernelSemanticDiagnostics = - server.getSemanticDiagnostics(kernelFilePath); - const kernelSyntacticDiagnostics = - server.getSyntacticDiagnostics(kernelFilePath); - const kernelDiagnostics = [ - ...kernelSyntacticDiagnostics, - ...kernelSemanticDiagnostics, - ]; + const kernelDiagnostics = collectItemDiagnostics( + server, + kernelFilePath, + buildKernelContext(sdcpn, transition, extensions), + ); if (kernelDiagnostics.length > 0) { itemDiagnostics.push({ @@ -121,7 +157,11 @@ export function checkSDCPN( } return { - isValid: itemDiagnostics.length === 0, + isValid: !itemDiagnostics.some((item) => + item.diagnostics.some( + (diagnostic) => diagnostic.category === TS_CATEGORY_ERROR, + ), + ), itemDiagnostics, }; } diff --git a/libs/@hashintel/petrinaut-core/src/lsp/lib/ts-to-lsp.ts b/libs/@hashintel/petrinaut-core/src/lsp/lib/ts-to-lsp.ts index f10c93baa53..703d0f44026 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/lib/ts-to-lsp.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/lib/ts-to-lsp.ts @@ -97,6 +97,6 @@ export function serializeDiagnostic( ), message: ts.flattenDiagnosticMessageText(diag.messageText, "\n"), code: diag.code, - source: "ts", + source: diag.source ?? "ts", }; } diff --git a/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts b/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts index d1437ec9d74..1d2e8e3f2b3 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/worker/language-server.worker.ts @@ -21,6 +21,7 @@ import { import { createWorkerThreadRuntime } from "../../environment"; import { DEFAULT_PETRINAUT_EXTENSIONS } from "../../extensions"; +import { compileHirArtifacts } from "../../hir"; import { checkSDCPN } from "../lib/checker"; import { SDCPNLanguageServer } from "../lib/create-sdcpn-language-service"; import { filePathToUri, uriToFilePath } from "../lib/document-uris"; @@ -367,6 +368,15 @@ workerRuntime.onMessage((data) => { // --- Requests (send response) --- + case "sdcpn/compileHirArtifacts": { + const { id } = data; + respond( + id, + compileHirArtifacts(data.params.sdcpn, data.params.extensions), + ); + break; + } + case "textDocument/completion": { const { id } = data; if (!server) { diff --git a/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts b/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts index ee97bd9db8c..9df7c2aee59 100644 --- a/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts +++ b/libs/@hashintel/petrinaut-core/src/lsp/worker/protocol.ts @@ -135,6 +135,15 @@ type ClientRequest = id: number; method: "textDocument/signatureHelp"; params: TextDocumentPositionParams; + } + | { + jsonrpc: "2.0"; + id: number; + method: "sdcpn/compileHirArtifacts"; + params: { + sdcpn: SDCPN; + extensions?: PetrinautExtensionSettings; + }; }; /** Any message from the main thread to the worker. */ diff --git a/libs/@hashintel/petrinaut-core/src/simulation/ARCHITECTURE.md b/libs/@hashintel/petrinaut-core/src/simulation/ARCHITECTURE.md index e002b18d1b0..9b93f44cd80 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/ARCHITECTURE.md +++ b/libs/@hashintel/petrinaut-core/src/simulation/ARCHITECTURE.md @@ -8,9 +8,15 @@ The simulation module is split into five boundaries: - `api.ts` defines the public Core contract. Consumers receive `SimulationFrameReader` and summary state, not engine storage objects. -- `authoring/metric/`, `authoring/scenario/`, and `authoring/user-code/` - compile user-authored inputs. Shared same-realm hardening helpers live in - `authoring/sandbox.ts`. +- `authoring/metric/` and `authoring/scenario/` compile metric/scenario + inputs (`new Function` with same-realm hardening from + `authoring/sandbox.ts`). Dynamics/lambda/kernel code is compiled solely + through the HIR pipeline (`src/hir/README.md`) into + `SimulationInput.hirArtifacts` — produced off-engine (LSP worker) and + instantiated dependency-free via `src/hir-runtime.ts`. Buffer-ABI programs + read/write the packed frame floats directly (`engine/buffer-transition.ts`); + object-convention programs are the per-item fallback for shapes the buffer + emitter cannot scalarize. - `engine/` builds SDCPN definitions into runnable state and advances internal `EngineFrame` state. `EngineFrame` is an `ArrayBuffer`; the SDCPN-specialized `EngineFrameLayout` lives on the `SimulationInstance`. diff --git a/libs/@hashintel/petrinaut-core/src/simulation/api.ts b/libs/@hashintel/petrinaut-core/src/simulation/api.ts index 87206b2121c..c63f54233b7 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/api.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/api.ts @@ -1,5 +1,6 @@ import type { AbortSignalLike, WorkerFactoryLike } from "../environment"; import type { PetrinautExtensionSettings } from "../extensions"; +import type { HirArtifacts } from "../hir-runtime"; import type { EventStream } from "../instance"; import type { ReadableStore } from "../store"; import type { Place, SDCPN, TokenRecord } from "../types/sdcpn"; @@ -67,6 +68,13 @@ export type SimulationConfig = { dt: number; /** Maximum simulation time. Null = no limit. */ maxTime: number | null; + /** + * Precompiled HIR artifacts for the net's user code, produced by + * `compileHirArtifacts` (or `LanguageClient.requestHirArtifacts`). The + * engine has no compiler of its own: items with user code and no artifact + * fail to build. + */ + hirArtifacts?: HirArtifacts; backpressure?: BackpressureConfig; /** Optional cancellation. Aborting tears down the simulation. */ signal?: AbortSignalLike; diff --git a/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.test.ts deleted file mode 100644 index d708c604be6..00000000000 --- a/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.test.ts +++ /dev/null @@ -1,368 +0,0 @@ -import { describe, expect, it } from "vitest"; - -import { compileUserCode } from "./compile-user-code"; - -describe("compileUserCode", () => { - describe("basic functionality", () => { - it("should compile a simple function with default export", () => { - const code = ` - export default Dynamics((a, b, c) => { - return a + b + c; - }); - `; - const fn = compileUserCode<[number, number, number]>(code, "Dynamics"); - const result = fn(1, 2, 3); - expect(result).toBe(6); - }); - - it("should compile an arrow function with implicit return", () => { - const code = ` - export default Dynamics((x, y) => x * y); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(5, 7); - expect(result).toBe(35); - }); - - it("should handle multiple parameters", () => { - const code = ` - export default Dynamics((a, b, c, d) => { - return a + b * c - d; - }); - `; - const fn = compileUserCode<[number, number, number, number]>( - code, - "Dynamics", - ); - const result = fn(10, 2, 5, 3); - expect(result).toBe(17); - }); - }); - - describe("TypeScript type annotations", () => { - it("should strip TypeScript type annotations from parameters", () => { - const code = ` - export default Dynamics((x: number, y: number) => { - return x + y; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(10, 20); - expect(result).toBe(30); - }); - - it("should strip return type annotations", () => { - const code = ` - export default Dynamics((x: number, y: number): number => { - return x * y; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(5, 6); - expect(result).toBe(30); - }); - - it("should strip type annotations from variable declarations", () => { - const code = ` - export default Dynamics((a: number, b: number): number => { - const sum: number = a + b; - const product: number = sum * 2; - return product; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(3, 7); - expect(result).toBe(20); - }); - - it("should handle complex TypeScript types", () => { - const code = ` - export default Dynamics((arr: number[]): number => { - return arr.reduce((sum: number, n: number) => sum + n, 0); - }); - `; - const fn = compileUserCode<[number[]]>(code, "Dynamics"); - const result = fn([1, 2, 3, 4, 5]); - expect(result).toBe(15); - }); - - it("should handle optional parameters", () => { - const code = ` - export default Dynamics((a: number, b?: number): number => { - return b !== undefined ? a + b : a; - }); - `; - const fn = compileUserCode<[number, number?]>(code, "Dynamics"); - expect(fn(5, 10)).toBe(15); - expect(fn(5)).toBe(5); - }); - }); - - describe("complex logic", () => { - it("should handle array operations", () => { - const code = ` - export default Dynamics((matrix) => { - return matrix.map(row => row.map(x => x * 2)); - }); - `; - const fn = compileUserCode<[number[][]]>(code, "Dynamics"); - const result = fn([ - [1, 2], - [3, 4], - ]); - expect(result).toEqual([ - [2, 4], - [6, 8], - ]); - }); - - it("should handle closures", () => { - const code = ` - export default Dynamics((multiplier) => { - return (value) => value * multiplier; - }); - `; - const fn = compileUserCode<[number]>(code, "Dynamics"); - const innerFn = fn(5); - expect(typeof innerFn).toBe("function"); - expect((innerFn as (v: number) => number)(10)).toBe(50); - }); - - it("should handle async functions", () => { - const code = ` - export default Dynamics(async (x, y) => { - return x + y; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(5, 10); - expect(result).toBeInstanceOf(Promise); - }); - - it("should handle destructuring", () => { - const code = ` - export default Dynamics(({ a, b }) => { - return a + b; - }); - `; - const fn = compileUserCode<[{ a: number; b: number }]>(code, "Dynamics"); - const result = fn({ a: 5, b: 10 }); - expect(result).toBe(15); - }); - }); - - describe("constructor function name validation", () => { - it("should accept correct constructor function name", () => { - const code = ` - export default Dynamics((x) => x * 2); - `; - expect(() => compileUserCode(code, "Dynamics")).not.toThrow(); - }); - - it("should work with different constructor function names", () => { - const code = ` - export default Flow((x, y) => x + y); - `; - const fn = compileUserCode<[number, number]>(code, "Flow"); - const result = fn(5, 10); - expect(result).toBe(15); - }); - - it("should throw error when constructor function name doesn't match", () => { - const code = ` - export default WrongName((x) => x * 2); - `; - expect(() => compileUserCode(code, "Dynamics")).toThrow( - "Default export must use the constructor function 'Dynamics'", - ); - }); - - it("should throw error when default export is missing", () => { - const code = ` - const fn = Dynamics((x) => x * 2); - `; - expect(() => compileUserCode(code, "Dynamics")).toThrow( - "Module must have a default export", - ); - }); - - it("should throw error when export is not default", () => { - const code = ` - export const fn = Dynamics((x) => x * 2); - `; - expect(() => compileUserCode(code, "Dynamics")).toThrow( - "Module must have a default export", - ); - }); - }); - - describe("edge cases", () => { - it("should handle whitespace variations", () => { - const code = ` - export default Dynamics( (x) => x * 2 ) ; - `; - const fn = compileUserCode<[number]>(code, "Dynamics"); - const result = fn(5); - expect(result).toBe(10); - }); - - it("should handle multiline functions", () => { - const code = ` - export default Dynamics((x, y) => { - const step1 = x * 2; - const step2 = y + 5; - const step3 = step1 + step2; - return step3; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(3, 7); - expect(result).toBe(18); - }); - - it("should handle functions that return objects", () => { - const code = ` - export default Dynamics((x, y) => { - return { sum: x + y, product: x * y }; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(3, 4); - expect(result).toEqual({ sum: 7, product: 12 }); - }); - - it("should handle functions that return arrays", () => { - const code = ` - export default Dynamics((x, y) => { - return [x, y, x + y]; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(5, 10); - expect(result).toEqual([5, 10, 15]); - }); - - it("should handle empty parameter list", () => { - const code = ` - export default Dynamics(() => { - return 42; - }); - `; - const fn = compileUserCode<[]>(code, "Dynamics"); - const result = fn(); - expect(result).toBe(42); - }); - - it("should preserve function behavior", () => { - const code = ` - export default Dynamics(function(multiplier) { - return multiplier * 2; - }); - `; - const fn = compileUserCode<[number]>(code, "Dynamics"); - const result = fn(5); - expect(result).toBe(10); - }); - }); - - describe("error handling", () => { - it("should throw descriptive error for invalid code", () => { - const code = ` - export default Dynamics((x) => { - invalid syntax here - }); - `; - expect(() => compileUserCode(code, "Dynamics")).toThrow(); - }); - - it("should throw error if constructor function doesn't return a function", () => { - const code = ` - export default Dynamics(42); - `; - expect(() => compileUserCode(code, "Dynamics")).toThrow(); - }); - - it("should throw error for empty code", () => { - const code = ""; - expect(() => compileUserCode(code, "Dynamics")).toThrow( - "Module must have a default export", - ); - }); - }); - - describe("real-world differential equations", () => { - it("should handle exponential growth equation", () => { - const code = ` - export default Dynamics((population: number, rate: number): number => { - return population * rate; - }); - `; - const fn = compileUserCode<[number, number]>(code, "Dynamics"); - const result = fn(100, 1.05); - expect(result).toBe(105); - }); - - it("should handle Lotka-Volterra predator-prey model", () => { - const code = ` - export default Dynamics((prey: number, predator: number, alpha: number, beta: number): number[] => { - const dPrey = alpha * prey - beta * prey * predator; - const dPredator = -beta * predator + alpha * prey * predator; - return [dPrey, dPredator]; - }); - `; - const fn = compileUserCode<[number, number, number, number]>( - code, - "Dynamics", - ); - const result = fn(100, 10, 0.1, 0.02); - expect(Array.isArray(result)).toBe(true); - expect((result as number[]).length).toBe(2); - }); - - it("should handle state vector operations", () => { - const code = ` - export default Dynamics((state: number[][]): number[][] => { - return state.map(vector => - vector.map(component => component * 0.5) - ); - }); - `; - const fn = compileUserCode<[number[][]]>(code, "Dynamics"); - const result = fn([ - [10, 20], - [30, 40], - ]); - expect(result).toEqual([ - [5, 10], - [15, 20], - ]); - }); - }); - - describe("Uuid runtime injection", () => { - it("exposes Uuid.generate() and Uuid.from() sentinels", () => { - const code = ` - export default TransitionKernel(() => { - return { generated: Uuid.generate(), derived: Uuid.from("order-1") }; - }); - `; - const fn = compileUserCode(code, "TransitionKernel"); - expect(fn()).toEqual({ - generated: { __petrinautUuid: "generate" }, - derived: { __petrinautUuid: "from", value: "order-1" }, - }); - }); - - it("keeps Uuid available when Distribution is disabled", () => { - const code = ` - export default TransitionKernel(() => { - return Uuid.generate(); - }); - `; - const fn = compileUserCode(code, "TransitionKernel", { - enableDistribution: false, - }); - expect(fn()).toEqual({ __petrinautUuid: "generate" }); - }); - }); -}); diff --git a/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.ts b/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.ts deleted file mode 100644 index 30ae6749a54..00000000000 --- a/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/compile-user-code.ts +++ /dev/null @@ -1,121 +0,0 @@ -import * as Babel from "@babel/standalone"; - -import { distributionRuntimeCode } from "./distribution"; -import { uuidRuntimeCode } from "./uuid-runtime"; - -type CompileUserCodeOptions = { - /** - * Whether user code should receive the stochastic `Distribution` helper. - * Disable this for document modes where stochastic behaviour is unavailable. - * @default true - */ - enableDistribution?: boolean; -}; - -/** - * Strips TypeScript type annotations from code to make it executable JavaScript. - * Uses Babel standalone (browser-compatible) to properly parse and transform TypeScript code. - * - * @param code - TypeScript code with type annotations - * @returns JavaScript code without type annotations - */ -function stripTypeAnnotations(code: string): string { - const result = Babel.transform(code, { - filename: "input.ts", - presets: [ - ["typescript", { allowDeclareFields: true, onlyRemoveTypeImports: true }], - ], - }); - return result.code ?? ""; -} - -/** - * Compiles TypeScript/JavaScript module code into an executable function. - * Expects a JavaScript module with a default export using a specific constructor function. - * - * @param code - The TypeScript/JavaScript module code with a default export. - * Should follow the pattern: `export default ConstructorFn((a, b, c) => { ... })` - * Can include TypeScript type annotations which will be stripped. - * @param constructorFnName - The name of the constructor function that wraps the user code. - * This must match the function used in the default export. - * @returns A compiled function that can be executed - * - * @example - * ```typescript - * const code = ` - * export default Dynamics((a: number, b: number, c: number) => { - * return a + b + c; - * }); - * `; - * const fn = compileUserCode(code, "Dynamics"); - * const result = fn(1, 2, 3); // returns 6 - * ``` - */ -export function compileUserCode( - code: string, - constructorFnName: string, - options: CompileUserCodeOptions = {}, -): (...args: T) => unknown { - // Strip TypeScript type annotations and remove leading/trailing whitespace - const sanitizedCode = stripTypeAnnotations(code.trim()); - - // Verify that the code has a default export - const defaultExportRegex = /export\s+default\s+/; - if (!defaultExportRegex.test(sanitizedCode)) { - throw new Error( - `Module must have a default export. Expected pattern: export default ${constructorFnName}(...)`, - ); - } - - // Verify that the default export uses the specified constructor function - const constructorPattern = new RegExp( - `export\\s+default\\s+${constructorFnName}\\s*\\(`, - ); - if (!constructorPattern.test(sanitizedCode)) { - throw new Error( - `Default export must use the constructor function '${constructorFnName}'. Expected pattern: export default ${constructorFnName}(...)`, - ); - } - - try { - // Create a mock constructor function that extracts the wrapped function - const mockConstructor = ` - function ${constructorFnName}(fn) { - if (typeof fn !== 'function') { - throw new Error('${constructorFnName} expects a function as argument'); - } - return fn; - } - `; - - // Create an executable module-like environment - const executableCode = ` - ${options.enableDistribution === false ? "" : distributionRuntimeCode} - ${uuidRuntimeCode} - ${mockConstructor} - let __default_export__; - ${sanitizedCode.replace(/export\s+default\s+/, "__default_export__ = ")} - return __default_export__; - `; - - // Use Function constructor to create and execute the module - // eslint-disable-next-line no-new-func, @typescript-eslint/no-implied-eval, @typescript-eslint/no-unsafe-call - const compiledFunction = new Function(executableCode)() as ( - ...args: T - ) => unknown; - - if (typeof compiledFunction !== "function") { - throw new Error( - `Expected default export to be a function, got ${typeof compiledFunction}`, - ); - } - - return compiledFunction; - } catch (error) { - throw new Error( - `Failed to compile user code: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - } -} diff --git a/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/distribution.ts b/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/distribution.ts index c7f881cde23..cd8f8bd8b39 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/distribution.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/authoring/user-code/distribution.ts @@ -40,27 +40,6 @@ export function isDistribution(value: unknown): value is RuntimeDistribution { ); } -/** - * JavaScript source code that defines the Distribution namespace at runtime. - * Injected into the compiled user code execution context so that - * Distribution.Gaussian() and Distribution.Uniform() are available. - */ -export const distributionRuntimeCode = ` - function __addMap(dist) { - dist.map = function(fn) { - return __addMap({ __brand: "distribution", type: "mapped", inner: dist, fn: fn }); - }; - return dist; - } - var Distribution = { - Gaussian: function(mean, deviation) { - return __addMap({ __brand: "distribution", type: "gaussian", mean: mean, deviation: deviation }); - }, - Uniform: function(min, max) { - return __addMap({ __brand: "distribution", type: "uniform", min: min, max: max }); - }, - Lognormal: function(mu, sigma) { - return __addMap({ __brand: "distribution", type: "lognormal", mu: mu, sigma: sigma }); - } - }; -`; +// Distribution construction for compiled user code lives in +// `hir/instantiate.ts` (`hirDistributionRuntime`) — emitted programs call it +// through the injected `__dist` binding instead of injected source code. diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/buffer-transition.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/buffer-transition.ts new file mode 100644 index 00000000000..8b5f51fe910 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/buffer-transition.ts @@ -0,0 +1,131 @@ +/** + * Shared buffer-ABI transition execution used by both the single-run engine + * (`compute-possible-transition.ts`) and the Monte Carlo path + * (`monte-carlo/transition-effect.ts`). + * + * Buffer-ABI lambdas/kernels read token attributes at statically-resolved + * offsets from the frame's packed token floats — no per-combination record + * decoding and no object allocation. The per-transition scratch buffers + * (`slotBases`, `kernelStaging`, pending-distribution arrays) live on the + * `CompiledTransition` and are reused across evaluations; the engine is + * single-threaded per simulation instance. + */ +import { sampleDistribution } from "./sample-distribution"; + +import type { RuntimeDistribution } from "../authoring/user-code/distribution"; +import type { CompiledTransition, CompiledTransitionBuffer } from "./types"; + +/** + * Fills `slotBases` with the float base offset of each selected token, in + * slot order (per colored non-inhibitor input arc, `weight` slots each — + * matching the emitter's layout, see `hir/surface-context.ts`). + * + * Returns `false` when the combination does not match the expected slot + * count (stale artifact) — callers must fall back to the object path. + */ +export function fillSlotBases( + slotBases: Int32Array, + combinationIndices: readonly (readonly number[])[], + places: readonly { offset: number; dimensions: number }[], +): boolean { + let slot = 0; + for (const [placeIndex, tokenIndices] of combinationIndices.entries()) { + const place = places[placeIndex]!; + for (const tokenIndex of tokenIndices) { + if (slot >= slotBases.length) { + return false; + } + // eslint-disable-next-line no-param-reassign -- writes into the reusable scratch buffer + slotBases[slot] = place.offset + tokenIndex * place.dimensions; + slot += 1; + } + } + return slot === slotBases.length; +} + +/** + * Runs a buffer-ABI kernel: fills the staging floats, then samples deferred + * distribution values ordered by output float index — reproducing the legacy + * (place, token, element) sampling order, and therefore the exact RNG + * stream. Shared distribution objects keep one draw via their sample cache. + * + * Returns the advanced RNG state; the staging buffer holds the final values. + */ +export function executeBufferKernel( + buffer: CompiledTransitionBuffer, + kernelFn: NonNullable, + tokenValues: Float64Array, + rngState: number, +): number { + const { kernelStaging, pendingSlots, pendingDists } = buffer; + pendingSlots.length = 0; + pendingDists.length = 0; + + kernelFn(tokenValues, buffer.slotBases, kernelStaging, (index, dist) => { + pendingSlots.push(index); + pendingDists.push(dist); + }); + + if (pendingSlots.length === 0) { + return rngState; + } + + // Sample ordered by output float index (kernels emit sinks in evaluation + // order, which may differ from element order within a token). + const order = pendingSlots.map((_, index) => index); + order.sort((left, right) => pendingSlots[left]! - pendingSlots[right]!); + + let currentRngState = rngState; + for (const pendingIndex of order) { + const dist: RuntimeDistribution = pendingDists[pendingIndex]!; + const [sampled, nextRngState] = sampleDistribution(dist, currentRngState); + currentRngState = nextRngState; + kernelStaging[pendingSlots[pendingIndex]!] = sampled; + } + // Clear per-call sample caches so the next firing draws fresh values: + // emitted kernels construct fresh distribution objects per call, so this + // is defensive only. + pendingSlots.length = 0; + pendingDists.length = 0; + + return currentRngState; +} + +/** + * Converts the kernel staging floats into per-place token value arrays + * (place-major, colored output arcs in arc order; uncolored outputs get + * empty tuples). Later arcs to the same place overwrite earlier ones, + * matching the legacy object semantics where one kernel result array served + * every arc to that place. + */ +export function stagingToAddMap( + transition: CompiledTransition, + staging: Float64Array, +): Record { + const addMap: Record = {}; + let floatBase = 0; + + for (const outputPlace of transition.outputPlaces) { + if (!outputPlace.elements) { + addMap[outputPlace.placeId] = Array.from( + { length: outputPlace.weight }, + () => [], + ); + continue; + } + const dimensions = outputPlace.elements.length; + const tokenArrays: number[][] = []; + for (let tokenIndex = 0; tokenIndex < outputPlace.weight; tokenIndex += 1) { + const start = floatBase + tokenIndex * dimensions; + const values: number[] = []; + for (let dimension = 0; dimension < dimensions; dimension += 1) { + values.push(staging[start + dimension]!); + } + tokenArrays.push(values); + } + addMap[outputPlace.placeId] = tokenArrays; + floatBase += outputPlace.weight * dimensions; + } + + return addMap; +} diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation-hir.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation-hir.test.ts new file mode 100644 index 00000000000..62330fdff95 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation-hir.test.ts @@ -0,0 +1,208 @@ +/** + * End-to-end tests for HIR-compiled artifacts: buffer-ABI programs must + * produce bit-identical frames (including RNG stream evolution) to the + * object-convention programs, and simulations without artifacts must fail + * with a per-item error. + */ +import { describe, expect, it } from "vitest"; + +import { compileHirArtifacts } from "../../hir"; +import { buildSimulation } from "./build-simulation"; +import { computeNextFrame } from "./compute-next-frame"; + +import type { HirArtifacts } from "../../hir-runtime"; +import type { SDCPN } from "../../types/sdcpn"; +import type { SimulationInput, SimulationInstance } from "./types"; + +const sdcpn: SDCPN = { + types: [ + { + id: "type1", + name: "Particle", + iconSlug: "circle", + displayColor: "#FF0000", + elements: [ + { elementId: "e1", name: "x", type: "real" }, + { elementId: "e2", name: "v", type: "real" }, + { elementId: "e3", name: "generation", type: "integer" }, + ], + }, + ], + differentialEquations: [ + { + id: "de1", + name: "Oscillator", + colorId: "type1", + code: `export default Dynamics((tokens, parameters) => { + return tokens.map(({ x, v }) => { + return { x: v, v: -parameters.k * x }; + }); +});`, + }, + ], + parameters: [ + { + id: "param1", + name: "Spring constant", + variableName: "k", + type: "real", + defaultValue: "2", + }, + { + id: "param2", + name: "Rate", + variableName: "rate", + type: "real", + defaultValue: "5", + }, + ], + places: [ + { + id: "p1", + name: "Source", + colorId: "type1", + dynamicsEnabled: true, + differentialEquationId: "de1", + x: 0, + y: 0, + }, + { + id: "p2", + name: "Target", + colorId: "type1", + dynamicsEnabled: false, + differentialEquationId: null, + x: 0, + y: 0, + }, + ], + transitions: [ + { + id: "t1", + name: "Hop", + lambdaType: "stochastic", + inputArcs: [{ placeId: "p1", weight: 1, type: "standard" }], + outputArcs: [{ placeId: "p2", weight: 1 }], + lambdaCode: `export default Lambda((input, parameters) => { + const { x } = input.Source[0]; + if (x > 0) return parameters.rate; + return 0.5; +});`, + transitionKernelCode: `export default TransitionKernel((input, parameters) => { + const noise = Distribution.Gaussian(0, 0.1); + return { + Target: [ + { + x: input.Source[0].x, + v: noise.map((value) => value + input.Source[0].v), + generation: input.Source[0].generation + 1, + }, + ], + }; +});`, + x: 0, + y: 0, + }, + ], +}; + +function makeInput(hirArtifacts?: SimulationInput["hirArtifacts"]) { + return { + sdcpn, + initialMarking: { + p1: [ + { x: 1, v: 0, generation: 0 }, + { x: -0.5, v: 2, generation: 0 }, + ], + p2: [], + }, + parameterValues: { k: "2", rate: "5" }, + seed: 1234, + dt: 0.05, + maxTime: null, + hirArtifacts, + } satisfies SimulationInput; +} + +function runFrames(instance: SimulationInstance, count: number): number[][] { + let simulation = instance; + const frames: number[][] = []; + for (let step = 0; step < count; step++) { + const result = computeNextFrame(simulation); + simulation = result.simulation; + const frame = simulation.frames[simulation.currentFrameNumber]!; + frames.push([...new Float64Array(frame)]); + } + return frames; +} + +/** Strips buffer programs so only the object convention runs. */ +function objectOnly(artifacts: HirArtifacts): HirArtifacts { + const strip = ( + entries: Record, + ): Record => + Object.fromEntries( + Object.entries(entries).map(([key, value]) => [ + key, + { object: value.object }, + ]), + ); + return { + version: 2, + dynamics: strip(artifacts.dynamics), + lambdas: strip(artifacts.lambdas), + kernels: strip(artifacts.kernels), + }; +} + +describe("buildSimulation with HIR artifacts", () => { + it("compiles buffer and object programs for all three surfaces", () => { + const { artifacts, failures } = compileHirArtifacts(sdcpn); + expect(failures).toEqual([]); + expect(artifacts.dynamics.de1!.buffer).toBeDefined(); + expect(artifacts.dynamics.de1!.object).toBeDefined(); + expect(artifacts.lambdas.t1!.buffer).toBeDefined(); + expect(artifacts.lambdas.t1!.buffer!.inputSlotCount).toBe(1); + expect(artifacts.kernels.t1!.buffer).toBeDefined(); + expect(artifacts.kernels.t1!.buffer!.outputFloatCount).toBe(3); + }); + + it("buffer programs produce bit-identical frames to the object convention", () => { + const { artifacts } = compileHirArtifacts(sdcpn); + + const bufferRun = buildSimulation(makeInput(artifacts)); + const objectRun = buildSimulation(makeInput(objectOnly(artifacts))); + + // Sanity: the buffer run actually uses buffer programs. + const compiled = bufferRun.compiledTransitions.get("t1")!; + expect(compiled.buffer?.lambdaFn).not.toBeNull(); + expect(compiled.buffer?.kernelFn).not.toBeNull(); + expect(objectRun.compiledTransitions.get("t1")!.buffer).toBeNull(); + + expect(runFrames(bufferRun, 50)).toEqual(runFrames(objectRun, 50)); + }); + + it("throws a per-item error when artifacts are missing", () => { + expect(() => buildSimulation(makeInput())).toThrow(/has not been compiled/); + }); + + it("falls back to object programs when buffer metadata is stale", () => { + const { artifacts } = compileHirArtifacts(sdcpn); + const stale: HirArtifacts = { + ...artifacts, + lambdas: { + t1: { + ...artifacts.lambdas.t1!, + buffer: { ...artifacts.lambdas.t1!.buffer!, inputSlotCount: 99 }, + }, + }, + }; + const simulation = buildSimulation(makeInput(stale)); + expect( + simulation.compiledTransitions.get("t1")!.buffer?.lambdaFn, + ).toBeNull(); + // Still simulates correctly via the object path. + const reference = buildSimulation(makeInput(objectOnly(artifacts))); + expect(runFrames(simulation, 10)).toEqual(runFrames(reference, 10)); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.test.ts index c1bb7803dc5..8a4f4c642b5 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.test.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.test.ts @@ -1,12 +1,26 @@ import { describe, expect, it } from "vitest"; +import { compileHirArtifacts } from "../../hir"; import { compileSimulationFrameReader } from "../frames/frame-reader"; import { materializeEngineFrame } from "../frames/internal-frame"; -import { buildSimulation } from "./build-simulation"; +import { buildSimulation as buildSimulationRaw } from "./build-simulation"; import { decodePlaceTokens } from "./token-layout.test-helpers"; import type { SimulationInput } from "./types"; +/** buildSimulation with HIR artifacts compiled from the input's SDCPN (the + * engine no longer compiles user code itself). */ +function buildSimulation( + input: SimulationInput, +): ReturnType { + return buildSimulationRaw({ + ...input, + hirArtifacts: + input.hirArtifacts ?? + compileHirArtifacts(input.sdcpn, input.extensions).artifacts, + }); +} + describe("buildSimulation", () => { it("packs and decodes integer and boolean token attributes", () => { const input: SimulationInput = { @@ -273,7 +287,7 @@ describe("buildSimulation", () => { dt: 0.1, maxTime: null, }), - ).toThrow("Failed to compile Lambda function"); + ).toThrow("The Lambda code for `Move` has not been compiled"); }); it("does not expose Distribution to transition kernels when stochasticity is disabled", () => { @@ -553,7 +567,7 @@ describe("buildSimulation", () => { id: "diffeq1", name: "Differential Equation 1", colorId: "type1", - code: "export default Dynamics((placeValues, t) => { return new Float64Array([0, 0]); });", + code: "export default Dynamics((tokens) => tokens.map(() => ({})));", }, ], parameters: [], @@ -650,13 +664,13 @@ describe("buildSimulation", () => { id: "diffeq1", name: "Differential Equation 1", colorId: "type1", - code: "export default Dynamics((placeValues, t) => { return new Float64Array([0]); });", + code: "export default Dynamics((tokens) => tokens.map(() => ({})));", }, { id: "diffeq2", name: "Differential Equation 2", colorId: "type2", - code: "export default Dynamics((placeValues, t) => { return new Float64Array([0, 0]); });", + code: "export default Dynamics((tokens) => tokens.map(() => ({})));", }, ], parameters: [], @@ -817,7 +831,7 @@ describe("buildSimulation", () => { id: "diffeq1", name: "Differential Equation 1", colorId: "type1", - code: "export default Dynamics((placeValues, t) => { return new Float64Array([0]); });", + code: "export default Dynamics((tokens) => tokens.map(() => ({})));", }, ], parameters: [], @@ -868,7 +882,7 @@ describe("buildSimulation", () => { id: "diffeq1", name: "Differential Equation 1", colorId: "type1", - code: "export default Dynamics((placeValues, t) => { return new Float64Array([0, 0]); });", + code: "export default Dynamics((tokens) => tokens.map(() => ({})));", }, ], parameters: [], diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.ts index f73e05f36a8..a89392990a6 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/build-simulation.ts @@ -7,11 +7,16 @@ import { sanitizeSDCPNForExtensions, type PetrinautExtensionSettings, } from "../../extensions"; +import { + instantiateHirBufferDynamics, + instantiateHirBufferKernel, + instantiateHirBufferLambda, + instantiateHirUserFn, +} from "../../hir-runtime"; import { deriveDefaultParameterValues, mergeParameterValues, } from "../../parameter-values"; -import { compileUserCode } from "../authoring/user-code/compile-user-code"; import { isDistribution } from "../authoring/user-code/distribution"; import { createEngineFrame, @@ -32,9 +37,17 @@ import { } from "./token-layout"; import { coerceTokenRecord } from "./token-values"; +import type { + HirCompiledBufferKernel, + HirCompiledBufferLambda, + HirDynamicsArtifact, + HirKernelArtifact, + HirLambdaArtifact, +} from "../../hir-runtime"; import type { TokenRecord } from "../../types/sdcpn"; import type { CompiledTransition, + CompiledTransitionBuffer, DifferentialEquationFn, LambdaFn, ParameterValues, @@ -250,17 +263,89 @@ function getPlaceElements( return type.elements; } +/** + * Recovers the pre-flattening item id: flattening scopes ids as + * `instancePath::originalId`, while HIR artifacts are keyed by the original + * (root or subnet-local) id. + */ +function sourceItemId(flattenedId: string): string { + const separatorIndex = flattenedId.lastIndexOf("::"); + return separatorIndex === -1 + ? flattenedId + : flattenedId.slice(separatorIndex + 2); +} + +/** Error for items whose code has no compiled artifact (outside the + * supported subset, or artifacts not supplied). */ +function missingArtifactError( + kind: string, + name: string, + itemId: string, +): SDCPNItemError { + return new SDCPNItemError( + `The ${kind} code for \`${name}\` has not been compiled. Either the code is outside the supported Petrinaut code subset (check the Diagnostics tab for details) or the simulation was started without compiled artifacts.`, + itemId, + ); +} + +/** + * Expected `slotBases.length` for a transition: one slot per token of each + * colored, non-inhibitor input arc whose color has at least one element — + * must match the emitter's layout (see `hir/surface-context.ts`). + */ +function computeInputSlotCount( + transition: SimulationInput["sdcpn"]["transitions"][number], + placesMap: ReadonlyMap, + typesMap: ReadonlyMap, +): number { + let slots = 0; + for (const arc of transition.inputArcs) { + if (arc.type === "inhibitor") { + continue; + } + const placeId = getArcEndpointPlaceId(arc); + const place = placeId ? placesMap.get(placeId) : undefined; + const color = place?.colorId ? typesMap.get(place.colorId) : undefined; + if (color && color.elements.length > 0) { + slots += arc.weight; + } + } + return slots; +} + +/** Expected kernel staging float count: colored output arcs place-major. */ +function computeKernelStagingSize( + transition: SimulationInput["sdcpn"]["transitions"][number], + placesMap: ReadonlyMap, + typesMap: ReadonlyMap, +): number { + let floats = 0; + for (const arc of transition.outputArcs) { + const placeId = getArcEndpointPlaceId(arc); + const place = placeId ? placesMap.get(placeId) : undefined; + const color = place?.colorId ? typesMap.get(place.colorId) : undefined; + if (color) { + floats += arc.weight * color.elements.length; + } + } + return floats; +} + function createLambdaFn({ transition, sdcpn, extensions, parameterValues, + artifact, + expectedSlotCount, }: { transition: SimulationInput["sdcpn"]["transitions"][number]; sdcpn: SimulationInput["sdcpn"]; extensions: PetrinautExtensionSettings; parameterValues: ParameterValues; -}): LambdaFn { + artifact: HirLambdaArtifact | undefined; + expectedSlotCount: number; +}): { lambdaFn: LambdaFn; bufferLambdaFn: HirCompiledBufferLambda | null } { const availability = getTransitionLogicAvailability( transition, sdcpn, @@ -269,20 +354,29 @@ function createLambdaFn({ const lambdaType = getEffectiveTransitionLambdaType(transition, availability); if (!availability.lambda || transition.lambdaCode.trim() === "") { - return lambdaType === "stochastic" ? () => Infinity : () => true; + return { + lambdaFn: lambdaType === "stochastic" ? () => Infinity : () => true, + bufferLambdaFn: null, + }; } - try { - const userFn = compileUserCode<[TransitionTokenValues, ParameterValues]>( - transition.lambdaCode, - "Lambda", - { enableDistribution: extensions.stochasticity }, - ) as UserLambdaFn; + if (!artifact?.object) { + throw missingArtifactError("Lambda", transition.name, transition.id); + } - return (tokenValues) => userFn(tokenValues, parameterValues); + try { + const userFn = instantiateHirUserFn(artifact.object) as UserLambdaFn; + const bufferLambdaFn = + artifact.buffer && artifact.buffer.inputSlotCount === expectedSlotCount + ? instantiateHirBufferLambda(artifact.buffer.source, parameterValues) + : null; + return { + lambdaFn: (tokenValues) => userFn(tokenValues, parameterValues), + bufferLambdaFn, + }; } catch (error) { throw new SDCPNItemError( - `Failed to compile Lambda function for transition \`${ + `Failed to instantiate the compiled Lambda for transition \`${ transition.name }\`:\n\n${error instanceof Error ? error.message : String(error)}`, transition.id, @@ -295,12 +389,21 @@ function createTransitionKernelFn({ extensions, placesMap, parameterValues, + artifact, + expectedSlotCount, + expectedStagingSize, }: { transition: SimulationInput["sdcpn"]["transitions"][number]; extensions: PetrinautExtensionSettings; placesMap: ReadonlyMap; parameterValues: ParameterValues; -}): TransitionKernelFn { + artifact: HirKernelArtifact | undefined; + expectedSlotCount: number; + expectedStagingSize: number; +}): { + transitionKernelFn: TransitionKernelFn; + bufferKernelFn: HirCompiledBufferKernel | null; +} { const hasTypedOutputPlace = transition.outputArcs.some((arc) => { const placeId = getArcEndpointPlaceId(arc); const place = placeId ? placesMap.get(placeId) : undefined; @@ -308,17 +411,33 @@ function createTransitionKernelFn({ }); if (!hasTypedOutputPlace) { - return () => ({}); + return { transitionKernelFn: () => ({}), bufferKernelFn: null }; + } + + if (!artifact?.object) { + throw missingArtifactError( + "transition kernel", + transition.name, + transition.id, + ); } try { - const userFn = compileUserCode<[TransitionTokenValues, ParameterValues]>( - transition.transitionKernelCode, - "TransitionKernel", - { enableDistribution: extensions.stochasticity }, + const userFn = instantiateHirUserFn( + artifact.object, ) as UserTransitionKernelFn; - - return (tokenValues) => { + const bufferKernelFn = + artifact.buffer && + artifact.buffer.inputSlotCount === expectedSlotCount && + artifact.buffer.outputFloatCount === expectedStagingSize && + // Distribution outputs are rejected statically when stochasticity is + // off, but a stale artifact could still carry them — the object path + // performs the runtime check, so only it may run in that case. + extensions.stochasticity + ? instantiateHirBufferKernel(artifact.buffer.source, parameterValues) + : null; + + const transitionKernelFn: TransitionKernelFn = (tokenValues) => { const output = userFn(tokenValues, parameterValues); if (!extensions.stochasticity) { for (const [placeName, tokens] of Object.entries(output)) { @@ -335,9 +454,11 @@ function createTransitionKernelFn({ } return output; }; + + return { transitionKernelFn, bufferKernelFn }; } catch (error) { throw new SDCPNItemError( - `Failed to compile transition kernel for transition \`${ + `Failed to instantiate the compiled transition kernel for transition \`${ transition.name }\`:\n\n${error instanceof Error ? error.message : String(error)}`, transition.id, @@ -347,19 +468,66 @@ function createTransitionKernelFn({ function createCompiledTransition({ transition, + sdcpn, + extensions, placesMap, typesMap, arcPlaceNameOverrides, - lambdaFn, - transitionKernelFn, + parameterValues, + lambdaArtifact, + kernelArtifact, }: { transition: SimulationInput["sdcpn"]["transitions"][number]; + sdcpn: SimulationInput["sdcpn"]; + extensions: PetrinautExtensionSettings; placesMap: ReadonlyMap; typesMap: ReadonlyMap; arcPlaceNameOverrides: ReadonlyMap; - lambdaFn: LambdaFn; - transitionKernelFn: TransitionKernelFn; + parameterValues: ParameterValues; + lambdaArtifact: HirLambdaArtifact | undefined; + kernelArtifact: HirKernelArtifact | undefined; }): CompiledTransition { + const expectedSlotCount = computeInputSlotCount( + transition, + placesMap, + typesMap, + ); + const expectedStagingSize = computeKernelStagingSize( + transition, + placesMap, + typesMap, + ); + + const { lambdaFn, bufferLambdaFn } = createLambdaFn({ + transition, + sdcpn, + extensions, + parameterValues, + artifact: lambdaArtifact, + expectedSlotCount, + }); + const { transitionKernelFn, bufferKernelFn } = createTransitionKernelFn({ + transition, + extensions, + placesMap, + parameterValues, + artifact: kernelArtifact, + expectedSlotCount, + expectedStagingSize, + }); + + const buffer: CompiledTransitionBuffer | null = + bufferLambdaFn || bufferKernelFn + ? { + lambdaFn: bufferLambdaFn, + kernelFn: bufferKernelFn, + slotBases: new Int32Array(expectedSlotCount), + kernelStaging: new Float64Array(expectedStagingSize), + pendingSlots: [], + pendingDists: [], + } + : null; + return { id: transition.id, name: transition.name, @@ -424,6 +592,7 @@ function createCompiledTransition({ }), lambdaFn, transitionKernelFn, + buffer, }; } @@ -525,8 +694,6 @@ export function buildSimulation(input: SimulationInput): SimulationInstance { `Differential equation with ID ${place.differentialEquationId} referenced by place ${place.id} does not exist in SDCPN`, ); } - const { code } = differentialEquation; - try { if (!place.colorId) { continue; @@ -543,18 +710,39 @@ export function buildSimulation(input: SimulationInput): SimulationInstance { continue; } - const userFn = compileUserCode<[TokenRecord[], ParameterValues]>( - code, - "Dynamics", - { enableDistribution: extensions.stochasticity }, - ) as UserDifferentialEquationFn; + const placeParameterValues = + flattened.placeParameterValues.get(place.id) ?? parameterValues; + + const artifact: HirDynamicsArtifact | undefined = + input.hirArtifacts?.dynamics[sourceItemId(differentialEquation.id)]; + if (!artifact || (!artifact.buffer && !artifact.object)) { + throw missingArtifactError( + "dynamics", + differentialEquation.name, + place.id, + ); + } + + // Prefer the buffer-native program: it matches the + // DifferentialEquationFn signature directly and skips per-token record + // decoding entirely. + if (artifact.buffer) { + differentialEquationFns.set( + place.id, + instantiateHirBufferDynamics(artifact.buffer, placeParameterValues), + ); + continue; + } + + const userFn = instantiateHirUserFn( + artifact.object!, + ) as unknown as UserDifferentialEquationFn; differentialEquationFns.set( place.id, createDifferentialEquationFn({ placeId: place.id, tokenLayout: computeTokenSlotLayout(type.elements), - parameterValues: - flattened.placeParameterValues.get(place.id) ?? parameterValues, + parameterValues: placeParameterValues, userFn, stringPool, }), @@ -576,25 +764,18 @@ export function buildSimulation(input: SimulationInput): SimulationInstance { transition.id, createCompiledTransition({ transition, + sdcpn, + extensions, placesMap, typesMap, arcPlaceNameOverrides: flattened.arcPlaceNameOverrides, - lambdaFn: createLambdaFn({ - transition, - sdcpn, - extensions, - parameterValues: - flattened.transitionParameterValues.get(transition.id) ?? - parameterValues, - }), - transitionKernelFn: createTransitionKernelFn({ - transition, - extensions, - placesMap, - parameterValues: - flattened.transitionParameterValues.get(transition.id) ?? - parameterValues, - }), + parameterValues: + flattened.transitionParameterValues.get(transition.id) ?? + parameterValues, + lambdaArtifact: + input.hirArtifacts?.lambdas[sourceItemId(transition.id)], + kernelArtifact: + input.hirArtifacts?.kernels[sourceItemId(transition.id)], }), ); } diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-next-frame.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-next-frame.test.ts index 76189db8692..2b07faf3ce8 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-next-frame.test.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-next-frame.test.ts @@ -1,11 +1,26 @@ import { describe, expect, it } from "vitest"; -import { buildSimulation } from "./build-simulation"; +import { compileHirArtifacts } from "../../hir"; +import { buildSimulation as buildSimulationRaw } from "./build-simulation"; import { computeNextFrame } from "./compute-next-frame"; import { decodePlaceTokens } from "./token-layout.test-helpers"; import { parseUuid } from "./uuid"; import type { SDCPN } from "../../types/sdcpn"; +import type { SimulationInput as SimulationInputForArtifacts } from "./types"; + +/** buildSimulation with HIR artifacts compiled from the input's SDCPN (the + * engine no longer compiles user code itself). */ +function buildSimulation( + input: SimulationInputForArtifacts, +): ReturnType { + return buildSimulationRaw({ + ...input, + hirArtifacts: + input.hirArtifacts ?? + compileHirArtifacts(input.sdcpn, input.extensions).artifacts, + }); +} describe("computeNextFrame", () => { it("should compute next frame with dynamics and transitions", () => { diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-possible-transition.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-possible-transition.test.ts index 2685b7056a6..2b515e8ff96 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-possible-transition.test.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/compute-possible-transition.test.ts @@ -125,6 +125,7 @@ function makeCompiledTransitions({ }), lambdaFn, transitionKernelFn, + buffer: null, }, ]; }), diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/execute-transitions.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/execute-transitions.test.ts index 81ec85ad87a..a3b10653a88 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/engine/execute-transitions.test.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/execute-transitions.test.ts @@ -137,6 +137,7 @@ function makeCompiledTransitions({ }), lambdaFn, transitionKernelFn, + buffer: null, }, ]; }), diff --git a/libs/@hashintel/petrinaut-core/src/simulation/engine/types.ts b/libs/@hashintel/petrinaut-core/src/simulation/engine/types.ts index 99f7c8a3564..d25d1231a45 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/engine/types.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/engine/types.ts @@ -6,6 +6,11 @@ */ import type { PetrinautExtensionSettings } from "../../extensions"; +import type { + HirArtifacts, + HirCompiledBufferKernel, + HirCompiledBufferLambda, +} from "../../hir-runtime"; import type { Color, InputArcType, @@ -93,6 +98,26 @@ export type CompiledTransitionInputPlace = CompiledTransitionPlace & { arcType: InputArcType; }; +/** + * Buffer-ABI programs and reusable scratch for one transition. Present when + * at least one of lambda/kernel compiled to the buffer ABI; scratch is + * shared across evaluations (the engine is single-threaded per simulation). + */ +export type CompiledTransitionBuffer = { + /** `(tokenValues, slotBases) => number | boolean`, parameters pre-bound. */ + lambdaFn: HirCompiledBufferLambda | null; + /** `(tokenValues, slotBases, out, distSink) => void`, parameters pre-bound. */ + kernelFn: HirCompiledBufferKernel | null; + /** One float base offset per input token slot (see `hir/surface-context.ts` + * slot layout invariant); filled per enumerated combination. */ + slotBases: Int32Array; + /** Kernel output staging: colored output arcs place-major in arc order. */ + kernelStaging: Float64Array; + /** Deferred distribution sinks, reused per kernel call. */ + pendingSlots: number[]; + pendingDists: RuntimeDistribution[]; +}; + export type CompiledTransition = { id: string; name: string; @@ -100,6 +125,7 @@ export type CompiledTransition = { outputPlaces: readonly CompiledTransitionPlace[]; lambdaFn: LambdaFn; transitionKernelFn: TransitionKernelFn; + buffer: CompiledTransitionBuffer | null; }; /** @@ -120,6 +146,14 @@ export type SimulationInput = { dt: number; /** Maximum simulation time (immutable once set). Null means no limit. */ maxTime: number | null; + /** + * Optional precompiled HIR artifacts (see `compileHirArtifacts`). When an + * item has an artifact it is instantiated directly — skipping the Babel + * compile and, for dynamics, running buffer-native without per-token record + * decoding. Items without artifacts use the legacy compiler. Artifacts must + * be produced from the same SDCPN snapshot. + */ + hirArtifacts?: HirArtifacts; }; /** diff --git a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/monte-carlo-simulator.test.ts b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/monte-carlo-simulator.test.ts index e44ba829264..3c348665938 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/monte-carlo-simulator.test.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/monte-carlo-simulator.test.ts @@ -1,9 +1,24 @@ import { describe, expect, it } from "vitest"; +import { compileHirArtifacts } from "../../hir"; import { createMonteCarloUserDefinedMetric } from "./metrics"; -import { createMonteCarloSimulator } from "./monte-carlo-simulator"; +import { createMonteCarloSimulator as createMonteCarloSimulatorRaw } from "./monte-carlo-simulator"; import type { SDCPN } from "../../types/sdcpn"; +import type { MonteCarloSimulatorConfig } from "./types"; + +/** createMonteCarloSimulator with HIR artifacts compiled from the config's + * SDCPN (the engine no longer compiles user code itself). */ +function createMonteCarloSimulator( + config: MonteCarloSimulatorConfig, +): ReturnType { + return createMonteCarloSimulatorRaw({ + ...config, + hirArtifacts: + config.hirArtifacts ?? + compileHirArtifacts(config.sdcpn, config.extensions).artifacts, + }); +} const sdcpn: SDCPN = { types: [ diff --git a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/run-state.ts b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/run-state.ts index f35f2135025..97624284ae8 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/run-state.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/run-state.ts @@ -95,6 +95,7 @@ export function createRunState( seed, dt: config.dt, maxTime: config.maxTime, + hirArtifacts: config.hirArtifacts, }); const initialFrame = simulation.frames[0]; if (!initialFrame) { diff --git a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/runtime/experiment.ts b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/runtime/experiment.ts index 4707570bcdb..daed29c3f20 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/runtime/experiment.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/runtime/experiment.ts @@ -7,6 +7,7 @@ import { createMonteCarloSimulator } from "../monte-carlo-simulator"; import type { AbortSignalLike } from "../../../environment"; import type { PetrinautExtensionSettings } from "../../../extensions"; +import type { HirArtifacts } from "../../../hir-runtime"; import type { EventStream } from "../../../instance"; import type { ReadableStore } from "../../../store"; import type { SDCPN } from "../../../types/sdcpn"; @@ -53,6 +54,9 @@ type CreateMonteCarloExperimentBaseConfig = { seed: number; dt: number; maxTime: number; + /** Precompiled HIR artifacts (`compileHirArtifacts`) — required for any + * dynamics/lambda/kernel user code in the net. */ + hirArtifacts?: HirArtifacts; runCount: number; batchSize?: number; signal?: AbortSignalLike; @@ -307,6 +311,7 @@ function createLocalMonteCarloExperiment( seed: config.seed, dt: config.dt, maxTime: config.maxTime, + hirArtifacts: config.hirArtifacts, runCount: config.runCount, metrics: userMetrics, }); @@ -628,6 +633,7 @@ export function createMonteCarloExperiment( seed: config.seed, dt: config.dt, maxTime: config.maxTime, + hirArtifacts: config.hirArtifacts, runCount: config.runCount, batchSize: config.batchSize, metricSpecs: "metricSpecs" in config ? config.metricSpecs : undefined, diff --git a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/types.ts b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/types.ts index bf7f66c655e..660f6ff6522 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/types.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/types.ts @@ -1,4 +1,5 @@ import type { PetrinautExtensionSettings } from "../../extensions"; +import type { HirArtifacts } from "../../hir-runtime"; import type { SDCPN } from "../../types/sdcpn"; import type { InitialMarking } from "../api"; import type { SimulationCompletionReason } from "../engine/compute-next-frame"; @@ -16,6 +17,9 @@ export type MonteCarloRunConfig = { export type MonteCarloSimulatorConfig = { sdcpn: SDCPN; extensions?: PetrinautExtensionSettings; + /** Precompiled HIR artifacts (`compileHirArtifacts`) — required for any + * dynamics/lambda/kernel user code in the net. */ + hirArtifacts?: HirArtifacts; runCount: number; initialMarking: InitialMarking; parameterValues?: Record; diff --git a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/messages.ts b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/messages.ts index 9b21c6fe472..4af3a65a41b 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/messages.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/messages.ts @@ -1,4 +1,5 @@ import type { PetrinautExtensionSettings } from "../../../extensions"; +import type { HirArtifacts } from "../../../hir-runtime"; import type { SDCPN } from "../../../types/sdcpn"; import type { InitialMarking } from "../../api"; import type { @@ -16,6 +17,9 @@ export type MonteCarloInitMessage = { seed: number; dt: number; maxTime: number; + /** Precompiled HIR artifacts (`compileHirArtifacts`) — required for any + * dynamics/lambda/kernel user code in the net. */ + hirArtifacts?: HirArtifacts; runCount: number; batchSize?: number; metricSpecs?: readonly MonteCarloMetricSpec[]; diff --git a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/monte-carlo.worker.ts b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/monte-carlo.worker.ts index 76409855541..4d3e29827dc 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/monte-carlo.worker.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/monte-carlo/worker/monte-carlo.worker.ts @@ -157,6 +157,7 @@ function initialize(message: MonteCarloInitMessage): void { seed: message.seed, dt: message.dt, maxTime: message.maxTime, + hirArtifacts: message.hirArtifacts, runCount: message.runCount, metrics: userMetrics, }); diff --git a/libs/@hashintel/petrinaut-core/src/simulation/runtime/simulation.ts b/libs/@hashintel/petrinaut-core/src/simulation/runtime/simulation.ts index fbe6b352855..6d30061e430 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/runtime/simulation.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/runtime/simulation.ts @@ -261,6 +261,7 @@ export function createSimulation( seed: config.seed, dt: config.dt, maxTime: config.maxTime, + hirArtifacts: config.hirArtifacts, maxFramesAhead: config.backpressure?.maxFramesAhead, batchSize: config.backpressure?.batchSize, }); diff --git a/libs/@hashintel/petrinaut-core/src/simulation/worker/messages.ts b/libs/@hashintel/petrinaut-core/src/simulation/worker/messages.ts index 45ee1f14835..d32263a1e34 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/worker/messages.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/worker/messages.ts @@ -5,6 +5,7 @@ */ import type { PetrinautExtensionSettings } from "../../extensions"; +import type { HirArtifacts } from "../../hir-runtime"; import type { SDCPN } from "../../types/sdcpn"; import type { InitialMarking } from "../api"; import type { SimulationFramePayload } from "./frame-payload"; @@ -33,6 +34,9 @@ export type InitMessage = { dt: number; /** Maximum simulation time (immutable once set). Null means no limit. */ maxTime: number | null; + /** Precompiled HIR artifacts (`compileHirArtifacts`) — required for any + * dynamics/lambda/kernel user code in the net. */ + hirArtifacts?: HirArtifacts; /** Maximum frames the worker can compute ahead before waiting for ack (backpressure) */ maxFramesAhead?: number; /** Number of frames to compute in each batch before checking for messages */ diff --git a/libs/@hashintel/petrinaut-core/src/simulation/worker/simulation.worker.ts b/libs/@hashintel/petrinaut-core/src/simulation/worker/simulation.worker.ts index 52f01510d1c..226324dfe96 100644 --- a/libs/@hashintel/petrinaut-core/src/simulation/worker/simulation.worker.ts +++ b/libs/@hashintel/petrinaut-core/src/simulation/worker/simulation.worker.ts @@ -193,6 +193,7 @@ workerRuntime.onMessage((message) => { seed: message.seed, dt: message.dt, maxTime: message.maxTime, + hirArtifacts: message.hirArtifacts, }); // Configure backpressure from init message or use defaults diff --git a/libs/@hashintel/petrinaut/docs/petri-net-extensions.md b/libs/@hashintel/petrinaut/docs/petri-net-extensions.md index 411f76167a3..c849696562d 100644 --- a/libs/@hashintel/petrinaut/docs/petri-net-extensions.md +++ b/libs/@hashintel/petrinaut/docs/petri-net-extensions.md @@ -221,8 +221,21 @@ Use read arcs when a transition needs to inspect shared state, permission tokens ## Diagnostics -The **Diagnostics** tab in the bottom panel shows TypeScript errors in your code (dynamics, firing rate, kernels, visualizers), grouped by entity. Click a diagnostic to select the relevant entity and see the error in context. +The **Diagnostics** tab in the bottom panel shows problems in your code (dynamics, firing rate, kernels, visualizers), grouped by entity. Click a diagnostic to select the relevant entity and see the problem in context. Petrinaut only reports diagnostics for code surfaces that are active for the current document and graph shape. For example, a hidden firing-time editor or a transition with no coloured outputs will not produce lambda or kernel diagnostics. -Diagnostics must be resolved before running a simulation -- pressing Play with unresolved errors opens the Diagnostics tab instead of starting the simulation. +Diagnostics come in two flavours: + +- **Errors** (red) -- TypeScript type/syntax errors, plus code that falls outside Petrinaut's supported code subset (see below). These must be resolved before running a simulation: pressing Play with unresolved errors opens the Diagnostics tab instead of starting the simulation. The status indicator in the bottom toolbar shows a red cross while errors remain. +- **Warnings and hints** (amber indicator) -- semantic advice that does not block simulation. Examples: `Math.random()` makes runs non-reproducible (prefer `Distribution.Uniform`), a firing rate that is always 0 so the transition can never fire, a `const` binding that is never used, or one distribution feeding several output attributes (they all receive the same sampled value). + +### Supported code subset + +Dynamics, firing-rate and transition-kernel code is compiled by Petrinaut's own compiler, which accepts a focused expression subset of TypeScript rather than arbitrary programs: + +- `const` bindings (including destructuring like `const { a, b } = parameters` or `const [first] = input.Place`), a final `return`, and guard clauses (`if (condition) return value;`). +- Arithmetic, comparisons, boolean logic, ternaries, and `Math.*` functions. +- Token access (`input.Place[0].attr`, `.length`), `.map(...)` over token arrays, and `Distribution.*` constructors (with `.map` transforms). + +Loops, `let`/`var`, object spread and arbitrary function calls are rejected with an error pointing at the offending code and suggesting the idiomatic alternative. This is what lets Petrinaut analyze your model (e.g. which parameters a rate depends on) and compile it to fast code that reads token values directly from the simulation's internal buffers. Metric and scenario code is not affected by this subset. diff --git a/libs/@hashintel/petrinaut/docs/simulation.md b/libs/@hashintel/petrinaut/docs/simulation.md index 1927de9535d..e68ca36ae94 100644 --- a/libs/@hashintel/petrinaut/docs/simulation.md +++ b/libs/@hashintel/petrinaut/docs/simulation.md @@ -59,7 +59,7 @@ Press **Play** in the bottom toolbar. The simulation: If you need multiple runs of the same configuration -- e.g. to compare a stochastic model under different conditions -- use a Monte Carlo [experiment](experiments.md). -If there are unresolved [diagnostics](petri-net-extensions.md#diagnostics) (code errors), pressing Play opens the Diagnostics tab instead of starting the simulation. Fix all errors first. +If there are unresolved error-severity [diagnostics](petri-net-extensions.md#diagnostics) (code errors), pressing Play opens the Diagnostics tab instead of starting the simulation. Fix all errors first -- warnings and hints don't block simulation. simulation-settings diff --git a/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx b/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx index 6c5b298629e..be95ce646c5 100644 --- a/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/experiments/provider.tsx @@ -16,6 +16,7 @@ import { createMonteCarloWorker } from "@hashintel/petrinaut-core/workers/monte- import { useBlockWindowClose } from "../hooks/use-block-window-close"; import { useLatest } from "../hooks/use-latest"; import { useStableCallback } from "../hooks/use-stable-callback"; +import { LanguageClientContext } from "../lsp/context"; import { NotificationsContext } from "../notifications/context"; import { SDCPNContext } from "../state/sdcpn-context"; import { @@ -163,6 +164,7 @@ export const ExperimentsProvider: React.FC = ({ workerFactory, }) => { const { extensions, petriNetDefinition } = use(SDCPNContext); + const { requestHirArtifacts } = use(LanguageClientContext); const { addNotification } = use(NotificationsContext); const petriNetDefinitionRef = useLatest(petriNetDefinition); const extensionsRef = useLatest(extensions); @@ -360,18 +362,25 @@ export const ExperimentsProvider: React.FC = ({ pendingRegistrationsRef.current.set(experimentId, { abortController }); const initializeExperiment = async () => { - const experimentConfigBase = { - sdcpn: experimentSdcpn, - extensions: extensionsRef.current, - initialMarking, - parameterValues, - seed: input.seed, - dt: input.dt, - maxTime: input.maxTime, - runCount: input.runCount, - }; - try { + // Compile the net's user code to HIR artifacts in the language + // worker — the simulation engine has no compiler of its own. + const { artifacts } = await requestHirArtifacts( + experimentSdcpn, + extensionsRef.current, + ); + const experimentConfigBase = { + sdcpn: experimentSdcpn, + extensions: extensionsRef.current, + initialMarking, + parameterValues, + seed: input.seed, + dt: input.dt, + maxTime: input.maxTime, + hirArtifacts: artifacts, + runCount: input.runCount, + }; + const handle = await createMonteCarloExperiment({ ...experimentConfigBase, createWorker: workerFactoryRef.current, diff --git a/libs/@hashintel/petrinaut/src/react/lsp/context.ts b/libs/@hashintel/petrinaut/src/react/lsp/context.ts index a7941800f1e..c81f11bc88d 100644 --- a/libs/@hashintel/petrinaut/src/react/lsp/context.ts +++ b/libs/@hashintel/petrinaut/src/react/lsp/context.ts @@ -4,8 +4,11 @@ import type { CompletionList, Diagnostic, DocumentUri, + HirCompileResult, Hover, + PetrinautExtensionSettings, Position, + SDCPN, SignatureHelp, } from "@hashintel/petrinaut-core"; import type { @@ -18,6 +21,12 @@ export interface LanguageClientContextValue { diagnosticsByUri: Map; /** Total number of diagnostics across all documents. */ totalDiagnosticsCount: number; + /** + * Error-severity diagnostics only. Warnings/hints (e.g. HIR semantic + * lints) are included in `totalDiagnosticsCount` but not here — use this + * to decide whether the net can be simulated. + */ + errorDiagnosticsCount: number; /** Notify the server that a document's content changed. */ notifyDocumentChanged: (uri: DocumentUri, text: string) => void; /** Request completions at a position within a document. */ @@ -32,6 +41,15 @@ export interface LanguageClientContextValue { uri: DocumentUri, position: Position, ) => Promise; + /** + * Compile the SDCPN's user code to HIR artifacts (in the language worker). + * Required before starting simulations/experiments — the engine has no + * compiler of its own. + */ + requestHirArtifacts: ( + sdcpn: SDCPN, + extensions?: PetrinautExtensionSettings, + ) => Promise; /** Initialize a temporary scenario editing session. */ initializeScenarioSession: (params: ScenarioSessionParams) => void; /** Update a scenario editing session. */ @@ -49,10 +67,16 @@ export interface LanguageClientContextValue { const DEFAULT_CONTEXT_VALUE: LanguageClientContextValue = { diagnosticsByUri: new Map(), totalDiagnosticsCount: 0, + errorDiagnosticsCount: 0, notifyDocumentChanged: () => {}, requestCompletion: () => Promise.resolve({ isIncomplete: false, items: [] }), requestHover: () => Promise.resolve(null), requestSignatureHelp: () => Promise.resolve(null), + requestHirArtifacts: () => + Promise.resolve({ + artifacts: { version: 2, dynamics: {}, lambdas: {}, kernels: {} }, + failures: [], + }), initializeScenarioSession: () => {}, updateScenarioSession: () => {}, killScenarioSession: () => {}, diff --git a/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx b/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx index c9b8694d77a..0cf2e64836c 100644 --- a/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/lsp/provider.tsx @@ -17,6 +17,7 @@ import { LanguageClientContext } from "./context"; const EMPTY_DIAGNOSTICS_SNAPSHOT: DiagnosticsSnapshot = { byUri: new Map(), total: 0, + errorCount: 0, }; const EMPTY_DIAGNOSTICS_STORE: ReadableStore = { @@ -81,9 +82,11 @@ export const LanguageClientProvider: React.FC<{ // Subscribe to diagnostics from the client. Use an empty fallback store // before the client is created so hook order stays stable. - const { byUri: diagnosticsByUri, total: totalDiagnosticsCount } = useStore( - client?.diagnostics ?? EMPTY_DIAGNOSTICS_STORE, - ); + const { + byUri: diagnosticsByUri, + total: totalDiagnosticsCount, + errorCount: errorDiagnosticsCount, + } = useStore(client?.diagnostics ?? EMPTY_DIAGNOSTICS_STORE); // Before the client lands (StrictMode's first effect cycle gets cleaned // up before the worker is wired), fall through to LanguageClientContext's @@ -96,10 +99,12 @@ export const LanguageClientProvider: React.FC<{ const value = { diagnosticsByUri, totalDiagnosticsCount, + errorDiagnosticsCount, notifyDocumentChanged: client.notifyDocumentChanged, requestCompletion: client.requestCompletion, requestHover: client.requestHover, requestSignatureHelp: client.requestSignatureHelp, + requestHirArtifacts: client.requestHirArtifacts, initializeScenarioSession: client.initializeScenarioSession, updateScenarioSession: client.updateScenarioSession, killScenarioSession: client.killScenarioSession, diff --git a/libs/@hashintel/petrinaut/src/react/simulation/provider.tsx b/libs/@hashintel/petrinaut/src/react/simulation/provider.tsx index ec060032cee..4a7aa4c80b0 100644 --- a/libs/@hashintel/petrinaut/src/react/simulation/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/simulation/provider.tsx @@ -15,6 +15,7 @@ import { createSimulationWorker } from "@hashintel/petrinaut-core/workers/simula import { deriveDefaultParameterValues } from "../hooks/use-default-parameter-values"; import { useLatest } from "../hooks/use-latest"; import { useStableCallback } from "../hooks/use-stable-callback"; +import { LanguageClientContext } from "../lsp/context"; import { NotificationsContext } from "../notifications/context"; import { SDCPNContext } from "../state/sdcpn-context"; import { useStore } from "../use-store"; @@ -159,6 +160,7 @@ export const SimulationProvider: React.FC = ({ workerFactory, }) => { const sdcpnContext = use(SDCPNContext); + const { requestHirArtifacts } = use(LanguageClientContext); const { extensions, petriNetDefinition } = sdcpnContext; const { addNotification } = use(NotificationsContext); @@ -374,9 +376,16 @@ export const SimulationProvider: React.FC = ({ let sim: Simulation; try { + // Compile the net's user code to HIR artifacts in the language worker — + // the simulation engine has no compiler of its own. + const { artifacts } = await requestHirArtifacts( + simulationSdcpn, + extensionsRef.current, + ); sim = await createSimulation({ sdcpn: simulationSdcpn, extensions: extensionsRef.current, + hirArtifacts: artifacts, // eslint-disable-next-line no-use-before-define -- closure; ref is defined later in render initialMarking: effectiveInitialMarkingRef.current, // eslint-disable-next-line no-use-before-define -- closure; ref is defined later in render diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/bottom-bar.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/bottom-bar.tsx index f91bc492837..aa264351c90 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/bottom-bar.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/bottom-bar.tsx @@ -89,8 +89,10 @@ export const BottomBar: React.FC = ({ toggleAiAssistant, } = use(EditorContext); - const { totalDiagnosticsCount } = use(LanguageClientContext); - const hasDiagnostics = totalDiagnosticsCount > 0; + // Only error-severity diagnostics block simulation — warnings and hints + // (e.g. HIR semantic lints) are informational. + const { errorDiagnosticsCount } = use(LanguageClientContext); + const hasDiagnostics = errorDiagnosticsCount > 0; const { activeSubnetId } = use(ActiveNetContext); const isInSubnet = activeSubnetId !== null; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/diagnostics-indicator.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/diagnostics-indicator.tsx index f369665bc6b..5773fc2cb6d 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/diagnostics-indicator.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/components/BottomBar/diagnostics-indicator.tsx @@ -21,6 +21,10 @@ const iconContainerStyle = cva({ backgroundColor: "[rgba(239, 68, 68, 0.1)]", color: "[#dc2626]", }, + warning: { + backgroundColor: "[rgba(245, 158, 11, 0.1)]", + color: "[#d97706]", + }, success: { backgroundColor: "[rgba(34, 197, 94, 0.1)]", color: "[#16a34a]", @@ -42,22 +46,26 @@ interface DiagnosticsIndicatorProps { /** * DiagnosticsIndicator shows the current SDCPN validation status. * - Green check icon if no issues - * - Red cross icon with count if issues found + * - Amber icon with count if only warnings/hints (simulation still allowed) + * - Red cross icon with count if errors found */ export const DiagnosticsIndicator: React.FC = ({ onClick, isExpanded, }) => { - const { totalDiagnosticsCount } = use(LanguageClientContext); + const { totalDiagnosticsCount, errorDiagnosticsCount } = use( + LanguageClientContext, + ); - const hasErrors = totalDiagnosticsCount > 0; + const hasErrors = errorDiagnosticsCount > 0; + const hasIssues = totalDiagnosticsCount > 0; return ( = ({ >
- {hasErrors ? ( + {hasIssues ? ( <> {totalDiagnosticsCount} From 057870fb4064650be2e5dafb33c3292acd594c42 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Sat, 4 Jul 2026 02:43:21 +0200 Subject: [PATCH 02/12] Add HIR playground Storybook story MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Interactive playground for the HIR pipeline: type dynamics/lambda/kernel TypeScript on the left, and inspect everything it compiles to on the right — the spanned HIR tree, diagnostics with exact ranges, the inferred return type, dependency sets, the distribution DAG (derivation edges, output sinks, shared draws), and both emitted programs (buffer-ABI fast path and object-convention fallback). The model schema (parameters, places, attribute types) is editable JSON and re-checks live. Exposes the compiler via new `@hashintel/petrinaut-core/hir` and `/hir-runtime` entries (typescript declared as an optional peer — only the ./hir entry needs it). Co-Authored-By: Claude Fable 5 --- libs/@hashintel/petrinaut-core/package.json | 16 + libs/@hashintel/petrinaut-core/vite.config.ts | 9 +- .../src/ui/hir-playground.stories.tsx | 586 ++++++++++++++++++ yarn.lock | 2 - 4 files changed, 609 insertions(+), 4 deletions(-) create mode 100644 libs/@hashintel/petrinaut/src/ui/hir-playground.stories.tsx diff --git a/libs/@hashintel/petrinaut-core/package.json b/libs/@hashintel/petrinaut-core/package.json index 352b36f113f..a2b43724781 100644 --- a/libs/@hashintel/petrinaut-core/package.json +++ b/libs/@hashintel/petrinaut-core/package.json @@ -34,6 +34,14 @@ "types": "./dist/examples/index.d.d.ts", "import": "./dist/examples/index.js" }, + "./hir": { + "types": "./dist/hir.d.d.ts", + "import": "./dist/hir.js" + }, + "./hir-runtime": { + "types": "./dist/hir-runtime.d.d.ts", + "import": "./dist/hir-runtime.js" + }, "./workers/lsp": { "types": "./dist/workers/lsp.d.d.ts", "import": "./dist/workers/lsp.js" @@ -76,5 +84,13 @@ "typescript": "5.9.3", "vite": "8.1.0", "vitest": "4.1.8" + }, + "peerDependencies": { + "typescript": ">=5.5" + }, + "peerDependenciesMeta": { + "typescript": { + "optional": true + } } } diff --git a/libs/@hashintel/petrinaut-core/vite.config.ts b/libs/@hashintel/petrinaut-core/vite.config.ts index df29a6553c9..23aebb484f1 100644 --- a/libs/@hashintel/petrinaut-core/vite.config.ts +++ b/libs/@hashintel/petrinaut-core/vite.config.ts @@ -13,8 +13,12 @@ export default defineConfig(({ command }) => ({ entry: { index: resolve(packageRoot, "src/index.ts"), // Dedicated edge-safe entry exposing only the AI prompt + tool schemas. - // Needed to avoid pulling in heavy and edge-incompatible deps (e.g. @babel/standalone) ai: resolve(packageRoot, "src/ai.ts"), + // HIR compiler (bundles the TypeScript frontend — heavy; used by the + // LSP worker internally and by tooling/playgrounds). + hir: resolve(packageRoot, "src/hir.ts"), + // Dependency-free instantiation of compiled HIR artifacts. + "hir-runtime": resolve(packageRoot, "src/hir-runtime.ts"), "examples/index": resolve(packageRoot, "src/examples/index.ts"), "workers/lsp": resolve(packageRoot, "src/workers/lsp.ts"), "workers/monte-carlo": resolve( @@ -28,7 +32,8 @@ export default defineConfig(({ command }) => ({ }, rolldownOptions: { external: [ - "@babel/standalone", + // Peer (optional): only the ./hir compiler entry needs it. + "typescript", "elkjs", "immer", "uuid", diff --git a/libs/@hashintel/petrinaut/src/ui/hir-playground.stories.tsx b/libs/@hashintel/petrinaut/src/ui/hir-playground.stories.tsx new file mode 100644 index 00000000000..ca9f8bcafe5 --- /dev/null +++ b/libs/@hashintel/petrinaut/src/ui/hir-playground.stories.tsx @@ -0,0 +1,586 @@ +import { useState } from "react"; + +import { css } from "@hashintel/ds-helpers/css"; +import { + emitBufferDynamicsJs, + emitBufferKernelJs, + emitBufferLambdaJs, + emitUserFunctionJs, + formatHirType, + lintHirUserCode, +} from "@hashintel/petrinaut-core/hir"; + +import { CodeEditor } from "./monaco/code-editor"; +import { MonacoProvider } from "./monaco/provider"; + +import type { + HirAnalysis, + HirDiagnostic, + HirFunction, + HirSurfaceContext, + HirSurfaceKind, + HirTokenElementInfo, + HirTypecheckResult, +} from "@hashintel/petrinaut-core/hir"; +import type { Meta, StoryObj } from "@storybook/react-vite"; + +// -- Playground schema (hand-editable, derives the HirSurfaceContext) --------- + +type PlaygroundPlace = { + name: string; + tokenCount: number; + elements: HirTokenElementInfo[]; +}; + +type PlaygroundSchema = { + parameters: { name: string; type: "real" | "integer" | "boolean" }[]; + /** Attributes of the colour a dynamics equation is attached to. */ + elements: HirTokenElementInfo[]; + /** Colored input places, in arc order (lambda/kernel). */ + inputPlaces: PlaygroundPlace[]; + /** Colored output places, in arc order (kernel). */ + outputPlaces: PlaygroundPlace[]; + lambdaType: "stochastic" | "predicate"; + stochasticity: boolean; +}; + +function toArcSlots(places: PlaygroundPlace[]) { + let slotStart = 0; + return places.map((place, index) => { + const slot = { + name: place.name, + colorId: `color-${index}`, + elements: place.elements, + tokenCount: place.tokenCount, + slotStart, + }; + slotStart += place.tokenCount; + return slot; + }); +} + +function toContext( + surface: HirSurfaceKind, + schema: PlaygroundSchema, +): HirSurfaceContext { + switch (surface) { + case "dynamics": + return { + surface, + parameters: schema.parameters, + elements: schema.elements, + }; + case "lambda": + return { + surface, + parameters: schema.parameters, + inputPlaces: toArcSlots(schema.inputPlaces), + inputSlots: toArcSlots(schema.inputPlaces), + lambdaType: schema.lambdaType, + }; + case "kernel": + return { + surface, + parameters: schema.parameters, + inputPlaces: toArcSlots(schema.inputPlaces), + inputSlots: toArcSlots(schema.inputPlaces), + outputPlaces: toArcSlots(schema.outputPlaces), + outputSlots: toArcSlots(schema.outputPlaces), + stochasticity: schema.stochasticity, + }; + } +} + +// -- Presets ------------------------------------------------------------------- + +const DEFAULT_SCHEMA: PlaygroundSchema = { + parameters: [ + { name: "k", type: "real" }, + { name: "rate", type: "real" }, + { name: "sigma", type: "real" }, + { name: "threshold", type: "real" }, + ], + elements: [ + { name: "x", type: "real" }, + { name: "v", type: "real" }, + { name: "generation", type: "integer" }, + { name: "alive", type: "boolean" }, + ], + inputPlaces: [ + { + name: "Pool", + tokenCount: 2, + elements: [ + { name: "x", type: "real" }, + { name: "v", type: "real" }, + { name: "alive", type: "boolean" }, + ], + }, + { + name: "Fuel", + tokenCount: 1, + elements: [{ name: "level", type: "real" }], + }, + ], + outputPlaces: [ + { + name: "Out", + tokenCount: 1, + elements: [ + { name: "x", type: "real" }, + { name: "generation", type: "integer" }, + { name: "alive", type: "boolean" }, + ], + }, + ], + lambdaType: "stochastic", + stochasticity: true, +}; + +const CODE_PRESETS: Record = { + dynamics: `export default Dynamics((tokens, parameters) => { + const stiffness = parameters.k; + + return tokens.map(({ x, v, alive }) => { + return { + x: alive ? v : 0, + v: -stiffness * x, + }; + }); +}); +`, + lambda: `export default Lambda((input, parameters) => { + const { rate, threshold } = parameters; + const { x, alive } = input.Pool[0]; + + if (!alive) return 0; + if (input.Fuel[0].level < threshold) return 0; + + return rate * Math.max(0.1, x); +}); +`, + kernel: `export default TransitionKernel((input, parameters) => { + // One draw shared by two attributes (same sample at fire time). + const noise = Distribution.Gaussian(0, parameters.sigma); + + return { + Out: [{ + x: noise.map((value) => input.Pool[0].x + value), + generation: input.Pool[0].v > 0 ? 1 : 0, + alive: input.Pool[0].alive, + }], + }; +}); +`, +}; + +// -- The pipeline result (recomputed every render — lowering is ~1ms) ---------- + +type PipelineResult = { + context: HirSurfaceContext | null; + schemaError: string | null; + fn: HirFunction | null; + diagnostics: HirDiagnostic[]; + typecheck: HirTypecheckResult | null; + analysis: HirAnalysis | null; + objectJs: string | null; + bufferJs: string | null; + bufferBailed: boolean; +}; + +function emitBuffer( + fn: HirFunction, + context: HirSurfaceContext, +): string | null { + switch (context.surface) { + case "dynamics": + return emitBufferDynamicsJs(fn, context.elements); + case "lambda": { + const program = emitBufferLambdaJs(fn, context); + return program + ? `// inputSlotCount: ${program.inputSlotCount}\n${program.source}` + : null; + } + case "kernel": { + const program = emitBufferKernelJs(fn, context); + return program + ? `// inputSlotCount: ${program.inputSlotCount}, outputFloatCount: ${program.outputFloatCount}\n${program.source}` + : null; + } + } +} + +function runPipeline( + surface: HirSurfaceKind, + code: string, + schemaJson: string, +): PipelineResult { + let context: HirSurfaceContext | null = null; + let schemaError: string | null = null; + try { + const schema = JSON.parse(schemaJson) as PlaygroundSchema; + context = toContext(surface, schema); + } catch (error) { + schemaError = error instanceof Error ? error.message : String(error); + } + + const empty: PipelineResult = { + context, + schemaError, + fn: null, + diagnostics: [], + typecheck: null, + analysis: null, + objectJs: null, + bufferJs: null, + bufferBailed: false, + }; + if (!context) { + return empty; + } + + // One call runs the full pipeline: lowering, typecheck, analyses, lints. + const lint = lintHirUserCode(code, context); + if (!lint.fn) { + return { ...empty, diagnostics: lint.diagnostics }; + } + + let objectJs: string | null = null; + let bufferJs: string | null = null; + try { + objectJs = emitUserFunctionJs(lint.fn); + bufferJs = emitBuffer(lint.fn, context); + } catch (error) { + schemaError = error instanceof Error ? error.message : String(error); + } + + return { + context, + schemaError, + fn: lint.fn, + diagnostics: lint.diagnostics, + typecheck: lint.typecheck ?? null, + analysis: lint.analysis ?? null, + objectJs, + bufferJs, + bufferBailed: bufferJs === null, + }; +} + +// -- Styles --------------------------------------------------------------------- + +const pageStyle = css({ + display: "grid", + gridTemplateColumns: "[1fr 1fr]", + gap: "[16px]", + width: "full", + fontFamily: "body", +}); + +const columnStyle = css({ + display: "flex", + flexDirection: "column", + gap: "[12px]", + minWidth: "[0]", +}); + +const panelStyle = css({ + borderWidth: "[1px]", + borderStyle: "solid", + borderColor: "neutral.bd.subtle", + borderRadius: "lg", + overflow: "hidden", +}); + +const panelTitleStyle = css({ + fontSize: "xs", + fontWeight: "semibold", + color: "neutral.s80", + textTransform: "uppercase", + letterSpacing: "wide", + padding: "[6px 10px]", + backgroundColor: "neutral.s10", + borderBottomWidth: "[1px]", + borderBottomStyle: "solid", + borderBottomColor: "neutral.bd.subtle", +}); + +const preStyle = css({ + margin: "[0]", + padding: "[10px]", + fontSize: "xs", + fontFamily: "mono", + lineHeight: "[1.5]", + whiteSpace: "pre-wrap", + wordBreak: "break-word", + maxHeight: "[320px]", + overflow: "auto", +}); + +const surfaceBarStyle = css({ + display: "flex", + gap: "[6px]", +}); + +const surfaceButtonStyle = css({ + fontSize: "sm", + padding: "[4px 12px]", + borderRadius: "md", + borderWidth: "[1px]", + borderStyle: "solid", + borderColor: "neutral.bd.subtle", + cursor: "pointer", + backgroundColor: "[transparent]", + '&[data-active="true"]': { + backgroundColor: "neutral.s90", + color: "neutral.s10", + }, +}); + +const diagnosticStyle = css({ + display: "flex", + gap: "[8px]", + alignItems: "baseline", + padding: "[6px 10px]", + fontSize: "xs", + fontFamily: "mono", + borderBottomWidth: "[1px]", + borderBottomStyle: "solid", + borderBottomColor: "neutral.bd.subtle", +}); + +const severityStyle = css({ + fontWeight: "semibold", + textTransform: "uppercase", + fontSize: "[10px]", + '&[data-severity="error"]': { color: "[#dc2626]" }, + '&[data-severity="warning"]': { color: "[#d97706]" }, + '&[data-severity="info"]': { color: "[#2563eb]" }, + '&[data-severity="hint"]': { color: "neutral.s70" }, +}); + +const schemaTextareaStyle = css({ + width: "full", + minHeight: "[180px]", + fontFamily: "mono", + fontSize: "xs", + padding: "[10px]", + border: "none", + outline: "none", + resize: "vertical", +}); + +const okStyle = css({ + padding: "[10px]", + fontSize: "xs", + color: "[#16a34a]", +}); + +// -- Components ------------------------------------------------------------------- + +const Panel = ({ + title, + children, +}: { + title: string; + children: React.ReactNode; +}) => ( +
+
{title}
+ {children} +
+); + +const Json = ({ value }: { value: unknown }) => ( + 
{JSON.stringify(value, null, 2)}
+); + +const DiagnosticsPanel = ({ + diagnostics, + code, +}: { + diagnostics: HirDiagnostic[]; + code: string; +}) => ( + + {diagnostics.length === 0 ? ( +
No diagnostics — compiles cleanly.
+ ) : ( + diagnostics.map((diagnostic) => ( +
+ + {diagnostic.severity} + + + {diagnostic.code} {diagnostic.message} +
+ + @{diagnostic.span.start}.. + {diagnostic.span.start + diagnostic.span.length}:{" "} + {JSON.stringify( + code.slice( + diagnostic.span.start, + diagnostic.span.start + Math.min(diagnostic.span.length, 80), + ), + )} + + +
+ )) + )} + +); + +const HirPlayground = () => { + const [surface, setSurface] = useState("kernel"); + const [codeBySurface, setCodeBySurface] = useState(CODE_PRESETS); + const [schemaJson, setSchemaJson] = useState( + JSON.stringify(DEFAULT_SCHEMA, null, 2), + ); + + const code = codeBySurface[surface]; + const result = runPipeline(surface, code, schemaJson); + + return ( +
+ {/* Left column — inputs */} +
+
+ {(["dynamics", "lambda", "kernel"] as const).map((candidate) => ( + + ))} +
+ + + + setCodeBySurface((previous) => ({ + ...previous, + [surface]: next ?? "", + })) + } + /> + + + +