From d4cdaf252337027615812a01c366c7ae0265ea95 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Fri, 9 Feb 2024 14:25:33 -0800
Subject: [PATCH] Fix assignment expression with context variables

Fixes the one case discovered in the previous PR; for AssignmentExpression we
correctly lowered the store instruction to a local/context, but then always used
a `LoadLocal` to read the result back.

The load instruction appears like it might be dangling - i think what was
happening is that DCE cleaned up the unused LoadLocal whereas it leaves the
LoadContext alone. But this works for now, we can always clean up the extra
instruction later since this case isn't too common.
---
 .../src/HIR/BuildHIR.ts                       |  3 +-
 ...epro-scope-missing-mutable-range.expect.md | 25 ---------
 ...epro-scope-missing-mutable-range.expect.md | 52 +++++++++++++++++++
 ...s => repro-scope-missing-mutable-range.js} |  0
 4 files changed, 54 insertions(+), 26 deletions(-)
 delete mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.expect.md
 rename compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/{error.todo-repro-scope-missing-mutable-range.js => repro-scope-missing-mutable-range.js} (100%)

diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts
index afffabe288..1900df6ced 100644
--- a/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts
@@ -1869,6 +1869,7 @@ function lowerExpression(
               type: null,
               loc: exprLoc,
             });
+            return { kind: "LoadLocal", place: identifier, loc: exprLoc };
           } else {
             lowerValueToTemporary(builder, {
               kind: "StoreContext",
@@ -1879,8 +1880,8 @@ function lowerExpression(
               value: { ...binaryPlace },
               loc: exprLoc,
             });
+            return { kind: "LoadContext", place: identifier, loc: exprLoc };
           }
-          return { kind: "LoadLocal", place: identifier, loc: exprLoc };
         }
         case "MemberExpression": {
           // a.b.c += <right>
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.expect.md
deleted file mode 100644
index 29f3639655..0000000000
--- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.expect.md
+++ /dev/null
@@ -1,25 +0,0 @@
-
-## Input
-
-```javascript
-function HomeDiscoStoreItemTileRating(props) {
-  const item = useFragment();
-  let count = 0;
-  const aggregates = item?.aggregates || [];
-  aggregates.forEach((aggregate) => {
-    count += aggregate.count || 0;
-  });
-
-  return <Text>{count}</Text>;
-}
-
-```
-
-
-## Error
-
-```
-[ReactForget] Invariant: Expected all references to a variable to be consistently local or context references. Identifier <unknown> count$6 is referenced as a local variable, but was previously referenced as a context variable (6:6)
-```
-          
-      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.expect.md
new file mode 100644
index 0000000000..2d78251da6
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.expect.md
@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+function HomeDiscoStoreItemTileRating(props) {
+  const item = useFragment();
+  let count = 0;
+  const aggregates = item?.aggregates || [];
+  aggregates.forEach((aggregate) => {
+    count += aggregate.count || 0;
+  });
+
+  return <Text>{count}</Text>;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function HomeDiscoStoreItemTileRating(props) {
+  const $ = useMemoCache(4);
+  const item = useFragment();
+  let count;
+  if ($[0] !== item) {
+    count = 0;
+    const aggregates = item?.aggregates || [];
+    aggregates.forEach((aggregate) => {
+      count = count + (aggregate.count || 0);
+      count;
+    });
+    $[0] = item;
+    $[1] = count;
+  } else {
+    count = $[1];
+  }
+
+  const t0 = count;
+  let t1;
+  if ($[2] !== t0) {
+    t1 = <Text>{t0}</Text>;
+    $[2] = t0;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.js
similarity index 100%
rename from compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.js
rename to compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.js