Change rules for context variable promotion

This PR updates the conditions for which variables are promoted to context 
variables. 

The previous rule was to promote any variable where the variable was reassigned 
in some function expression other than the function that declared the variable. 
Notably, this meant that we did not use context variables for variables which 
were captured in a function expression, but reassigned _outside_ a function 
expression. 

The new rule is more consistent: we promote any variable which is a) reassigned 
somewhere and b) referenced in some function expression outside of their 
declaring function. The implementation builds two sets of identifiers, one for 
each criteria, then takes the union of these two sets. 

## Motivation 

The motivation for this change is to unblock additional validations and 
optimizations of function expressions. It's currently difficult to translate 
metadata that we infer about identifiers outside of a function expression into 
metadata about the identifiers within a function expression — for example to 
infer types within function expression bodies based on type information outside, 
propagate constants into functions, infer reference effects, etc. 

After this change, the only free variables inside function expressions will be 
variables that are effectively `const` - never reassigned anywhere. Thus it will 
be safe to renumber those identifiers to match the outer context (during 
EnterSSA), making it trivial to map metadata from outside the function into the 
function. 

This change also more closely models the runtime representation — any variable 
referenced in a function, and reassigned somewhere, would have to be compiled 
(ie in a JS engine) to use a context variable.

compiler/forget/packages/babel-plugin-react-forget/src/HIR/FindContextIdentifiers.ts

@@ -1,25 +1,26 @@
 import type { NodePath } from "@babel/traverse";
 import type * as t from "@babel/types";
 import { CompilerError } from "../CompilerError";
+import { Set_union } from "../Utils/utils";
 import { GeneratedSource } from "./HIR";
 
 type FindContextIdentifierState = {
-  inLambda: number;
   currentLambda: Array<
     | NodePath<t.FunctionDeclaration>
     | NodePath<t.FunctionExpression>
     | NodePath<t.ArrowFunctionExpression>
   >;
-  contextIdentifiers: Set<t.Identifier>;
+  reassigned: Set<t.Identifier>;
+  referenced: Set<t.Identifier>;
 };
 
 export function findContextIdentifiers(
   func: NodePath<t.Function>
 ): Set<t.Identifier> {
   const state: FindContextIdentifierState = {
-    inLambda: 0,
     currentLambda: [],
-    contextIdentifiers: new Set(),
+    reassigned: new Set(),
+    referenced: new Set(),
   };
 
   func.traverse<FindContextIdentifierState>(
@@ -70,25 +71,43 @@ export function findContextIdentifiers(
       AssignmentExpression(
         path: NodePath<t.AssignmentExpression>,
         state: FindContextIdentifierState
+      ): void {
+        const left = path.get("left");
+        handleAssignment(state.reassigned, left);
+      },
+      Identifier(
+        path: NodePath<t.Identifier>,
+        state: FindContextIdentifierState
       ): void {
         const currentLambda = state.currentLambda.at(-1);
-        if (currentLambda) {
-          const left = path.get("left");
-          handleAssignment(currentLambda, state.contextIdentifiers, left);
-        }
+        if (currentLambda !== undefined)
+          handleIdentifier(currentLambda, state.referenced, path);
       },
     },
     state
   );
-  return state.contextIdentifiers;
+  return Set_union(state.reassigned, state.referenced);
 }
 
-function handleAssignment(
+function handleIdentifier(
   currentLambda:
     | NodePath<t.FunctionDeclaration>
     | NodePath<t.FunctionExpression>
     | NodePath<t.ArrowFunctionExpression>,
-  contextIdentifiers: Set<t.Identifier>,
+  referenced: Set<t.Identifier>,
+  path: NodePath<t.Identifier>
+): void {
+  const name = path.node.name;
+  const binding = path.scope.getBinding(name);
+  const bindingAboveLambdaScope = currentLambda.scope.parent.getBinding(name);
+
+  if (binding != null && binding === bindingAboveLambdaScope) {
+    referenced.add(binding.identifier);
+  }
+}
+
+function handleAssignment(
+  reassigned: Set<t.Identifier>,
   lvalPath: NodePath<t.LVal>
 ): void {
   // Find all reassignments to identifiers declared outside of currentLambda
@@ -98,12 +117,9 @@ function handleAssignment(
     case "Identifier": {
       const path = lvalPath as NodePath<t.Identifier>;
       const name = path.node.name;
-      const ownBinding = path.scope.getBinding(name);
-      const bindingAboveLambdaScope =
-        currentLambda.scope.parent.getBinding(name);
-
-      if (ownBinding != null && ownBinding === bindingAboveLambdaScope) {
-        contextIdentifiers.add(ownBinding.identifier);
+      const binding = path.scope.getBinding(name);
+      if (binding != null) {
+        reassigned.add(binding.identifier);
       }
       break;
     }
@@ -111,7 +127,7 @@ function handleAssignment(
       const path = lvalPath as NodePath<t.ArrayPattern>;
       for (const element of path.get("elements")) {
         if (nonNull(element)) {
-          handleAssignment(currentLambda, contextIdentifiers, element);
+          handleAssignment(reassigned, element);
         }
       }
       break;
@@ -127,7 +143,7 @@ function handleAssignment(
               valuePath.node.loc ?? GeneratedSource
             );
           }
-          handleAssignment(currentLambda, contextIdentifiers, valuePath);
+          handleAssignment(reassigned, valuePath);
         } else {
           if (!property.isRestElement()) {
             CompilerError.invariant(
@@ -135,7 +151,7 @@ function handleAssignment(
               property.node.loc ?? GeneratedSource
             );
           }
-          handleAssignment(currentLambda, contextIdentifiers, property);
+          handleAssignment(reassigned, property);
         }
       }
       break;
@@ -143,12 +159,12 @@ function handleAssignment(
     case "AssignmentPattern": {
       const path = lvalPath as NodePath<t.AssignmentPattern>;
       const left = path.get("left");
-      handleAssignment(currentLambda, contextIdentifiers, left);
+      handleAssignment(reassigned, left);
       break;
     }
     case "RestElement": {
       const path = lvalPath as NodePath<t.RestElement>;
-      handleAssignment(currentLambda, contextIdentifiers, path.get("argument"));
+      handleAssignment(reassigned, path.get("argument"));
       break;
     }
     case "MemberExpression": {

compiler/forget/packages/babel-plugin-react-forget/src/Utils/utils.ts

@@ -57,3 +57,13 @@ export function getOrInsertDefault<U, V>(
     return defaultValue;
   }
 }
+
+export function Set_union<T>(a: Set<T>, b: Set<T>): Set<T> {
+  const union = new Set<T>();
+  for (const item of a) {
+    if (b.has(item)) {
+      union.add(item);
+    }
+  }
+  return union;
+}

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-declaration-reassign.expect.md

@@ -17,15 +17,24 @@ function component() {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function component() {
-  const $ = useMemoCache(1);
+  const $ = useMemoCache(2);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = {};
+    t0 = function x(a) {
+      a.foo();
+    };
     $[0] = t0;
   } else {
     t0 = $[0];
   }
-  const x = t0;
+  let x;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    x = t0;
+    x = {};
+    $[1] = x;
+  } else {
+    x = $[1];
+  }
   return x;
 }
 

compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md

@@ -0,0 +1,46 @@
+
+## Input
+
+```javascript
+// @debug
+function Component(props) {
+  let x = null;
+  const onChange = (e) => {
+    console.log(x);
+  };
+  x = {};
+  return <Foo onChange={onChange} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @debug
+function Component(props) {
+  const $ = useMemoCache(2);
+  let onChange;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    let x;
+    x = null;
+    onChange = (e) => {
+      console.log(x);
+    };
+    x = {};
+    $[0] = onChange;
+  } else {
+    onChange = $[0];
+  }
+  let t0;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Foo onChange={onChange} />;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js

@@ -0,0 +1,9 @@
+// @debug
+function Component(props) {
+  let x = null;
+  const onChange = (e) => {
+    console.log(x);
+  };
+  x = {};
+  return <Foo onChange={onChange} />;
+}