[valueblocks] Support conditional expressions (ternary)

Support conditional expressions from AST -> HIR -> ReactiveFunction -> AST. This 
also helps make the patterns for value block handling more clear, so i was able 
to extract some reusable logic in the HIR -> ReactiveFunction conversion phase.

compiler/forget/src/HIR/BuildHIR.ts

@@ -889,6 +889,69 @@ function lowerExpression(
         loc: exprLoc,
       };
     }
+    case "ConditionalExpression": {
+      const expr = exprPath as NodePath<t.ConditionalExpression>;
+      const exprLoc = expr.node.loc ?? GeneratedSource;
+
+      //  Block for code following the if
+      const continuationBlock = builder.reserve(builder.currentBlockKind());
+      const testBlock = builder.reserve("value");
+      const place = buildTemporaryPlace(builder, exprLoc);
+
+      //  Block for the consequent (if the test is truthy)
+      const consequentBlock = builder.enter("block", (blockId) => {
+        builder.push({
+          id: makeInstructionId(0),
+          lvalue: { kind: InstructionKind.Reassign, place: { ...place } },
+          value: lowerExpressionToPlace(builder, expr.get("consequent")),
+          loc: exprLoc,
+        });
+        return {
+          kind: "goto",
+          block: continuationBlock.id,
+          variant: GotoVariant.Break,
+          id: makeInstructionId(0),
+        };
+      });
+      //  Block for the alternate (if the test is not truthy)
+      const alternateBlock = builder.enter("block", (blockId) => {
+        builder.push({
+          id: makeInstructionId(0),
+          lvalue: { kind: InstructionKind.Reassign, place: { ...place } },
+          value: lowerExpressionToPlace(builder, expr.get("alternate")),
+          loc: exprLoc,
+        });
+        return {
+          kind: "goto",
+          block: continuationBlock.id,
+          variant: GotoVariant.Break,
+          id: makeInstructionId(0),
+        };
+      });
+
+      builder.terminateWithContinuation(
+        {
+          kind: "ternary",
+          fallthrough: continuationBlock.id,
+          id: makeInstructionId(0),
+          test: testBlock.id,
+          loc: exprLoc,
+        },
+        testBlock
+      );
+      const testPlace = lowerExpressionToPlace(builder, expr.get("test"));
+      builder.terminateWithContinuation(
+        {
+          kind: "branch",
+          test: { ...testPlace },
+          consequent: consequentBlock,
+          alternate: alternateBlock,
+          id: makeInstructionId(0),
+        },
+        continuationBlock
+      );
+      return place;
+    }
     case "LogicalExpression": {
       const expr = exprPath as NodePath<t.LogicalExpression>;
       const exprLoc = expr.node.loc ?? GeneratedSource;

compiler/forget/src/HIR/HIR.ts

@@ -96,7 +96,8 @@ export type ReactiveInstruction = {
 export type ReactiveValue =
   | InstructionValue
   | ReactiveLogicalValue
-  | ReactiveSequenceValue;
+  | ReactiveSequenceValue
+  | ReactiveTernaryValue;
 
 export type ReactiveLogicalValue = {
   kind: "LogicalExpression";
@@ -106,6 +107,14 @@ export type ReactiveLogicalValue = {
   loc: SourceLocation;
 };
 
+export type ReactiveTernaryValue = {
+  kind: "ConditionalExpression";
+  test: ReactiveValue;
+  consequent: ReactiveValue;
+  alternate: ReactiveValue;
+  loc: SourceLocation;
+};
+
 export type ReactiveSequenceValue = {
   kind: "SequenceExpression";
   instructions: Array<ReactiveInstruction>;

compiler/forget/src/HIR/PrintHIR.ts

@@ -354,6 +354,14 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       ].join("\n");
       break;
     }
+    case "ConditionalExpression": {
+      value = `Ternary ${printInstructionValue(
+        instrValue.test
+      )} ? ${printInstructionValue(
+        instrValue.consequent
+      )} : ${printInstructionValue(instrValue.alternate)}`;
+      break;
+    }
     default: {
       assertExhaustive(
         instrValue,

compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts

@@ -16,6 +16,7 @@ import {
   Place,
   ReactiveBlock,
   ReactiveValueBlock,
+  SourceLocation,
 } from "../HIR";
 import {
   HIRFunction,
@@ -28,6 +29,7 @@ import {
   ReactiveLogicalValue,
   ReactiveSequenceValue,
   ReactiveTerminalStatement,
+  ReactiveTernaryValue,
   ReactiveValue,
   Terminal,
 } from "../HIR/HIR";
@@ -465,6 +467,68 @@ class Driver {
     }
   }
 
+  visitValueBlockNew(
+    id: BlockId,
+    loc: SourceLocation
+  ): { block: BlockId; value: ReactiveValue; place: Place } {
+    let block: BasicBlock;
+    let value: ReactiveValue | null = null;
+    let place: Place | null = null;
+    const defaultBlock = this.cx.ir.blocks.get(id)!;
+    if (
+      defaultBlock.terminal.kind === "goto" ||
+      defaultBlock.terminal.kind === "branch"
+    ) {
+      block = defaultBlock;
+    } else {
+      const result = this.visitValueTerminal(defaultBlock.terminal);
+      block = this.cx.ir.blocks.get(result.fallthrough)!;
+      place = result.place;
+      value = result.value;
+    }
+    const instructions: Array<ReactiveInstruction> = block.instructions;
+    if (place !== null && value !== null) {
+      instructions.forEach((instr) =>
+        mapInstructionOperands(instr as Instruction, (place) => {
+          return place.identifier === place.identifier
+            ? (value as Place)
+            : place;
+        })
+      );
+    }
+    if (instructions.length === 0) {
+      invariant(
+        block.terminal.kind === "branch",
+        "Expected instructions for non-branch terminal"
+      );
+      return {
+        block: block.id,
+        place: block.terminal.test,
+        value: value ?? block.terminal.test,
+      };
+    } else if (instructions.length === 1) {
+      const instr = instructions[0]!;
+      return {
+        block: block.id,
+        place: instr.lvalue!.place,
+        value: instr.value,
+      };
+    } else {
+      const instr = instructions.at(-1)!;
+      const sequence: ReactiveSequenceValue = {
+        kind: "SequenceExpression",
+        instructions: instructions.slice(0, -1),
+        value: instr.value,
+        loc: loc,
+      };
+      return {
+        block: block.id,
+        place: instr.lvalue!.place,
+        value: sequence,
+      };
+    }
+  }
+
   visitValueTerminal(terminal: Terminal): {
     value: ReactiveValue;
     place: Place;
@@ -503,40 +567,7 @@ class Driver {
             })
           );
         }
-
-        let rightBlock: BasicBlock;
-        let rightValue: ReactiveValue | null = null;
-        let rightPlace: Place | null = null;
-        const defaultRightBlock = this.cx.ir.blocks.get(
-          testBlock.terminal.alternate
-        )!;
-        if (defaultRightBlock.terminal.kind === "goto") {
-          rightBlock = defaultRightBlock;
-        } else {
-          const rightResult = this.visitValueTerminal(
-            defaultRightBlock.terminal
-          );
-          rightBlock = this.cx.ir.blocks.get(rightResult.fallthrough)!;
-          rightPlace = rightResult.place;
-          rightValue = rightResult.value;
-        }
-        const rightInstructions: Array<ReactiveInstruction> =
-          rightBlock.instructions;
         const place = leftInstructions.at(-1)!.lvalue!.place;
-        invariant(
-          place.identifier ===
-            rightInstructions.at(-1)!.lvalue!.place.identifier,
-          "Expected both branches of a logical expression to store to the same temporary"
-        );
-        if (rightPlace !== null && rightValue !== null) {
-          rightInstructions.forEach((instr) =>
-            mapInstructionOperands(instr as Instruction, (place) => {
-              return place.identifier === rightPlace!.identifier
-                ? (rightValue as Place)
-                : place;
-            })
-          );
-        }
 
         let left: ReactiveValue;
         if (leftInstructions.length === 1) {
@@ -550,23 +581,19 @@ class Driver {
           };
           left = sequence;
         }
-        let right: ReactiveValue;
-        if (rightInstructions.length === 1) {
-          right = rightInstructions[0]!.value;
-        } else {
-          const sequence: ReactiveSequenceValue = {
-            kind: "SequenceExpression",
-            instructions: rightInstructions.slice(0, -1),
-            value: rightInstructions.at(-1)!.value,
-            loc: terminal.loc,
-          };
-          right = sequence;
-        }
+        const right = this.visitValueBlockNew(
+          testBlock.terminal.alternate,
+          terminal.loc
+        );
+        invariant(
+          place.identifier === right.place.identifier,
+          "Expected the left and right side of a logical expression to store a value to the same place"
+        );
         const value: ReactiveLogicalValue = {
           kind: "LogicalExpression",
           operator: terminal.operator,
           left,
-          right,
+          right: right.value,
           loc: terminal.loc,
         };
         return {
@@ -575,6 +602,39 @@ class Driver {
           fallthrough: terminal.fallthrough,
         };
       }
+      case "ternary": {
+        const test = this.visitValueBlockNew(terminal.test, terminal.loc);
+        const testBlock = this.cx.ir.blocks.get(test.block)!;
+        invariant(
+          testBlock.terminal.kind === "branch",
+          "Unexpected terminal kind '%s' for ternary test block",
+          testBlock.terminal.kind
+        );
+        const consequent = this.visitValueBlockNew(
+          testBlock.terminal.consequent,
+          terminal.loc
+        );
+        const alternate = this.visitValueBlockNew(
+          testBlock.terminal.alternate,
+          terminal.loc
+        );
+        const value: ReactiveTernaryValue = {
+          kind: "ConditionalExpression",
+          test: test.value,
+          consequent: consequent.value,
+          alternate: alternate.value,
+          loc: terminal.loc,
+        };
+        invariant(
+          consequent.place.identifier === alternate.place.identifier,
+          "Expected the consquent and alternate of a ternary to store a value to the same place"
+        );
+        return {
+          place: { ...consequent.place },
+          value,
+          fallthrough: terminal.fallthrough,
+        };
+      }
       default: {
         invariant(
           false,

compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -432,6 +432,7 @@ const createWhileStatement = withLoc(t.whileStatement);
 const createTaggedTemplateExpression = withLoc(t.taggedTemplateExpression);
 const createLogicalExpression = withLoc(t.logicalExpression);
 const createSequenceExpression = withLoc(t.sequenceExpression);
+const createConditionalExpression = withLoc(t.conditionalExpression);
 
 type Temporaries = Map<IdentifierId, t.Expression>;
 
@@ -693,6 +694,15 @@ function codegenInstructionValue(
       );
       break;
     }
+    case "ConditionalExpression": {
+      value = createConditionalExpression(
+        instrValue.loc,
+        codegenInstructionValue(cx, instrValue.test),
+        codegenInstructionValue(cx, instrValue.consequent),
+        codegenInstructionValue(cx, instrValue.alternate)
+      );
+      break;
+    }
     case "SequenceExpression": {
       const body = codegenBlock(
         cx,
@@ -705,6 +715,15 @@ function codegenInstructionValue(
         if (stmt.type === "ExpressionStatement") {
           return stmt.expression;
         } else {
+          if (t.isVariableDeclaration(stmt)) {
+            const declarator = stmt.declarations[0];
+            todoInvariant(
+              false,
+              `Cannot declare variables in a value block, tried to declare '${
+                (declarator.id as t.Identifier).name
+              }'`
+            );
+          }
           todoInvariant(
             false,
             `Handle conversion of ${stmt.type} to expression`

compiler/forget/src/ReactiveScopes/PrintReactiveFunction.ts

@@ -124,6 +124,20 @@ function printValueBlock(writer: Writer, block: ReactiveValueBlock): void {
 
 function printReactiveValue(writer: Writer, value: ReactiveValue): void {
   switch (value.kind) {
+    case "ConditionalExpression": {
+      writer.append(`Ternary `);
+      printReactiveValue(writer, value.test);
+      writer.newline();
+      writer.indented(() => {
+        writer.write(`? `);
+        printReactiveValue(writer, value.consequent);
+        writer.newline();
+        writer.write(`: `);
+        printReactiveValue(writer, value.alternate);
+        writer.newline();
+      });
+      break;
+    }
     case "LogicalExpression": {
       writer.append(`Logical ${value.operator} `);
       printReactiveValue(writer, value.left);

compiler/forget/src/ReactiveScopes/visitors.ts

@@ -98,6 +98,12 @@ export function* eachReactiveValueOperand(
       yield* eachReactiveValueOperand(instrValue.value);
       break;
     }
+    case "ConditionalExpression": {
+      yield* eachReactiveValueOperand(instrValue.test);
+      yield* eachReactiveValueOperand(instrValue.consequent);
+      yield* eachReactiveValueOperand(instrValue.alternate);
+      break;
+    }
     default: {
       yield* eachInstructionValueOperand(instrValue);
     }

compiler/forget/src/SSA/LeaveSSA.ts

@@ -148,7 +148,7 @@ export function leaveSSA(fn: HIRFunction) {
       pushPhis(rewritePhis, update);
       update.phis.clear();
     }
-    if (terminal.kind === "logical") {
+    if (terminal.kind === "logical" || terminal.kind === "ternary") {
       const fallthrough = fn.body.blocks.get(terminal.fallthrough)!;
       pushPhis(rewritePhis, fallthrough);
       fallthrough.phis.clear();

compiler/forget/src/__tests__/fixtures/hir/logical-expression.expect.md

@@ -5,7 +5,7 @@
 function component(props) {
   let a = props.a || (props.b && props.c && props.d);
   let b = (props.a && props.b && props.c) || props.d;
-  return a ?? b;
+  return a ? b : props.c;
 }
 
 ```
@@ -37,16 +37,18 @@ function component(props) {
   const b = t3;
   const c_4 = $[4] !== a;
   const c_5 = $[5] !== b;
-  let t6;
-  if (c_4 || c_5) {
-    t6 = a ?? b;
+  const c_6 = $[6] !== props;
+  let t7;
+  if (c_4 || c_5 || c_6) {
+    t7 = a ? b : props.c;
     $[4] = a;
     $[5] = b;
-    $[6] = t6;
+    $[6] = props;
+    $[7] = t7;
   } else {
-    t6 = $[6];
+    t7 = $[7];
   }
-  return t6;
+  return t7;
 }
 
 ```

compiler/forget/src/__tests__/fixtures/hir/logical-expression.js

@@ -1,5 +1,5 @@
 function component(props) {
   let a = props.a || (props.b && props.c && props.d);
   let b = (props.a && props.b && props.c) || props.d;
-  return a ?? b;
+  return a ? b : props.c;
 }

compiler/forget/src/__tests__/fixtures/hir/sequence-expression.expect.md

@@ -0,0 +1,51 @@
+
+## Input
+
+```javascript
+// @only
+function sequence(props) {
+  let x = (null, Math.max(1, 2), sequence({}));
+  if (((x = x + 1), x < 10)) {
+    x = 10;
+  }
+  //   while (((x = x * 2), x < 20)) {
+  //     x = ((x = x + 1), x + 1);
+  //   }
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+// @only
+function sequence(props) {
+  const $ = React.useMemoCache();
+  null;
+  Math.max(1, 2);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = sequence({});
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  const x$0 = x + 1;
+  let x$1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    x$1 = x$0;
+    if (x$0 < 10) {
+      const x$2 = 10;
+      x$1 = x$2;
+    }
+    $[1] = x$1;
+  } else {
+    x$1 = $[1];
+  }
+  return x$1;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/sequence-expression.js

@@ -0,0 +1,12 @@
+// @skip
+function sequence(props) {
+  let x = (null, Math.max(1, 2), sequence({}));
+  if (((x = x + 1), x < 10)) {
+    x = 10;
+  }
+  x = ((x = x + 1), x > 15) && x < 20 ? ((x = x + 1), x) : 42;
+  while (((x = x * 2), x < 20)) {
+    x = ((x = x + 1), x + 1);
+  }
+  return x;
+}