From 35fee3135561ec3259cce695f4e2bb0863f7d7e0 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 6 Nov 2023 08:33:30 -0800
Subject: [PATCH] Handle IIFE with logical mutated later

The previous changes mostly meant that we removed the label terminal and didn't
have instructions for the same scope split in a way that we couldn't merge. But
logicals were still causing a split because MergeConsecutiveScopes can't merge
the blocks in that case. Here we move PruneUnusedLabels earlier in the pipeline
to ensure that instructions from IIFEs have floated up to the parent block scope
level.
---
 .../src/Entrypoint/Pipeline.ts                | 14 +++---
 ...fe-return-modified-later-logical.expect.md | 49 +++++++++++++++++++
 .../iife-return-modified-later-logical.js     | 12 +++++
 3 files changed, 68 insertions(+), 7 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.js

diff --git a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
index 7982942333..f57e96e4e6 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
@@ -210,6 +210,13 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
+  pruneUnusedLabels(reactiveFunction);
+  yield log({
+    kind: "reactive",
+    name: "PruneUnusedLabels",
+    value: reactiveFunction,
+  });
+
   memoizeFbtOperandsInSameScope(reactiveFunction);
   yield log({
     kind: "reactive",
@@ -320,13 +327,6 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
-  pruneUnusedLabels(reactiveFunction);
-  yield log({
-    kind: "reactive",
-    name: "PruneUnusedLabels",
-    value: reactiveFunction,
-  });
-
   pruneUnusedLValues(reactiveFunction);
   yield log({
     kind: "reactive",
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.expect.md
new file mode 100644
index 0000000000..dfa5077e00
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.expect.md
@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const items = (() => {
+    return foo() ?? [];
+  })();
+  items.push(props.a);
+  return items;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: {} }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t10;
+  let items;
+  if ($[0] !== props.a) {
+    t10 = foo() ?? [];
+    items = t10;
+
+    items.push(props.a);
+    $[0] = props.a;
+    $[1] = items;
+    $[2] = t10;
+  } else {
+    items = $[1];
+    t10 = $[2];
+  }
+  return items;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: {} }],
+};
+
+```
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.js
new file mode 100644
index 0000000000..adf84e69e4
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/iife-return-modified-later-logical.js
@@ -0,0 +1,12 @@
+function Component(props) {
+  const items = (() => {
+    return foo() ?? [];
+  })();
+  items.push(props.a);
+  return items;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: {} }],
+};