[babel][contextvar] Patch context identifier babel logic; only use referenced

identifiers 

--- 

A few fixes for finding context identifiers: 

Previously, we counted every babel identifier as a reference. This is 
problematic because babel counts every string symbol as an identifier. 

```js 

print(x);  // x is an identifier as expected 

obj.x      // x is.. also an identifier here 

{x: 2}     // x is also an identifier here 

``` 

This PR adds a check for `isReferencedIdentifier`. Note that only non-lval 
references pass this check 

```js 

print(x);  // isReferencedIdentifier(x) -> true 

obj.x      // isReferencedIdentifier(x) -> false 

{x: 2}     // isReferencedIdentifier(x) -> false 

x = 2      // isReferencedIdentifier(x) -> false 

``` 

Which brings us to change #2. 

Previously, we counted assignments as references due to the identifier visiting 
+ checking logic. The logic was roughly the following (from facebook#1691) 

```js 

contextVars = intersection(reassigned, referencedByInnerFn); 

``` 

Now that assignments (lvals) and references (rvals) are tracked separately, the 
equivalent logic is this. Note that assignment to a context variable does not 
need to be modeled as a read (`console.log(x = 5)` always will evaluates and 
prints 5, regardless of the previous value of x). 

``` 

contextVars = union(reassignedByInnerFn, intersection(reassigned, 
referencedByInnerFn)) 

``` 

--- 

Note that variables that are never read do not need to be modeled as context 
variables, but this is unlikely to be a common pattern. 

```js 

function fn() { 

let x = 2; 

const inner = () => { 

x = 3; 

} 

} 

```

Author: mofeiZ

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

@@ -8,18 +8,28 @@
 import type { NodePath } from "@babel/traverse";
 import type * as t from "@babel/types";
 import { CompilerError } from "../CompilerError";
-import { Set_union } from "../Utils/utils";
+import { getOrInsertDefault } from "../Utils/utils";
 import { GeneratedSource } from "./HIR";
 
+type IdentifierInfo = {
+  reassigned: boolean;
+  reassignedByInnerFn: boolean;
+  referencedByInnerFn: boolean;
+};
+const DEFAULT_IDENTIFIER_INFO: IdentifierInfo = {
+  reassigned: false,
+  reassignedByInnerFn: false,
+  referencedByInnerFn: false,
+};
+
 type BabelFunction =
   | NodePath<t.FunctionDeclaration>
   | NodePath<t.FunctionExpression>
   | NodePath<t.ArrowFunctionExpression>
   | NodePath<t.ObjectMethod>;
 type FindContextIdentifierState = {
   currentFn: Array<BabelFunction>;
-  reassigned: Set<t.Identifier>;
-  referenced: Set<t.Identifier>;
+  identifiers: Map<t.Identifier, IdentifierInfo>;
 };
 
 const withFunctionScope = {
@@ -39,8 +49,7 @@ export function findContextIdentifiers(
 ): Set<t.Identifier> {
   const state: FindContextIdentifierState = {
     currentFn: [],
-    reassigned: new Set(),
-    referenced: new Set(),
+    identifiers: new Map(),
   };
 
   func.traverse<FindContextIdentifierState>(
@@ -54,38 +63,59 @@ export function findContextIdentifiers(
         state: FindContextIdentifierState
       ): void {
         const left = path.get("left");
-        handleAssignment(state.reassigned, left);
+        const currentFn = state.currentFn.at(-1) ?? null;
+        handleAssignment(currentFn, state.identifiers, left);
       },
       Identifier(
         path: NodePath<t.Identifier>,
         state: FindContextIdentifierState
       ): void {
-        const currentFn = state.currentFn.at(-1);
-        if (currentFn !== undefined)
-          handleIdentifier(currentFn, state.referenced, path);
+        const currentFn = state.currentFn.at(-1) ?? null;
+        if (path.isReferencedIdentifier()) {
+          handleIdentifier(currentFn, state.identifiers, path);
+        }
       },
     },
     state
   );
-  return Set_union(state.reassigned, state.referenced);
+
+  const result = new Set<t.Identifier>();
+  for (const [id, info] of state.identifiers.entries()) {
+    if (info.reassignedByInnerFn) {
+      result.add(id);
+    } else if (info.reassigned && info.referencedByInnerFn) {
+      result.add(id);
+    }
+  }
+  return result;
 }
 
 function handleIdentifier(
-  currentFn: BabelFunction,
-  referenced: Set<t.Identifier>,
+  currentFn: BabelFunction | null,
+  identifiers: Map<t.Identifier, IdentifierInfo>,
   path: NodePath<t.Identifier>
 ): void {
   const name = path.node.name;
   const binding = path.scope.getBinding(name);
-  const bindingAboveLambdaScope = currentFn.scope.parent.getBinding(name);
+  if (binding == null) {
+    return;
+  }
+  const identifier = getOrInsertDefault(identifiers, binding.identifier, {
+    ...DEFAULT_IDENTIFIER_INFO,
+  });
 
-  if (binding != null && binding === bindingAboveLambdaScope) {
-    referenced.add(binding.identifier);
+  if (currentFn != null) {
+    const bindingAboveLambdaScope = currentFn.scope.parent.getBinding(name);
+
+    if (binding === bindingAboveLambdaScope) {
+      identifier.referencedByInnerFn = true;
+    }
   }
 }
 
 function handleAssignment(
-  reassigned: Set<t.Identifier>,
+  currentFn: BabelFunction | null,
+  identifiers: Map<t.Identifier, IdentifierInfo>,
   lvalPath: NodePath<t.LVal>
 ): void {
   /*
@@ -98,16 +128,28 @@ function handleAssignment(
       const path = lvalPath as NodePath<t.Identifier>;
       const name = path.node.name;
       const binding = path.scope.getBinding(name);
-      if (binding != null) {
-        reassigned.add(binding.identifier);
+      if (binding == null) {
+        break;
+      }
+      const state = getOrInsertDefault(identifiers, binding.identifier, {
+        ...DEFAULT_IDENTIFIER_INFO,
+      });
+      state.reassigned = true;
+
+      if (currentFn != null) {
+        const bindingAboveLambdaScope = currentFn.scope.parent.getBinding(name);
+
+        if (binding === bindingAboveLambdaScope) {
+          state.reassignedByInnerFn = true;
+        }
       }
       break;
     }
     case "ArrayPattern": {
       const path = lvalPath as NodePath<t.ArrayPattern>;
       for (const element of path.get("elements")) {
         if (nonNull(element)) {
-          handleAssignment(reassigned, element);
+          handleAssignment(currentFn, identifiers, element);
         }
       }
       break;
@@ -123,28 +165,28 @@ function handleAssignment(
             loc: valuePath.node.loc ?? GeneratedSource,
             suggestions: null,
           });
-          handleAssignment(reassigned, valuePath);
+          handleAssignment(currentFn, identifiers, valuePath);
         } else {
           CompilerError.invariant(property.isRestElement(), {
             reason: `[FindContextIdentifiers] Invalid assumptions for babel types.`,
             description: null,
             loc: property.node.loc ?? GeneratedSource,
             suggestions: null,
           });
-          handleAssignment(reassigned, property);
+          handleAssignment(currentFn, identifiers, property);
         }
       }
       break;
     }
     case "AssignmentPattern": {
       const path = lvalPath as NodePath<t.AssignmentPattern>;
       const left = path.get("left");
-      handleAssignment(reassigned, left);
+      handleAssignment(currentFn, identifiers, left);
       break;
     }
     case "RestElement": {
       const path = lvalPath as NodePath<t.RestElement>;
-      handleAssignment(reassigned, path.get("argument"));
+      handleAssignment(currentFn, identifiers, path.get("argument"));
       break;
     }
     case "MemberExpression": {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-3.expect.md

@@ -17,8 +17,10 @@ function bar(a, b) {
 
 export const FIXTURE_ENTRYPOINT = {
   fn: bar,
-  params: ["TodoAdd"],
-  isComponent: "TodoAdd",
+  params: [
+    [1, 2],
+    [2, 3],
+  ],
 };
 
 ```
@@ -52,9 +54,13 @@ function bar(a, b) {
 
 export const FIXTURE_ENTRYPOINT = {
   fn: bar,
-  params: ["TodoAdd"],
-  isComponent: "TodoAdd",
+  params: [
+    [1, 2],
+    [2, 3],
+  ],
 };
 
 ```
-      
+      
+### Eval output
+(kind: ok) 2

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-3.js

@@ -13,6 +13,8 @@ function bar(a, b) {
 
 export const FIXTURE_ENTRYPOINT = {
   fn: bar,
-  params: ["TodoAdd"],
-  isComponent: "TodoAdd",
+  params: [
+    [1, 2],
+    [2, 3],
+  ],
 };

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/chained-assignment-context-variable.expect.md

@@ -2,22 +2,31 @@
 ## Input
 
 ```javascript
+import { makeArray } from "shared-runtime";
+
 function Component() {
   let x,
     y = (x = {});
   const foo = () => {
-    x = getObject();
+    x = makeArray();
   };
   foo();
   return [y, x];
 }
 
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
 ```
 
 ## Code
 
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
+import { makeArray } from "shared-runtime";
+
 function Component() {
   const $ = useMemoCache(3);
   let x;
@@ -26,7 +35,7 @@ function Component() {
     y = x = {};
 
     const foo = () => {
-      x = getObject();
+      x = makeArray();
     };
 
     foo();
@@ -46,5 +55,12 @@ function Component() {
   return t0;
 }
 
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
 ```
-      
+      
+### Eval output
+(kind: ok) [{},[]]

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/chained-assignment-context-variable.js

@@ -1,9 +1,16 @@
+import { makeArray } from "shared-runtime";
+
 function Component() {
   let x,
     y = (x = {});
   const foo = () => {
-    x = getObject();
+    x = makeArray();
   };
   foo();
   return [y, x];
 }
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-prop-colliding-identifier.expect.md

@@ -0,0 +1,44 @@
+
+## Input
+
+```javascript
+import { invoke } from "shared-runtime";
+
+function Component() {
+  let x = 2;
+  const fn = () => {
+    return { x: "value" };
+  };
+  invoke(fn);
+  x = 3;
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { invoke } from "shared-runtime";
+
+function Component() {
+  const fn = () => ({ x: "value" });
+
+  invoke(fn);
+  return 3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) 3

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-prop-colliding-identifier.js

@@ -0,0 +1,16 @@
+import { invoke } from "shared-runtime";
+
+function Component() {
+  let x = 2;
+  const fn = () => {
+    return { x: "value" };
+  };
+  invoke(fn);
+  x = 3;
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};