Improve decorated private method check (babel#16226)

* refactor: extract createPrivateMethodUpdateError

* fix: check private method update after transform

* fix: do not throw when a private name shadow decorated method

* cleanup typings

* review comments

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

@@ -5,6 +5,10 @@ import splitExportDeclaration from "@babel/helper-split-export-declaration";
 import * as charCodes from "charcodes";
 import type { PluginAPI, PluginObject, PluginPass } from "@babel/core";
 import { skipTransparentExprWrappers } from "@babel/helper-skip-transparent-expression-wrappers";
+import {
+  privateNameVisitorFactory,
+  type PrivateNameVisitorState,
+} from "./fields.ts";
 
 interface Options {
   /** @deprecated use `constantSuper` assumption instead. Only supported in 2021-12 version. */
@@ -532,8 +536,55 @@ function createToPropertyKeyCall(state: PluginPass, propertyKey: t.Expression) {
   return t.callExpression(state.addHelper("toPropertyKey"), [propertyKey]);
 }
 
+function checkPrivateMethodUpdateError(
+  path: NodePath<t.Class>,
+  decoratedPrivateMethods: Set<string>,
+) {
+  const privateNameVisitor = privateNameVisitorFactory<
+    PrivateNameVisitorState<null>,
+    null
+  >({
+    PrivateName(path, state) {
+      if (!state.privateNamesMap.has(path.node.id.name)) return;
+
+      const parentPath = path.parentPath;
+      const parentParentPath = parentPath.parentPath;
+
+      if (
+        // this.bar().#x = 123;
+        (parentParentPath.node.type === "AssignmentExpression" &&
+          parentParentPath.node.left === parentPath.node) ||
+        // this.#x++;
+        parentParentPath.node.type === "UpdateExpression" ||
+        // ([...this.#x] = foo);
+        parentParentPath.node.type === "RestElement" ||
+        // ([this.#x] = foo);
+        parentParentPath.node.type === "ArrayPattern" ||
+        // ({ a: this.#x } = bar);
+        (parentParentPath.node.type === "ObjectProperty" &&
+          parentParentPath.node.value === parentPath.node &&
+          parentParentPath.parentPath.type === "ObjectPattern") ||
+        // for (this.#x of []);
+        (parentParentPath.node.type === "ForOfStatement" &&
+          parentParentPath.node.left === parentPath.node)
+      ) {
+        throw path.buildCodeFrameError(
+          `Decorated private methods are read-only, but "#${path.node.id.name}" is updated via this expression.`,
+        );
+      }
+    },
+  });
+  const privateNamesMap = new Map<string, null>();
+  for (const name of decoratedPrivateMethods) {
+    privateNamesMap.set(name, null);
+  }
+  path.traverse(privateNameVisitor, {
+    privateNamesMap: privateNamesMap,
+  });
+}
+
 function transformClass(
-  path: NodePath<t.ClassExpression | t.ClassDeclaration>,
+  path: NodePath<t.Class>,
   state: PluginPass,
   constantSuper: boolean,
   version: DecoratorVersionKind,
@@ -1024,40 +1075,6 @@ function transformClass(
     elementLocals.push(staticInitLocal);
   }
 
-  if (decoratedPrivateMethods.size > 0) {
-    path.traverse({
-      PrivateName(path) {
-        if (!decoratedPrivateMethods.has(path.node.id.name)) return;
-
-        const parentPath = path.parentPath;
-        const parentParentPath = parentPath.parentPath;
-
-        if (
-          // this.bar().#x = 123;
-          (parentParentPath.node.type === "AssignmentExpression" &&
-            parentParentPath.node.left === parentPath.node) ||
-          // this.#x++;
-          parentParentPath.node.type === "UpdateExpression" ||
-          // ([...this.#x] = foo);
-          parentParentPath.node.type === "RestElement" ||
-          // ([this.#x] = foo);
-          parentParentPath.node.type === "ArrayPattern" ||
-          // ({ a: this.#x } = bar);
-          (parentParentPath.node.type === "ObjectProperty" &&
-            parentParentPath.node.value === parentPath.node &&
-            parentParentPath.parentPath.type === "ObjectPattern") ||
-          // for (this.#x of []);
-          (parentParentPath.node.type === "ForOfStatement" &&
-            parentParentPath.node.left === parentPath.node)
-        ) {
-          throw path.buildCodeFrameError(
-            `Decorated private methods are not updatable, but "#${path.node.id.name}" is updated via this expression.`,
-          );
-        }
-      },
-    });
-  }
-
   const classLocals: t.Identifier[] = [];
   let classInitInjected = false;
   const classInitCall =
@@ -1204,6 +1221,10 @@ function transformClass(
     );
   }
 
+  if (decoratedPrivateMethods.size > 0) {
+    checkPrivateMethodUpdateError(path, decoratedPrivateMethods);
+  }
+
   // Recrawl the scope to make sure new identifiers are properly synced
   path.scope.crawl();
 

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

@@ -26,7 +26,9 @@ interface PrivateNameMetadata {
   setterDeclared?: boolean;
 }
 
-type PrivateNamesMap = Map<string, PrivateNameMetadata>;
+type PrivateNamesMapGeneric<V> = Map<string, V>;
+
+type PrivateNamesMap = PrivateNamesMapGeneric<PrivateNameMetadata>;
 
 export function buildPrivateNamesMap(props: PropPath[]) {
   const privateNamesMap: PrivateNamesMap = new Map();
@@ -102,17 +104,16 @@ export function buildPrivateNamesNodes(
   return initNodes;
 }
 
-interface PrivateNameVisitorState {
-  privateNamesMap: PrivateNamesMap;
-  privateFieldsAsProperties: boolean;
+export interface PrivateNameVisitorState<V> {
+  privateNamesMap: PrivateNamesMapGeneric<V>;
   redeclared?: string[];
 }
 
 // Traverses the class scope, handling private name references. If an inner
 // class redeclares the same private name, it will hand off traversal to the
 // restricted visitor (which doesn't traverse the inner class's inner scope).
-function privateNameVisitorFactory<S>(
-  visitor: Visitor<PrivateNameVisitorState & S>,
+export function privateNameVisitorFactory<S, V>(
+  visitor: Visitor<PrivateNameVisitorState<V> & S>,
 ) {
   // Traverses the outer portion of a class, without touching the class's inner
   // scope, for private names.
@@ -123,7 +124,10 @@ function privateNameVisitorFactory<S>(
     environmentVisitor,
   ]);
 
-  const privateNameVisitor: Visitor<PrivateNameVisitorState & S> = {
+  // @ts-expect-error: TS2590: Expression produces a union type that is too complex to represent.
+  const privateNameVisitor: Visitor<
+    PrivateNameVisitorState<PrivateNameMetadata> & S
+  > = {
     ...visitor,
 
     Class(path) {
@@ -175,7 +179,8 @@ interface PrivateNameState {
 }
 
 const privateNameVisitor = privateNameVisitorFactory<
-  HandlerState<PrivateNameState> & PrivateNameState
+  HandlerState<PrivateNameState> & PrivateNameState,
+  PrivateNameMetadata
 >({
   PrivateName(path, { noDocumentAll }) {
     const { privateNamesMap, redeclared } = this;
@@ -222,11 +227,15 @@ export function buildCheckInRHS(
   return t.callExpression(file.addHelper("checkInRHS"), [rhs]);
 }
 
-const privateInVisitor = privateNameVisitorFactory<{
-  classRef: t.Identifier;
-  file: File;
-  innerBinding?: t.Identifier;
-}>({
+const privateInVisitor = privateNameVisitorFactory<
+  {
+    classRef: t.Identifier;
+    file: File;
+    innerBinding?: t.Identifier;
+    privateFieldsAsProperties: boolean;
+  },
+  PrivateNameMetadata
+>({
   BinaryExpression(path, { file }) {
     const { operator, left, right } = path.node;
     if (operator !== "in") return;

packages/babel-plugin-proposal-decorators/test/fixtures/2021-12-misc/setting-private-method-via-array-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2021-12-misc/setting-private-method-via-for-of/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2021-12-misc/setting-private-method-via-object-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2021-12-misc/setting-private-method-via-rest/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2021-12-misc/setting-private-method-via-update/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2021-12-misc/setting-private-method/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2022-03-misc/setting-private-method-via-array-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2022-03-misc/setting-private-method-via-for-of/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2022-03-misc/setting-private-method-via-object-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2022-03-misc/setting-private-method-via-rest/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2022-03-misc/setting-private-method-via-update/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2022-03-misc/setting-private-method/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-01-misc/setting-private-method-via-array-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-01-misc/setting-private-method-via-for-of/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-01-misc/setting-private-method-via-object-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-01-misc/setting-private-method-via-rest/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-01-misc/setting-private-method-via-update/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-01-misc/setting-private-method/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-in-body/input.js

@@ -0,0 +1,6 @@
+const dec = () => {};
+class Foo {
+  @dec #x() {
+    this.#x = 123;
+  }
+}

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-in-body/options.json

@@ -0,0 +1,3 @@
+{
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
+}

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-via-array-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-via-for-of/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-via-object-pattern/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-via-rest/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method-via-update/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-private-method/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "Decorated private methods are not updatable, but \"#x\" is updated via this expression."
+  "throws": "Decorated private methods are read-only, but \"#x\" is updated via this expression."
 }

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-shadowed-private-method-valid/input.js

@@ -0,0 +1,11 @@
+const dec = () => {};
+class Foo {
+  @dec #x() {
+    class Nested {
+      static #x;
+      static set x(v) {
+        this.#x = v;
+      }
+    }
+  }
+}

packages/babel-plugin-proposal-decorators/test/fixtures/2023-05-misc/setting-shadowed-private-method-valid/output.js

@@ -0,0 +1,18 @@
+var _call_x, _initProto;
+const dec = () => {};
+class Foo {
+  static {
+    [_call_x, _initProto] = babelHelpers.applyDecs2305(this, [[dec, 2, "x", function () {
+      class Nested {
+        static #x;
+        static set x(v) {
+          this.#x = v;
+        }
+      }
+    }]], [], 0, _ => #x in _).e;
+  }
+  constructor() {
+    _initProto(this);
+  }
+  #x = _call_x;
+}