[patch][contextvars] Patch for variables reassigned after objectmethod

definitions

Author: mofeiZ

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

@@ -11,91 +11,58 @@ import { CompilerError } from "../CompilerError";
 import { Set_union } from "../Utils/utils";
 import { GeneratedSource } from "./HIR";
 
+type BabelFunction =
+  | NodePath<t.FunctionDeclaration>
+  | NodePath<t.FunctionExpression>
+  | NodePath<t.ArrowFunctionExpression>
+  | NodePath<t.ObjectMethod>;
 type FindContextIdentifierState = {
-  currentLambda: Array<
-    | NodePath<t.FunctionDeclaration>
-    | NodePath<t.FunctionExpression>
-    | NodePath<t.ArrowFunctionExpression>
-    | NodePath<t.ObjectMethod>
-  >;
+  currentFn: Array<BabelFunction>;
   reassigned: Set<t.Identifier>;
   referenced: Set<t.Identifier>;
 };
 
+const withFunctionScope = {
+  enter: function (
+    path: BabelFunction,
+    state: FindContextIdentifierState
+  ): void {
+    state.currentFn.push(path);
+  },
+  exit: function (_: BabelFunction, state: FindContextIdentifierState): void {
+    state.currentFn.pop();
+  },
+};
+
 export function findContextIdentifiers(
   func: NodePath<t.Function>
 ): Set<t.Identifier> {
   const state: FindContextIdentifierState = {
-    currentLambda: [],
+    currentFn: [],
     reassigned: new Set(),
     referenced: new Set(),
   };
 
   func.traverse<FindContextIdentifierState>(
     {
-      FunctionDeclaration: {
-        enter(
-          fn: NodePath<t.FunctionDeclaration>,
-          state: FindContextIdentifierState
-        ): void {
-          state.currentLambda.push(fn);
-        },
-        exit(
-          fn: NodePath<t.FunctionDeclaration>,
-          state: FindContextIdentifierState
-        ): void {
-          state.currentLambda.pop();
-        },
-      },
-      FunctionExpression: {
-        enter(
-          fn: NodePath<t.FunctionExpression>,
-          state: FindContextIdentifierState
-        ): void {
-          state.currentLambda.push(fn);
-        },
-        exit(
-          _fn: NodePath<t.FunctionExpression>,
-          state: FindContextIdentifierState
-        ): void {
-          state.currentLambda.pop();
-        },
-      },
-
-      ArrowFunctionExpression: {
-        enter(
-          fn: NodePath<t.ArrowFunctionExpression>,
-          state: FindContextIdentifierState
-        ): void {
-          state.currentLambda.push(fn);
-        },
-        exit(
-          _fn: NodePath<t.ArrowFunctionExpression>,
-          state: FindContextIdentifierState
-        ): void {
-          state.currentLambda.pop();
-        },
-      },
+      FunctionDeclaration: withFunctionScope,
+      FunctionExpression: withFunctionScope,
+      ArrowFunctionExpression: withFunctionScope,
+      ObjectMethod: withFunctionScope,
       AssignmentExpression(
         path: NodePath<t.AssignmentExpression>,
         state: FindContextIdentifierState
       ): void {
         const left = path.get("left");
         handleAssignment(state.reassigned, left);
       },
-      ObjectMethod(
-        fn: NodePath<t.ObjectMethod>,
-        state: FindContextIdentifierState
-      ): void {
-        state.currentLambda.push(fn);
-      },
       Identifier(
         path: NodePath<t.Identifier>,
         state: FindContextIdentifierState
       ): void {
-        const currentLambda = state.currentLambda.at(-1);
-        if (currentLambda !== undefined)
-          handleIdentifier(currentLambda, state.referenced, path);
+        const currentFn = state.currentFn.at(-1);
+        if (currentFn !== undefined)
+          handleIdentifier(currentFn, state.referenced, path);
       },
     },
     state
@@ -104,17 +71,13 @@ export function findContextIdentifiers(
 }
 
 function handleIdentifier(
-  currentLambda:
-    | NodePath<t.FunctionDeclaration>
-    | NodePath<t.FunctionExpression>
-    | NodePath<t.ArrowFunctionExpression>
-    | NodePath<t.ObjectMethod>,
+  currentFn: BabelFunction,
   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);
+  const bindingAboveLambdaScope = currentFn.scope.parent.getBinding(name);
 
   if (binding != null && binding === bindingAboveLambdaScope) {
     referenced.add(binding.identifier);
@@ -126,7 +89,7 @@ function handleAssignment(
   lvalPath: NodePath<t.LVal>
 ): void {
   /*
-   * Find all reassignments to identifiers declared outside of currentLambda
+   * Find all reassignments to identifiers declared outside of currentFn
    * This closely follows destructuring assignment assumptions and logic in BuildHIR
    */
   const lvalNode = lvalPath.node;

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-prop-across-objectmethod-def.expect.md

@@ -0,0 +1,51 @@
+
+## Input
+
+```javascript
+import { identity } from "shared-runtime";
+
+// repro for context identifier scoping bug, in which x was
+// inferred as a context variable.
+
+function Component() {
+  let x = 2;
+  const obj = {
+    method() {},
+  };
+  x = 4;
+  identity(obj);
+  // constant propagation should return 4 here
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { identity } from "shared-runtime";
+
+// repro for context identifier scoping bug, in which x was
+// inferred as a context variable.
+
+function Component() {
+  const obj = { method() {} };
+
+  identity(obj);
+  return 4;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) 4

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-prop-across-objectmethod-def.js

@@ -0,0 +1,20 @@
+import { identity } from "shared-runtime";
+
+// repro for context identifier scoping bug, in which x was
+// inferred as a context variable.
+
+function Component() {
+  let x = 2;
+  const obj = {
+    method() {},
+  };
+  x = 4;
+  identity(obj);
+  // constant propagation should return 4 here
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};

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

@@ -0,0 +1,64 @@
+
+## Input
+
+```javascript
+import { invoke } from "shared-runtime";
+
+function Component({ cond }) {
+  let x = 2;
+  const obj = {
+    method(cond) {
+      if (cond) {
+        x = 4;
+      }
+    },
+  };
+  invoke(obj.method, cond);
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+import { invoke } from "shared-runtime";
+
+function Component(t24) {
+  const $ = useMemoCache(2);
+  const { cond } = t24;
+  let x;
+  if ($[0] !== cond) {
+    x = 2;
+    const obj = {
+      method(cond_0) {
+        if (cond_0) {
+          x = 4;
+        }
+      },
+    };
+
+    invoke(obj.method, cond);
+    $[0] = cond;
+    $[1] = x;
+  } else {
+    x = $[1];
+  }
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true }],
+};
+
+```
+      
+### Eval output
+(kind: ok) 4

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-objectmethod.js

@@ -0,0 +1,19 @@
+import { invoke } from "shared-runtime";
+
+function Component({ cond }) {
+  let x = 2;
+  const obj = {
+    method(cond) {
+      if (cond) {
+        x = 4;
+      }
+    },
+  };
+  invoke(obj.method, cond);
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true }],
+};