[WIP] Trying to make ternaries do constant propagation

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

@@ -190,9 +190,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

@@ -624,6 +624,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/HIR/MergeConsecutiveBlocks.ts

@@ -47,8 +47,8 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
     if (
       // Can only merge blocks with a single predecessor
       block.preds.size !== 1 ||
-      // Value blocks cannot merge
-      block.kind !== 'block' ||
+      // Loop blocks cannot merge
+      (block.kind !== 'block' && block.kind !== 'value') ||
       // Merging across fallthroughs could move the predecessor out of its block scope
       fallthroughBlocks.has(block.id)
     ) {
@@ -63,7 +63,13 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
       loc: null,
       suggestions: null,
     });
-    if (predecessor.terminal.kind !== 'goto' || predecessor.kind !== 'block') {
+
+    const predecessorValueWithNoInstructions =
+      predecessor.kind === 'value' && predecessor.instructions.length === 0;
+    if (
+      predecessor.terminal.kind !== 'goto' ||
+      (predecessor.kind !== 'block' && !predecessorValueWithNoInstructions)
+    ) {
       /*
        * The predecessor is not guaranteed to transfer control to this block,
        * they aren't consecutive.

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.
@@ -120,6 +119,8 @@ function applyConstantPropagation(
       }
     }
 
+    const localReassignments = new Map<IdentifierId, Place>();
+
     for (let i = 0; i < block.instructions.length; i++) {
       if (block.kind === 'sequence' && i === block.instructions.length - 1) {
         /*
@@ -133,8 +134,28 @@ function applyConstantPropagation(
       if (value !== null) {
         constants.set(instr.lvalue.identifier.id, value);
       }
+
+      switch (instr.value.kind) {
+        case 'StoreLocal': {
+          const identifierValue = localReassignments.get(
+            instr.value.value.identifier.id,
+          );
+          if (identifierValue != null) {
+            // constants.set(value.lvalue.place.identifier.id, placeValue);
+            instr.value.value = identifierValue;
+          }
+
+          localReassignments.set(
+            instr.value.lvalue.place.identifier.id,
+            instr.value.value,
+          );
+        }
+      }
     }
+  }
 
+  let hasChanges = false;
+  for (const [, block] of fn.body.blocks) {
     const terminal = block.terminal;
     switch (terminal.kind) {
       case 'if': {
@@ -154,6 +175,41 @@ 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;
+
+          const chosenBlock = fn.body.blocks.get(targetBlockId);
+          if (chosenBlock?.terminal.kind === 'goto') {
+            block.terminal = {
+              kind: 'goto',
+              variant: GotoVariant.Break,
+              block: targetBlockId,
+              id: terminal.id,
+              loc: terminal.loc,
+            };
+          }
+        }
+
+        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,
+};