From b686b5dd9d72705e84bd7a18fd6f4ece3a566921 Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Wed, 16 Aug 2023 15:39:00 +0100
Subject: [PATCH] [hir] Refactor constant propagation of phis

This PR moves the phi evaluation to a separate function.

Most importantly, it inverts the default case to _not_ constant propagate unless
we have explicit validation of the phi operands.
---
 .../src/Optimization/ConstantPropagation.ts   | 65 +++++++++++++------
 .../constant-propagate-global-phis.expect.md  | 14 ++--
 2 files changed, 54 insertions(+), 25 deletions(-)

diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts
index 093e773675..8b7691ee1d 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts
@@ -6,6 +6,7 @@
  */
 
 import { isValidIdentifier } from "@babel/types";
+import { CompilerError } from "../CompilerError";
 import {
   Environment,
   GotoVariant,
@@ -14,6 +15,7 @@ import {
   Instruction,
   InstructionValue,
   LoadGlobal,
+  Phi,
   Place,
   Primitive,
   assertConsistentIdentifiers,
@@ -116,25 +118,7 @@ function applyConstantPropagation(
     // Note that this analysis uses a single-pass only, so it will never fill in
     // phi values for blocks that have a back-edge.
     for (const phi of block.phis) {
-      let value: Primitive | LoadGlobal | null = null;
-      for (const [, operand] of phi.operands) {
-        const operandValue = constants.get(operand.id) ?? null;
-        if (operandValue === null) {
-          value = null;
-          break;
-        }
-        if (value === null) {
-          value = operandValue;
-        } else if (
-          operandValue.kind !== value.kind ||
-          (operandValue.kind === "Primitive" &&
-            value.kind === "Primitive" &&
-            operandValue.value !== value.value)
-        ) {
-          value = null;
-          break;
-        }
-      }
+      let value = evaluatePhi(phi, constants);
       if (value !== null) {
         constants.set(phi.id.id, value);
       }
@@ -187,6 +171,49 @@ function applyConstantPropagation(
   return hasChanges;
 }
 
+function evaluatePhi(phi: Phi, constants: Constants): Constant | null {
+  let value: Constant | null = null;
+  for (const [, operand] of phi.operands) {
+    const operandValue = constants.get(operand.id) ?? null;
+    // did not find a constant, can't constant propogate
+    if (operandValue === null) {
+      return null;
+    }
+
+    // first iteration of the loop, let's store the operand and continue
+    // looping.
+    if (value === null) {
+      value = operandValue;
+      continue;
+    }
+
+    // found different kinds of constants, can't constant propogate
+    if (operandValue.kind !== value.kind) {
+      return null;
+    }
+
+    switch (operandValue.kind) {
+      case "Primitive": {
+        CompilerError.invariant(value.kind === "Primitive", {
+          reason: "value kind expected to be Primitive",
+          loc: null,
+          suggestions: null,
+        });
+
+        // different constant values, can't constant propogate
+        if (operandValue.value !== value.value) {
+          return null;
+        }
+        break;
+      }
+      default:
+        return null;
+    }
+  }
+
+  return value;
+}
+
 function evaluateInstruction(
   env: Environment,
   constants: Constants,
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagate-global-phis.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagate-global-phis.expect.md
index 02520f234d..01554113e7 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagate-global-phis.expect.md
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagate-global-phis.expect.md
@@ -16,15 +16,17 @@ function Test() {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Test() {
-  const $ = useMemoCache(1);
+  const $ = useMemoCache(2);
   const { tab } = useFoo();
-  tab === WAT ? WAT : BAR;
+  const currentTab = tab === WAT ? WAT : BAR;
+  const c_0 = $[0] !== currentTab;
   let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = <Foo value={WAT} />;
-    $[0] = t0;
+  if (c_0) {
+    t0 = <Foo value={currentTab} />;
+    $[0] = currentTab;
+    $[1] = t0;
   } else {
-    t0 = $[0];
+    t0 = $[1];
   }
   return t0;
 }