From e33c9c43ccd36113f2d85d700f231679c7068612 Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Mon, 7 Aug 2023 17:41:35 -0400
Subject: [PATCH] [ssa] Patch: propagate every rewrite to function expressions

---

#1899 only propagated eliminated phi nodes to function expressions on the first
iteration through blocks. However, this was buggy as later iterations could
introduce rewrites that need to be propagated.

[playground
repro](https://0xeac7-forget.vercel.app/#eyJzb3VyY2UiOiJmdW5jdGlvbiBDb21wb25lbnQoKSB7XG4gIGNvbnN0IHggPSA0O1xuXG4gIGNvbnN0IGdldDQgPSAoKSA9PiB7XG4gICAgd2hpbGUgKGJhcigpKSB7XG4gICAgICBpZiAoYmF6KSB7XG4gICAgICAgIGJhcigpO1xuICAgICAgfVxuICAgIH1cbiAgICByZXR1cm4gKCkgPT4geDtcbiAgfTtcblxuICByZXR1cm4gZ2V0NDtcbn0ifQ==).

I manually synced #1907 to check that this fix works for the VR Store codebase.
---
 .../src/SSA/EliminateRedundantPhi.ts          |  9 ++--
 .../src/__tests__/e2e/constant-prop.e2e.js    |  5 +++
 ...e-phis-in-lambda-capture-context.expect.md | 42 ++++++++++---------
 .../rewrite-phis-in-lambda-capture-context.js | 16 ++++---
 4 files changed, 43 insertions(+), 29 deletions(-)

diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts
index 0cffa408c7..2d91fd5ad1 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts
@@ -46,7 +46,6 @@ export function eliminateRedundantPhi(
   // compare to see if any new rewrites were added in that iteration.
   let size = rewrites.size;
   do {
-    const isFirstIteration = !hasBackEdge;
     size = rewrites.size;
     for (const [blockId, block] of ir.blocks) {
       // On the first iteration of the loop check for any back-edges.
@@ -116,10 +115,10 @@ export function eliminateRedundantPhi(
             rewritePlace(place, rewrites);
           }
 
-          // visit function expressions on first iteration of each block
-          if (isFirstIteration) {
-            eliminateRedundantPhi(instr.value.loweredFunc, rewrites);
-          }
+          // recursive call to:
+          // - eliminate phi nodes in child node
+          // - propagate rewrites, which may have changed between iterations
+          eliminateRedundantPhi(instr.value.loweredFunc, rewrites);
         }
       }
 
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/e2e/constant-prop.e2e.js b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/e2e/constant-prop.e2e.js
index 3446030107..dfb68b2ce5 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/e2e/constant-prop.e2e.js
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/e2e/constant-prop.e2e.js
@@ -96,6 +96,11 @@ test("lambda-constant-propagation-of-phi-node", () => {
     if (constantValue) {
       noopCallback();
     }
+    for (let i = 0; i < 5; i++) {
+      if (!constantValue) {
+        noopCallback();
+      }
+    }
     const getDiv = () => <div>{x}</div>;
     return getDiv();
   }
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.expect.md
index fbf001317a..a94b4e7c94 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.expect.md
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.expect.md
@@ -2,13 +2,19 @@
 ## Input
 
 ```javascript
-function ConstantPropagationBug() {
-  const x = CONSTANT1;
-  const createPhiNode = CONSTANT2 || 5;
+function Component() {
+  const x = 4;
 
-  const getFoo = () => <Foo x={x} y={createPhiNode} />;
+  const get4 = () => {
+    while (bar()) {
+      if (baz) {
+        bar();
+      }
+    }
+    return () => x;
+  };
 
-  return getFoo();
+  return get4;
 }
 
 ```
@@ -17,26 +23,24 @@ function ConstantPropagationBug() {
 
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
-function ConstantPropagationBug() {
-  const $ = useMemoCache(2);
-
-  const createPhiNode = CONSTANT2 || 5;
+function Component() {
+  const $ = useMemoCache(1);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = () => <Foo x={CONSTANT1} y={createPhiNode} />;
+    t0 = () => {
+      while (bar()) {
+        if (baz) {
+          bar();
+        }
+      }
+      return () => 4;
+    };
     $[0] = t0;
   } else {
     t0 = $[0];
   }
-  const getFoo = t0;
-  let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = getFoo();
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  return t1;
+  const get4 = t0;
+  return get4;
 }
 
 ```
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.js b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.js
index 8ed131ece9..9d6fec2f61 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.js
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/rewrite-phis-in-lambda-capture-context.js
@@ -1,8 +1,14 @@
-function ConstantPropagationBug() {
-  const x = CONSTANT1;
-  const createPhiNode = CONSTANT2 || 5;
+function Component() {
+  const x = 4;
 
-  const getFoo = () => <Foo x={x} y={createPhiNode} />;
+  const get4 = () => {
+    while (bar()) {
+      if (baz) {
+        bar();
+      }
+    }
+    return () => x;
+  };
 
-  return getFoo();
+  return get4;
 }