Fix handling of this&co in computed keys in arrows transform (#14005)

packages/babel-helper-create-class-features-plugin/package.json

@@ -19,6 +19,7 @@
   ],
   "dependencies": {
     "@babel/helper-annotate-as-pure": "workspace:^",
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-function-name": "workspace:^",
     "@babel/helper-member-expression-to-functions": "workspace:^",
     "@babel/helper-optimise-call-expression": "workspace:^",

packages/babel-helper-create-class-features-plugin/src/fields.ts

@@ -1,9 +1,8 @@
 import { template, traverse, types as t } from "@babel/core";
 import type { File } from "@babel/core";
 import type { NodePath, Visitor, Scope } from "@babel/traverse";
-import ReplaceSupers, {
+import ReplaceSupers from "@babel/helper-replace-supers";
-  environmentVisitor,
+import environmentVisitor from "@babel/helper-environment-visitor";
-} from "@babel/helper-replace-supers";
 import memberExpressionToFunctions from "@babel/helper-member-expression-to-functions";
 import type {
   Handler,
@@ -851,7 +850,11 @@ function buildPrivateMethodDeclaration(
   );
 }
 
-const thisContextVisitor = traverse.visitors.merge([
+const thisContextVisitor = traverse.visitors.merge<{
+  classRef: t.Identifier;
+  needsClassRef: boolean;
+  innerBinding: t.Identifier;
+}>([
   {
     ThisExpression(path, state) {
       state.needsClassRef = true;

packages/babel-helper-create-class-features-plugin/src/misc.ts

@@ -1,9 +1,9 @@
 import { template, traverse, types as t } from "@babel/core";
 import type { File } from "@babel/core";
 import type { NodePath, Scope, Visitor, Binding } from "@babel/traverse";
-import { environmentVisitor } from "@babel/helper-replace-supers";
+import environmentVisitor from "@babel/helper-environment-visitor";
 
-const findBareSupers = traverse.visitors.merge([
+const findBareSupers = traverse.visitors.merge<NodePath<t.CallExpression>[]>([
   {
     Super(path: NodePath<t.Super>) {
       const { node, parentPath } = path;

packages/babel-helper-environment-visitor/.npmignore

@@ -0,0 +1,3 @@
+src
+test
+*.log

packages/babel-helper-environment-visitor/README.md

@@ -0,0 +1,19 @@
+# @babel/helper-environment-visitor
+
+> Helper visitor to only visit nodes in the current 'this' context
+
+See our website [@babel/helper-environment-visitor](https://babeljs.io/docs/en/babel-helper-environment-visitor) for more information.
+
+## Install
+
+Using npm:
+
+```sh
+npm install --save-dev @babel/helper-environment-visitor
+```
+
+or using yarn:
+
+```sh
+yarn add @babel/helper-environment-visitor --dev
+```

packages/babel-helper-environment-visitor/package.json

@@ -0,0 +1,29 @@
+{
+  "name": "@babel/helper-environment-visitor",
+  "version": "7.16.0",
+  "description": "Helper visitor to only visit nodes in the current 'this' context",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/babel/babel.git",
+    "directory": "packages/babel-helper-environment-visitor"
+  },
+  "homepage": "https://babel.dev/docs/en/next/babel-helper-environment-visitor",
+  "license": "MIT",
+  "publishConfig": {
+    "access": "public"
+  },
+  "main": "./lib/index.js",
+  "exports": {
+    ".": "./lib/index.js"
+  },
+  "dependencies": {
+    "@babel/types": "workspace:^"
+  },
+  "devDependencies": {
+    "@babel/traverse": "workspace:^"
+  },
+  "engines": {
+    "node": ">=6.9.0"
+  },
+  "author": "The Babel Team (https://babel.dev/team)"
+}

packages/babel-helper-environment-visitor/src/index.ts

@@ -0,0 +1,38 @@
+import type { NodePath } from "@babel/traverse";
+import { VISITOR_KEYS, staticBlock } from "@babel/types";
+import type * as t from "@babel/types";
+
+// TODO (Babel 8): Don't export this function.
+export function skipAllButComputedKey(
+  path: NodePath<t.Method | t.ClassProperty>,
+) {
+  // If the path isn't computed, just skip everything.
+  if (!path.node.computed) {
+    path.skip();
+    return;
+  }
+
+  // So it's got a computed key. Make sure to skip every other key the
+  // traversal would visit.
+  const keys = VISITOR_KEYS[path.type];
+  for (const key of keys) {
+    if (key !== "key") path.skipKey(key);
+  }
+}
+
+// Methods are handled by the Method visitor; arrows are not skipped because they inherit the context.
+const skipKey = process.env.BABEL_8_BREAKING
+  ? "StaticBlock|ClassPrivateProperty|TypeAnnotation|FunctionDeclaration|FunctionExpression"
+  : (staticBlock ? "StaticBlock|" : "") +
+    "ClassPrivateProperty|TypeAnnotation|FunctionDeclaration|FunctionExpression";
+
+// environmentVisitor should be used when traversing the whole class and not for specific class elements/methods.
+// For perf reasons, the environmentVisitor might be traversed with `{ noScope: true }`, which means `path.scope` is undefined.
+// Avoid using `path.scope` here
+export default {
+  [skipKey]: path => path.skip(),
+
+  "Method|ClassProperty"(path: NodePath<t.Method | t.ClassProperty>) {
+    skipAllButComputedKey(path);
+  },
+};

packages/babel-helper-member-expression-to-functions/src/index.ts

@@ -538,7 +538,7 @@ export interface HandlerState<State = {}> extends Handler<State> {
   handle(
     this: HandlerState<State> & State,
     member: Member,
-    noDocumentAll: boolean,
+    noDocumentAll?: boolean,
   ): void;
   memoiser: AssignmentMemoiser;
 }

packages/babel-helper-module-transforms/package.json

@@ -15,8 +15,8 @@
   },
   "main": "./lib/index.js",
   "dependencies": {
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-module-imports": "workspace:^",
-    "@babel/helper-replace-supers": "workspace:^",
     "@babel/helper-simple-access": "workspace:^",
     "@babel/helper-split-export-declaration": "workspace:^",
     "@babel/helper-validator-identifier": "workspace:^",

packages/babel-helper-module-transforms/src/rewrite-this.ts

@@ -1,4 +1,4 @@
-import { environmentVisitor } from "@babel/helper-replace-supers";
+import environmentVisitor from "@babel/helper-environment-visitor";
 import traverse from "@babel/traverse";
 import { numericLiteral, unaryExpression } from "@babel/types";
 import type * as t from "@babel/types";

packages/babel-helper-replace-supers/package.json

@@ -14,6 +14,7 @@
   },
   "main": "./lib/index.js",
   "dependencies": {
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-member-expression-to-functions": "workspace:^",
     "@babel/helper-optimise-call-expression": "workspace:^",
     "@babel/traverse": "workspace:^",

packages/babel-helper-replace-supers/src/index.ts

@@ -1,9 +1,10 @@
 import type { HubInterface, NodePath, Scope } from "@babel/traverse";
 import traverse from "@babel/traverse";
 import memberExpressionToFunctions from "@babel/helper-member-expression-to-functions";
+import type { HandlerState } from "@babel/helper-member-expression-to-functions";
 import optimiseCall from "@babel/helper-optimise-call-expression";
+import environmentVisitor from "@babel/helper-environment-visitor";
 import {
-  VISITOR_KEYS,
   assignmentExpression,
   booleanLiteral,
   callExpression,
@@ -11,12 +12,17 @@ import {
   identifier,
   memberExpression,
   sequenceExpression,
-  staticBlock,
   stringLiteral,
   thisExpression,
 } from "@babel/types";
 import type * as t from "@babel/types";
 
+// TODO (Babel 8): Don't export this.
+export {
+  default as environmentVisitor,
+  skipAllButComputedKey,
+} from "@babel/helper-environment-visitor";
+
 /**
  * Creates an expression which result is the proto of objectRef.
  *
@@ -38,49 +44,9 @@ function getPrototypeOfExpression(objectRef, isStatic, file, isPrivateMethod) {
   return callExpression(file.addHelper("getPrototypeOf"), [targetRef]);
 }
 
-export function skipAllButComputedKey(
+const visitor = traverse.visitors.merge<
-  path: NodePath<t.Method | t.ClassProperty | t.ClassPrivateProperty>,
+  HandlerState<ReplaceState> & ReplaceState
-) {
+>([
-  // If the path isn't computed, just skip everything.
-  // @ts-expect-error todo(flow->ts) check node type before cheking the property
-  if (!path.node.computed) {
-    path.skip();
-    return;
-  }
-
-  // So it's got a computed key. Make sure to skip every other key the
-  // traversal would visit.
-  const keys = VISITOR_KEYS[path.type];
-  for (const key of keys) {
-    if (key !== "key") path.skipKey(key);
-  }
-}
-
-// environmentVisitor should be used when traversing the whole class and not for specific class elements/methods.
-// For perf reasons, the environmentVisitor will be traversed with `{ noScope: true }`, which means `path.scope` is undefined.
-// Avoid using `path.scope` here
-export const environmentVisitor = {
-  // todo (Babel 8): remove StaticBlock brand checks
-  [`${staticBlock ? "StaticBlock|" : ""}ClassPrivateProperty|TypeAnnotation`](
-    path: NodePath,
-  ) {
-    path.skip();
-  },
-
-  Function(path: NodePath) {
-    // Methods will be handled by the Method visit
-    if (path.isMethod()) return;
-    // Arrow functions inherit their parent's environment
-    if (path.isArrowFunctionExpression()) return;
-    path.skip();
-  },
-
-  "Method|ClassProperty"(path: NodePath<t.Method | t.ClassProperty>) {
-    skipAllButComputedKey(path);
-  },
-};
-
-const visitor = traverse.visitors.merge([
   environmentVisitor,
   {
     Super(path, state) {

packages/babel-plugin-transform-classes/package.json

@@ -15,6 +15,7 @@
   "main": "./lib/index.js",
   "dependencies": {
     "@babel/helper-annotate-as-pure": "workspace:^",
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-function-name": "workspace:^",
     "@babel/helper-optimise-call-expression": "workspace:^",
     "@babel/helper-plugin-utils": "workspace:^",

packages/babel-plugin-transform-classes/src/transformClass.ts

@@ -1,8 +1,7 @@
 import type { NodePath, Visitor } from "@babel/traverse";
 import nameFunction from "@babel/helper-function-name";
-import ReplaceSupers, {
+import ReplaceSupers from "@babel/helper-replace-supers";
-  environmentVisitor,
+import environmentVisitor from "@babel/helper-environment-visitor";
-} from "@babel/helper-replace-supers";
 import optimiseCall from "@babel/helper-optimise-call-expression";
 import { traverse, template, types as t } from "@babel/core";
 import annotateAsPure from "@babel/helper-annotate-as-pure";

packages/babel-traverse/package.json

@@ -18,6 +18,7 @@
   "dependencies": {
     "@babel/code-frame": "workspace:^",
     "@babel/generator": "workspace:^",
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-function-name": "workspace:^",
     "@babel/helper-hoist-variables": "workspace:^",
     "@babel/helper-split-export-declaration": "workspace:^",

packages/babel-traverse/src/path/conversion.ts

@@ -26,7 +26,9 @@ import {
   unaryExpression,
 } from "@babel/types";
 import type * as t from "@babel/types";
+import environmentVisitor from "@babel/helper-environment-visitor";
 import nameFunction from "@babel/helper-function-name";
+import { merge as mergeVisitors } from "../visitors";
 import type NodePath from "./index";
 
 export function toComputedKey(this: NodePath) {
@@ -190,6 +192,18 @@ export function arrowFunctionToExpression(
   }
 }
 
+const getSuperCallsVisitor = mergeVisitors<{
+  allSuperCalls: NodePath<t.CallExpression>[];
+}>([
+  {
+    CallExpression(child, { allSuperCalls }) {
+      if (!child.get("callee").isSuper()) return;
+      allSuperCalls.push(child);
+    },
+  },
+  environmentVisitor,
+]);
+
 /**
  * Given a function, traverse its contents, and if there are references to "this", "arguments", "super",
  * or "new.target", ensure that these references reference the parent environment around this function.
@@ -252,20 +266,8 @@ function hoistFunctionEnvironment(
         "Unable to handle nested super() usage in arrow",
       );
     }
-    const allSuperCalls = [];
+    const allSuperCalls: NodePath<t.CallExpression>[] = [];
-    thisEnvFn.traverse({
+    thisEnvFn.traverse(getSuperCallsVisitor, { allSuperCalls });
-      Function(child) {
-        if (child.isArrowFunctionExpression()) return;
-        child.skip();
-      },
-      ClassProperty(child) {
-        child.skip();
-      },
-      CallExpression(child) {
-        if (!child.get("callee").isSuper()) return;
-        allSuperCalls.push(child);
-      },
-    });
     const superBinding = getSuperBinding(thisEnvFn);
     allSuperCalls.forEach(superCall => {
       const callee = identifier(superBinding);
@@ -508,34 +510,33 @@ function hasSuperClass(thisEnvFn) {
   );
 }
 
+const assignSuperThisVisitor = mergeVisitors<{
+  supers: WeakSet<t.CallExpression>;
+  thisBinding: string;
+}>([
+  {
+    CallExpression(child, { supers, thisBinding }) {
+      if (!child.get("callee").isSuper()) return;
+      if (supers.has(child.node)) return;
+      supers.add(child.node);
+
+      child.replaceWithMultiple([
+        child.node,
+        assignmentExpression("=", identifier(thisBinding), identifier("this")),
+      ]);
+    },
+  },
+  environmentVisitor,
+]);
+
 // Create a binding that evaluates to the "this" of the given function.
 function getThisBinding(thisEnvFn, inConstructor) {
   return getBinding(thisEnvFn, "this", thisBinding => {
     if (!inConstructor || !hasSuperClass(thisEnvFn)) return thisExpression();
 
-    const supers = new WeakSet();
+    thisEnvFn.traverse(assignSuperThisVisitor, {
-    thisEnvFn.traverse({
+      supers: new WeakSet(),
-      Function(child) {
+      thisBinding,
-        if (child.isArrowFunctionExpression()) return;
-        child.skip();
-      },
-      ClassProperty(child) {
-        child.skip();
-      },
-      CallExpression(child) {
-        if (!child.get("callee").isSuper()) return;
-        if (supers.has(child.node)) return;
-        supers.add(child.node);
-
-        child.replaceWithMultiple([
-          child.node,
-          assignmentExpression(
-            "=",
-            identifier(thisBinding),
-            identifier("this"),
-          ),
-        ]);
-      },
     });
   });
 }
@@ -601,25 +602,18 @@ function getBinding(thisEnvFn, key, init) {
   return data;
 }
 
-function getScopeInformation(fnPath) {
+const getScopeInformationVisitor = mergeVisitors<{
-  const thisPaths = [];
+  thisPaths: NodePath<t.ThisExpression | t.JSXIdentifier>[];
-  const argumentsPaths = [];
+  superCalls: NodePath<t.CallExpression>[];
-  const newTargetPaths = [];
+  superProps: NodePath<t.MemberExpression>[];
-  const superProps = [];
+  argumentsPaths: NodePath<t.Identifier | t.JSXIdentifier>[];
-  const superCalls = [];
+  newTargetPaths: NodePath<t.MetaProperty>[];
-
+}>([
-  fnPath.traverse({
+  {
-    ClassProperty(child) {
+    ThisExpression(child, { thisPaths }) {
-      child.skip();
-    },
-    Function(child) {
-      if (child.isArrowFunctionExpression()) return;
-      child.skip();
-    },
-    ThisExpression(child) {
       thisPaths.push(child);
     },
-    JSXIdentifier(child) {
+    JSXIdentifier(child, { thisPaths }) {
       if (child.node.name !== "this") return;
       if (
         !child.parentPath.isJSXMemberExpression({ object: child.node }) &&
@@ -630,14 +624,14 @@ function getScopeInformation(fnPath) {
 
       thisPaths.push(child);
     },
-    CallExpression(child) {
+    CallExpression(child, { superCalls }) {
       if (child.get("callee").isSuper()) superCalls.push(child);
     },
-    MemberExpression(child) {
+    MemberExpression(child, { superProps }) {
       if (child.get("object").isSuper()) superProps.push(child);
     },
-    ReferencedIdentifier(child) {
+    Identifier(child, { argumentsPaths }) {
-      if (child.node.name !== "arguments") return;
+      if (!child.isReferencedIdentifier({ name: "arguments" })) return;
 
       let curr = child.scope;
       do {
@@ -652,12 +646,29 @@ function getScopeInformation(fnPath) {
 
       argumentsPaths.push(child);
     },
-    MetaProperty(child) {
+    MetaProperty(child, { newTargetPaths }) {
       if (!child.get("meta").isIdentifier({ name: "new" })) return;
       if (!child.get("property").isIdentifier({ name: "target" })) return;
 
       newTargetPaths.push(child);
     },
+  },
+  environmentVisitor,
+]);
+
+function getScopeInformation(fnPath) {
+  const thisPaths = [];
+  const argumentsPaths = [];
+  const newTargetPaths = [];
+  const superProps = [];
+  const superCalls = [];
+
+  fnPath.traverse(getScopeInformationVisitor, {
+    thisPaths,
+    argumentsPaths,
+    newTargetPaths,
+    superProps,
+    superCalls,
   });
 
   return {

packages/babel-traverse/src/visitors.ts

@@ -1,5 +1,6 @@
 import * as virtualTypes from "./path/lib/virtual-types";
 import { DEPRECATED_KEYS, FLIPPED_ALIAS_KEYS, TYPES } from "@babel/types";
+import type { Visitor } from "./types";
 
 /**
  * explode() will take a visitor object with all of the various shorthands
@@ -175,6 +176,12 @@ function validateVisitorMethods(path, val) {
   }
 }
 
+export function merge<State>(visitors: Visitor<State>[]): Visitor<State>;
+export function merge(
+  visitors: Visitor<unknown>[],
+  states?: any[],
+  wrapper?: Function | null,
+): Visitor<unknown>;
 export function merge(
   visitors: any[],
   states: any[] = [],

packages/babel-traverse/test/arrow-transform.js

@@ -71,6 +71,30 @@ describe("arrow function conversion", () => {
     );
   });
 
+  it("should convert super calls in constructors used in computed method keys", () => {
+    assertConversion(
+      `
+      () => ({
+        [super()]: 123,
+        get [super() + "1"]() {},
+        [super()]() {},
+      });
+    `,
+      `
+      var _supercall = (..._args) => super(..._args);
+
+      (function () {
+        return {
+          [_supercall()]: 123,
+          get [_supercall() + "1"]() {},
+          [_supercall()]() {},
+        };
+      });
+    `,
+      { methodName: "constructor", extend: true },
+    );
+  });
+
   it("should convert super calls and this references in constructors", () => {
     assertConversion(
       `
@@ -127,6 +151,32 @@ describe("arrow function conversion", () => {
     );
   });
 
+  it("should convert this references in constructors with super() used in computed method keys", () => {
+    assertConversion(
+      `
+      () => this;
+
+      ({
+        [super()]: 123,
+        get [super() + "1"]() {},
+        [super()]() {},
+      });
+    `,
+      `
+      var _this;
+
+      (function () { return _this; });
+
+      ({
+        [(super(), _this = this)]: 123,
+        get [(super(), _this = this) + "1"]() {},
+        [(super(), _this = this)]() {},
+      });
+    `,
+      { methodName: "constructor", extend: true },
+    );
+  });
+
   it("should convert this references in constructors with spec compliance", () => {
     assertConversion(
       `
@@ -273,6 +323,29 @@ describe("arrow function conversion", () => {
     );
   });
 
+  it("should convert this references used in computed method keys", () => {
+    assertConversion(
+      `
+      () => ({
+        [this]: 123,
+        get [this + "1"]() {},
+        [this]() {},
+      });
+    `,
+      `
+      var _this = this;
+
+      (function () {
+        return {
+          [_this]: 123,
+          get [_this + "1"]() {},
+          [_this]() {},
+        };
+      });
+    `,
+    );
+  });
+
   it("should convert arguments references", () => {
     assertConversion(
       `

tsconfig.json

@@ -12,6 +12,7 @@
     "./packages/babel-helper-create-class-features-plugin/src/**/*.ts",
     "./packages/babel-helper-create-regexp-features-plugin/src/**/*.ts",
     "./packages/babel-helper-define-map/src/**/*.ts",
+    "./packages/babel-helper-environment-visitor/src/**/*.ts",
     "./packages/babel-helper-explode-assignable-expression/src/**/*.ts",
     "./packages/babel-helper-fixtures/src/**/*.ts",
     "./packages/babel-helper-function-name/src/**/*.ts",
@@ -177,6 +178,9 @@
       "@babel/helper-define-map": [
         "./packages/babel-helper-define-map/src"
       ],
+      "@babel/helper-environment-visitor": [
+        "./packages/babel-helper-environment-visitor/src"
+      ],
       "@babel/helper-explode-assignable-expression": [
         "./packages/babel-helper-explode-assignable-expression/src"
       ],

yarn.lock

@@ -546,6 +546,7 @@ __metadata:
   dependencies:
     "@babel/core": "workspace:^"
     "@babel/helper-annotate-as-pure": "workspace:^"
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-function-name": "workspace:^"
     "@babel/helper-member-expression-to-functions": "workspace:^"
     "@babel/helper-optimise-call-expression": "workspace:^"
@@ -611,6 +612,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@babel/helper-environment-visitor@workspace:^, @babel/helper-environment-visitor@workspace:packages/babel-helper-environment-visitor":
+  version: 0.0.0-use.local
+  resolution: "@babel/helper-environment-visitor@workspace:packages/babel-helper-environment-visitor"
+  dependencies:
+    "@babel/traverse": "workspace:^"
+    "@babel/types": "workspace:^"
+  languageName: unknown
+  linkType: soft
+
 "@babel/helper-explode-assignable-expression@npm:^7.16.0":
   version: 7.16.0
   resolution: "@babel/helper-explode-assignable-expression@npm:7.16.0"
@@ -735,8 +745,8 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@babel/helper-module-transforms@workspace:packages/babel-helper-module-transforms"
   dependencies:
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-module-imports": "workspace:^"
-    "@babel/helper-replace-supers": "workspace:^"
     "@babel/helper-simple-access": "workspace:^"
     "@babel/helper-split-export-declaration": "workspace:^"
     "@babel/helper-validator-identifier": "workspace:^"
@@ -849,6 +859,7 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@babel/helper-replace-supers@workspace:packages/babel-helper-replace-supers"
   dependencies:
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-member-expression-to-functions": "workspace:^"
     "@babel/helper-optimise-call-expression": "workspace:^"
     "@babel/traverse": "workspace:^"
@@ -2193,6 +2204,7 @@ __metadata:
   dependencies:
     "@babel/core": "workspace:^"
     "@babel/helper-annotate-as-pure": "workspace:^"
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-function-name": "workspace:^"
     "@babel/helper-optimise-call-expression": "workspace:^"
     "@babel/helper-plugin-test-runner": "workspace:^"
@@ -3667,6 +3679,7 @@ __metadata:
   dependencies:
     "@babel/code-frame": "workspace:^"
     "@babel/generator": "workspace:^"
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-function-name": "workspace:^"
     "@babel/helper-hoist-variables": "workspace:^"
     "@babel/helper-plugin-test-runner": "workspace:^"