Adding a test for nested ternaries

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -189,9 +189,17 @@ function runWithEnvironment(
 
   assertConsistentIdentifiers(hir);
 
+  const wasTernaryConstantPropagationEnabled =
+    env.config.enableTernaryConstantPropagation;
+  env.config.enableTernaryConstantPropagation = false;
   constantPropagation(hir);
   log({kind: 'hir', name: 'ConstantPropagation', value: hir});
 
+  env.config.enableTernaryConstantPropagation =
+    wasTernaryConstantPropagationEnabled;
+  constantPropagation(hir);
+  log({kind: 'hir', name: 'ConstantPropagationTernary', value: hir});
+
   inferTypes(hir);
   log({kind: 'hir', name: 'InferTypes', value: hir});
 

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

@@ -622,6 +622,17 @@ const EnvironmentConfigSchema = z.object({
    * ```
    */
   lowerContextAccess: ExternalFunctionSchema.nullable().default(null),
+
+  /**
+   * If enabled, ConstantPropgation will try to resolve ternaries.
+   *
+   * // input
+   * const x = true ? b : c;
+   *
+   * // output
+   * const x = b;
+   */
+  enableTernaryConstantPropagation: z.boolean().default(true),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

compiler/packages/babel-plugin-react-compiler/src/Optimization/ConstantPropagation.ts

@@ -106,7 +106,6 @@ function applyConstantPropagation(
   fn: HIRFunction,
   constants: Constants,
 ): boolean {
-  let hasChanges = false;
   for (const [, block] of fn.body.blocks) {
     /*
      * Initialize phi values if all operands have the same known constant value.
@@ -134,7 +133,10 @@ function applyConstantPropagation(
         constants.set(instr.lvalue.identifier.id, value);
       }
     }
+  }
 
+  let hasChanges = false;
+  for (const [, block] of fn.body.blocks) {
     const terminal = block.terminal;
     switch (terminal.kind) {
       case 'if': {
@@ -154,6 +156,64 @@ function applyConstantPropagation(
         }
         break;
       }
+      case 'ternary': {
+        if (!fn.env.config.enableTernaryConstantPropagation) {
+          break;
+        }
+        const branchBlock = fn.body.blocks.get(terminal.test);
+        if (branchBlock === undefined) {
+          break;
+        }
+
+        if (branchBlock.terminal.kind !== 'branch') {
+          // TODO: could be other kinds like logical
+          break;
+        }
+
+        const testValue = read(constants, branchBlock.terminal.test);
+
+        if (testValue !== null && testValue.kind === 'Primitive') {
+          hasChanges = true;
+          const targetBlockId = testValue.value
+            ? branchBlock.terminal.consequent
+            : branchBlock.terminal.alternate;
+
+          // I think I can only set this if the block isn't
+          // used in a value position by its predecessor
+          block.kind = 'block';
+
+          block.terminal = {
+            kind: 'goto',
+            variant: GotoVariant.Break,
+            block: targetBlockId,
+            id: terminal.id,
+            loc: terminal.loc,
+          };
+
+          const fallthrough = fn.body.blocks.get(
+            branchBlock.terminal.fallthrough,
+          );
+
+          if (fallthrough?.terminal.kind == 'goto') {
+            fallthrough.kind = 'block';
+          }
+
+          const consequent = fn.body.blocks.get(
+            branchBlock.terminal.consequent,
+          )!;
+          const alternate = fn.body.blocks.get(branchBlock.terminal.alternate)!;
+
+          if (consequent.terminal.kind === 'goto') {
+            consequent.kind = 'block';
+          }
+
+          if (alternate.terminal.kind === 'goto') {
+            alternate.kind = 'block';
+          }
+        }
+
+        break;
+      }
       default: {
         // no-op
       }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts

@@ -216,7 +216,7 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
         if (node == null) {
           // Transition from block->value block, derive the outer block range
           CompilerError.invariant(fallthrough !== null, {
-            reason: `Expected a fallthrough for value block`,
+            reason: `Expected a fallthrough for value block ${terminal.id}`,
             loc: terminal.loc,
           });
           const fallthroughBlock = fn.body.blocks.get(fallthrough)!;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/constant-propagation-ternary-nested-unary.expect.md

@@ -0,0 +1,83 @@
+
+## Input
+
+```javascript
+// @enableTernaryConstantPropagation
+import {Stringify} from 'shared-runtime';
+
+function foo() {
+  let _b;
+  const b = true;
+  _b = !b ? 'bar' : b ? 'foo' : 'baz';
+
+  return (
+    <Stringify
+      value={{
+        _b,
+        b0: !true,
+        n0: !0,
+        n1: !1,
+        n2: !2,
+        n3: !-1,
+        s0: !'',
+        s1: !'a',
+        s2: !'ab',
+        u: !undefined,
+        n: !null,
+      }}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: foo,
+  params: [],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTernaryConstantPropagation
+import { Stringify } from "shared-runtime";
+
+function foo() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = (
+      <Stringify
+        value={{
+          _b: "foo",
+          b0: false,
+          n0: true,
+          n1: false,
+          n2: false,
+          n3: !-1,
+          s0: true,
+          s1: false,
+          s2: false,
+          u: !undefined,
+          n: true,
+        }}
+      />
+    );
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: foo,
+  params: [],
+  isComponent: false,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"value":{"_b":"foo","b0":false,"n0":true,"n1":false,"n2":false,"n3":false,"s0":true,"s1":false,"s2":false,"u":true,"n":true}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/constant-propagation-ternary-nested-unary.js

@@ -0,0 +1,32 @@
+// @enableTernaryConstantPropagation
+import {Stringify} from 'shared-runtime';
+
+function foo() {
+  let _b;
+  const b = true;
+  _b = !b ? 'bar' : b ? 'foo' : 'baz';
+
+  return (
+    <Stringify
+      value={{
+        _b,
+        b0: !true,
+        n0: !0,
+        n1: !1,
+        n2: !2,
+        n3: !-1,
+        s0: !'',
+        s1: !'a',
+        s2: !'ab',
+        u: !undefined,
+        n: !null,
+      }}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: foo,
+  params: [],
+  isComponent: false,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/constant-propagation-ternary-unary.expect.md

@@ -0,0 +1,83 @@
+
+## Input
+
+```javascript
+// @enableTernaryConstantPropagation
+import {Stringify} from 'shared-runtime';
+
+function foo() {
+  let _b;
+  const b = true;
+  _b = !b ? 'bar' : 'baz';
+
+  return (
+    <Stringify
+      value={{
+        _b,
+        b0: !true,
+        n0: !0,
+        n1: !1,
+        n2: !2,
+        n3: !-1,
+        s0: !'',
+        s1: !'a',
+        s2: !'ab',
+        u: !undefined,
+        n: !null,
+      }}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: foo,
+  params: [],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTernaryConstantPropagation
+import { Stringify } from "shared-runtime";
+
+function foo() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = (
+      <Stringify
+        value={{
+          _b: "baz",
+          b0: false,
+          n0: true,
+          n1: false,
+          n2: false,
+          n3: !-1,
+          s0: true,
+          s1: false,
+          s2: false,
+          u: !undefined,
+          n: true,
+        }}
+      />
+    );
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: foo,
+  params: [],
+  isComponent: false,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"value":{"_b":"baz","b0":false,"n0":true,"n1":false,"n2":false,"n3":false,"s0":true,"s1":false,"s2":false,"u":true,"n":true}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/constant-propagation-ternary-unary.js

@@ -0,0 +1,32 @@
+// @enableTernaryConstantPropagation
+import {Stringify} from 'shared-runtime';
+
+function foo() {
+  let _b;
+  const b = true;
+  _b = !b ? 'bar' : 'baz';
+
+  return (
+    <Stringify
+      value={{
+        _b,
+        b0: !true,
+        n0: !0,
+        n1: !1,
+        n2: !2,
+        n3: !-1,
+        s0: !'',
+        s1: !'a',
+        s2: !'ab',
+        u: !undefined,
+        n: !null,
+      }}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: foo,
+  params: [],
+  isComponent: false,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/memoize-value-block-value-conditional.expect.md

@@ -22,12 +22,14 @@ import { c as _c } from "react/compiler-runtime";
 function Foo(props) {
   const $ = _c(1);
   let x;
+  let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    true ? (x = []) : (x = {});
-    $[0] = x;
+    t0 = [];
+    $[0] = t0;
   } else {
-    x = $[0];
+    t0 = $[0];
   }
+  x = t0;
   return x;
 }