FE-629: Fix HIR kernel compilation for conditional distributions

.changeset/quiet-ternaries-compile.md

@@ -0,0 +1,5 @@
+---
+"@hashintel/petrinaut-core": patch
+---
+
+Fix compilation of real transition-kernel attributes that conditionally return a number or a distribution.

libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.test.ts

@@ -282,6 +282,33 @@ describe("emitBufferKernelJs (token format v2)", () => {
     expect(sinkCalls[2]!.payload).toBe("order-1");
   });
 
+  it("writes conditional scalar and distribution values through their selected paths", () => {
+    const code = `export default TransitionKernel((input) => {
+  const sampledWhenActive = input.Pool[0].alive ? Distribution.Uniform(10, 20) : 0;
+  const sampledWhenInactive = input.Pool[1].alive ? Distribution.Uniform(30, 40) : 0;
+  return {
+    Out: [
+      { a: sampledWhenActive, b: 1, label: "first", id: Uuid.from("first"), flag: true },
+      { a: sampledWhenInactive, b: 2, label: "second", id: Uuid.from("second"), flag: false },
+    ],
+  };
+});`;
+
+    const { stagingViews, sinkCalls } = runKernel(code, new StringPool());
+    const stride = outLayout.strideBytes;
+    expect(stagingViews.f64[fieldOffset("a") / 8]).toBe(0);
+    expect(stagingViews.f64[(stride + fieldOffset("a")) / 8]).toBe(0);
+    const distributions = sinkCalls.filter(({ kind }) => kind === "dist");
+    expect(distributions).toHaveLength(1);
+    const [distribution] = distributions;
+    expect(distribution?.index).toBe(fieldOffset("a") / 8);
+    expect(distribution?.payload as RuntimeDistribution).toMatchObject({
+      type: "uniform",
+      min: 10,
+      max: 20,
+    });
+  });
+
   it("forwards whole input tokens, deferring uuid copies through the sink", () => {
     const forwardContext: HirKernelContext = {
       ...kernelContext,

libs/@hashintel/petrinaut-core/src/hir/emit-buffer-js.ts

@@ -78,6 +78,7 @@ const RESERVED_NAMES = [
   "Math",
   "Number",
   "RangeError",
+  "TypeError",
   "Infinity",
   "NaN",
 ];
@@ -91,11 +92,19 @@ type MetricPlaceRef = {
   elements: HirTokenElementInfo[];
 };
 
+/** A JS expression producing a number or boolean. */
+type ScalarValue = { kind: "scalar"; code: string };
+
+/** A JS expression producing a RuntimeDistribution (a hoisted const). */
+type DistributionValue = { kind: "dist"; code: string };
+
+/** A materialized scalar-or-distribution kernel output. */
+type MixedOutputValue = { kind: "mixed"; code: string };
+
 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 }
+  | ScalarValue
+  | DistributionValue
+  | MixedOutputValue
   /** `Uuid.generate()` / `Uuid.from(...)` — resolved by the engine sink. */
   | { kind: "uuidSentinel"; mode: "generate" | "from"; code: string }
   /** The lambda/kernel first parameter (`tokensByPlace`). */
@@ -117,6 +126,14 @@ type Value =
   /** `a.concat(b)` over place token arrays (parts in source order). */
   | { kind: "placeTokensConcat"; parts: MetricPlaceRef[] };
 
+function isMixedOutputValue(
+  value: Value,
+): value is ScalarValue | DistributionValue | MixedOutputValue {
+  return (
+    value.kind === "scalar" || value.kind === "dist" || value.kind === "mixed"
+  );
+}
+
 function quoteKey(key: string): string {
   return JSON.stringify(key);
 }
@@ -560,6 +577,20 @@ class BufferEmitter {
           );
           return { kind: thenValue.kind, code: result };
         }
+        if (isMixedOutputValue(thenValue) && isMixedOutputValue(elseValue)) {
+          const result = this.names.allocate("__conditionalValue");
+          this.lines.push(
+            `let ${result};`,
+            `if (${condition}) {`,
+            ...thenLines.map((line) => `  ${line}`),
+            `  ${result} = ${thenValue.code};`,
+            `} else {`,
+            ...elseLines.map((line) => `  ${line}`),
+            `  ${result} = ${elseValue.code};`,
+            `}`,
+          );
+          return { kind: "mixed", code: result };
+        }
         // Structural conditionals (records/tuples) stay on the object path.
         throw new BailError();
       }
@@ -909,6 +940,21 @@ export function emitBufferKernelJs(
             return null;
           }
 
+          if (value.kind === "mixed") {
+            if (element.type !== "real") {
+              return null;
+            }
+            writes.push(
+              `if (typeof ${value.code} === "number") {`,
+              `  outF64[${at >> 3}] = ${value.code};`,
+              `} else if (${value.code} !== null && typeof ${value.code} === "object" && ${value.code}.__brand === "distribution") {`,
+              `  __sink("dist", ${at >> 3}, ${value.code});`,
+              `} else {`,
+              `  throw new TypeError("Conditional real kernel outputs must produce a number or Distribution");`,
+              `}`,
+            );
+            continue;
+          }
           if (value.kind === "dist") {
             if (element.type !== "real") {
               return null;