Validate that all variable references are consistently local/context

Validates that all references to a variable (pre-SSA) are consistently "local" 
references or "context" references. Ie, if a variable is declared as 
DeclareContext, any accesses must be eg LoadContext or StoreContext, not 
LoadLocal/StoreLocal. This will help with the issue from #2577 (assuming that we 
know a variable _is_ a context variable) but also provides a more precise 
bailout for an existing case with destructuring assignment to a context 
variable.

compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts

@@ -69,6 +69,7 @@ import {
 } from "../Utils/logger";
 import { assertExhaustive } from "../Utils/utils";
 import {
+  validateContextVariableLValues,
   validateFrozenLambdas,
   validateHooksUsage,
   validateMemoizedEffectDependencies,
@@ -117,6 +118,7 @@ function* runWithEnvironment(
   pruneMaybeThrows(hir);
   yield log({ kind: "hir", name: "PruneMaybeThrows", value: hir });
 
+  validateContextVariableLValues(hir);
   validateUseMemo(hir);
 
   dropManualMemoization(hir);

compiler/packages/babel-plugin-react-forget/src/Validation/ValidateContextVariableLValues.ts

@@ -0,0 +1,106 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError } from "..";
+import { HIRFunction, IdentifierId, Place } from "../HIR";
+import { printPlace } from "../HIR/PrintHIR";
+import {
+  eachInstructionValueLValue,
+  eachPatternOperand,
+} from "../HIR/visitors";
+
+/**
+ * Validates that all store/load references to a given named identifier align with the
+ * "kind" of that variable (normal variable or context variable). For example, a context
+ * variable may not be loaded/stored with regular StoreLocal/LoadLocal/Destructure instructions.
+ */
+export function validateContextVariableLValues(fn: HIRFunction): void {
+  const identifierKinds: IdentifierKinds = new Map();
+  validateContextVariableLValuesImpl(fn, identifierKinds);
+}
+
+function validateContextVariableLValuesImpl(
+  fn: HIRFunction,
+  identifierKinds: IdentifierKinds
+): void {
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      const { value } = instr;
+      switch (value.kind) {
+        case "DeclareContext":
+        case "StoreContext": {
+          visit(identifierKinds, value.lvalue.place, "context");
+          break;
+        }
+        case "LoadContext": {
+          visit(identifierKinds, value.place, "context");
+          break;
+        }
+        case "StoreLocal":
+        case "DeclareLocal": {
+          visit(identifierKinds, value.lvalue.place, "local");
+          break;
+        }
+        case "LoadLocal": {
+          visit(identifierKinds, value.place, "local");
+          break;
+        }
+        case "PostfixUpdate":
+        case "PrefixUpdate": {
+          visit(identifierKinds, value.lvalue, "local");
+          break;
+        }
+        case "Destructure": {
+          for (const lvalue of eachPatternOperand(value.lvalue.pattern)) {
+            visit(identifierKinds, lvalue, "local");
+          }
+          break;
+        }
+        case "ObjectMethod":
+        case "FunctionExpression": {
+          validateContextVariableLValuesImpl(
+            value.loweredFunc.func,
+            identifierKinds
+          );
+          break;
+        }
+        default: {
+          for (const _ of eachInstructionValueLValue(value)) {
+            CompilerError.throwTodo({
+              reason:
+                "ValidateContextVariableLValues: unhandled instruction variant",
+              loc: value.loc,
+              description: `Handle '${value.kind} lvalues`,
+              suggestions: null,
+            });
+          }
+        }
+      }
+    }
+  }
+}
+
+type IdentifierKinds = Map<IdentifierId, "local" | "context">;
+
+function visit(
+  identifiers: IdentifierKinds,
+  place: Place,
+  kind: "local" | "context"
+): void {
+  const prevKind = identifiers.get(place.identifier.id);
+  if (prevKind !== undefined && prevKind !== kind) {
+    CompilerError.invariant(false, {
+      reason: `Expected all references to a variable to be consistently local or context references`,
+      loc: place.loc,
+      description: `Identifier ${printPlace(
+        place
+      )} is referenced as a ${kind} variable, but was previously referenced as a ${prevKind} variable`,
+      suggestions: null,
+    });
+  }
+  identifiers.set(place.identifier.id, kind);
+}

compiler/packages/babel-plugin-react-forget/src/Validation/index.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+export { validateContextVariableLValues } from "./ValidateContextVariableLValues";
 export { validateFrozenLambdas } from "./ValidateFrozenLambdas";
 export { validateHooksUsage } from "./ValidateHooksUsage";
 export { validateMemoizedEffectDependencies } from "./ValidateMemoizedEffectDependencies";

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-repro-scope-missing-mutable-range.expect.md

@@ -0,0 +1,25 @@
+
+## 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)
+```
+          
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo.destructure-assignment-to-context-var.expect.md

@@ -18,7 +18,7 @@ function useFoo(props) {
 ## Error
 
 ```
-[ReactForget] Invariant: [InferReferenceEffects] Context variables are always mutable. (5:5)
+[ReactForget] Invariant: Expected all references to a variable to be consistently local or context references. Identifier <unknown> x$1 is referenced as a local variable, but was previously referenced as a context variable (3:3)
 ```
           
       

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-scope-missing-mutable-range.expect.md

@@ -1,51 +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>;
-}
-
-```
-
-## 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);
-    });
-    $[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;
-}
-
-```
-