[WIP] Trying to make ternaries do constant propagation

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(false),
 });
 
 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,46 @@ 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
+            ? /*
+               * Do I need to change these from being value blocks?
+               * Currently getting
+               * Invariant: Expected a fallthrough for value block (6:6)
+               */
+              branchBlock.terminal.consequent
+            : branchBlock.terminal.alternate;
+          block.terminal = {
+            kind: 'goto',
+            variant: GotoVariant.Break,
+            block: targetBlockId,
+            id: terminal.id,
+            loc: terminal.loc,
+          };
+
+          fn.body.blocks.get(branchBlock.terminal.consequent)!.kind = 'block';
+          fn.body.blocks.get(branchBlock.terminal.alternate)!.kind = 'block';
+        }
+
+        break;
+      }
       default: {
         // no-op
       }

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,
+};