From 200d2c379dd3acb024ec0b2ef9d649bb1d149bdd Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Fri, 22 Mar 2024 14:51:36 -0700
Subject: [PATCH] Fix for loops with value block index initial value

---
 .../src/SSA/LeaveSSA.ts                       | 33 +++++---
 ...oop-with-value-block-initializer.expect.md | 77 +++++++++++++++++++
 .../for-loop-with-value-block-initializer.js  | 22 ++++++
 3 files changed, 121 insertions(+), 11 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.js

diff --git a/compiler/packages/babel-plugin-react-forget/src/SSA/LeaveSSA.ts b/compiler/packages/babel-plugin-react-forget/src/SSA/LeaveSSA.ts
index 58616adece..7d75ec729b 100644
--- a/compiler/packages/babel-plugin-react-forget/src/SSA/LeaveSSA.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/SSA/LeaveSSA.ts
@@ -23,6 +23,7 @@ import {
   eachInstructionValueOperand,
   eachPatternOperand,
   eachTerminalOperand,
+  eachTerminalSuccessor,
   terminalFallthrough,
 } from "../HIR/visitors";
 
@@ -340,8 +341,14 @@ export function leaveSSA(fn: HIRFunction): void {
        * To avoid generating a let binding for the initializer prior to the loop,
        * check to see if the for declares an iterator variable.
        */
-      while (init.id !== initContinuation) {
-        for (const instr of init.instructions) {
+      const queue: Array<BlockId> = [init.id];
+      while (queue.length !== 0) {
+        const blockId = queue.shift()!;
+        if (blockId === initContinuation) {
+          break;
+        }
+        const block = fn.body.blocks.get(blockId)!;
+        for (const instr of block.instructions) {
           if (
             instr.value.kind === "StoreLocal" &&
             instr.value.lvalue.kind !== InstructionKind.Reassign
@@ -362,25 +369,29 @@ export function leaveSSA(fn: HIRFunction): void {
           }
         }
 
-        let next: BlockId | null = null;
-        switch (init.terminal.kind) {
+        switch (block.terminal.kind) {
           case "maybe-throw": {
-            next = init.terminal.continuation;
+            queue.push(block.terminal.continuation);
             break;
           }
           case "goto": {
-            next = init.terminal.block;
+            queue.push(block.terminal.block);
+            break;
+          }
+          case "branch":
+          case "logical":
+          case "optional":
+          case "ternary":
+          case "label": {
+            for (const successor of eachTerminalSuccessor(block.terminal)) {
+              queue.push(successor);
+            }
             break;
           }
           default: {
             break;
           }
         }
-        if (next === null) {
-          break;
-        }
-
-        init = fn.body.blocks.get(next)!;
       }
 
       if (terminal.kind === "for" && terminal.update !== null) {
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.expect.md
new file mode 100644
index 0000000000..8affe41443
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.expect.md
@@ -0,0 +1,77 @@
+
+## Input
+
+```javascript
+const TOTAL = 10;
+function Component(props) {
+  const items = [];
+  for (let i = props.start ?? 0; i < props.items.length; i++) {
+    const item = props.items[i];
+    items.push(<div key={item.id}>{item.value}</div>);
+  }
+  return <div>{items}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [
+    {
+      start: null,
+      items: [
+        { id: 0, value: "zero" },
+        { id: 1, value: "one" },
+      ],
+    },
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+const TOTAL = 10;
+function Component(props) {
+  const $ = useMemoCache(5);
+  let items;
+  if ($[0] !== props.start || $[1] !== props.items) {
+    items = [];
+    for (let i = props.start ?? 0; i < props.items.length; i++) {
+      const item = props.items[i];
+      items.push(<div key={item.id}>{item.value}</div>);
+    }
+    $[0] = props.start;
+    $[1] = props.items;
+    $[2] = items;
+  } else {
+    items = $[2];
+  }
+  let t0;
+  if ($[3] !== items) {
+    t0 = <div>{items}</div>;
+    $[3] = items;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [
+    {
+      start: null,
+      items: [
+        { id: 0, value: "zero" },
+        { id: 1, value: "one" },
+      ],
+    },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div><div>zero</div><div>one</div></div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.js
new file mode 100644
index 0000000000..b966dc60d1
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-loop-with-value-block-initializer.js
@@ -0,0 +1,22 @@
+const TOTAL = 10;
+function Component(props) {
+  const items = [];
+  for (let i = props.start ?? 0; i < props.items.length; i++) {
+    const item = props.items[i];
+    items.push(<div key={item.id}>{item.value}</div>);
+  }
+  return <div>{items}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [
+    {
+      start: null,
+      items: [
+        { id: 0, value: "zero" },
+        { id: 1, value: "one" },
+      ],
+    },
+  ],
+};