Fix this arg and optional call after optional member expr

babel · Sep 20, 2024 · 2d9cd27 · 2d9cd27
1 parent 4296491
commit 2d9cd27
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 67 deletions.
diff --git a/packages/babel-plugin-polyfill-corejs3/src/index.ts b/packages/babel-plugin-polyfill-corejs3/src/index.ts
@@ -15,13 +15,15 @@ import * as BabelRuntimePaths from "./babel-runtime-corejs3-paths";
 import canSkipPolyfill from "./usage-filters";
 
 import type { NodePath } from "@babel/traverse";
-import { types as t } from "@babel/core";
+import { types as t, template } from "@babel/core";
 import {
   callMethod,
   coreJSModule,
   isCoreJSSource,
   coreJSPureHelper,
   BABEL_RUNTIME,
+  extractOptionalCheck,
+  maybeMemoizeContext,
 } from "./utils";
 
 import defineProvider from "@babel/helper-define-polyfill-provider";
@@ -344,64 +346,40 @@ export default defineProvider<Options>(function (
         >;
 
         if (t.isOptionalCallExpression(parent) && parent.callee === node) {
-          if (parent.optional) {
-            parent.optional = false;
-            callMethod(path, id, true);
-          } else {
-            let optionalNode = node;
-            while (
-              !optionalNode.optional &&
-              t.isOptionalMemberExpression(optionalNode.object)
-            ) {
-              optionalNode = optionalNode.object;
-            }
-            parent.optional = true;
-            optionalNode.optional = false;
-
-            const ctx = path.scope.generateDeclaredUidIdentifier("context");
-            const assign = t.assignmentExpression(
-              "=",
-              ctx,
-              optionalNode.object,
+          const wasOptional = parent.optional;
+          parent.optional = !wasOptional;
+
+          if (!wasOptional) {
+            const check = extractOptionalCheck(
+              path.scope,
+              node as t.OptionalMemberExpression,
             );
-            optionalNode.object = t.cloneNode(ctx);
+            const [thisArg, thisArg2] = maybeMemoizeContext(node, path.scope);
 
             path.replaceWith(
-              t.conditionalExpression(
-                t.binaryExpression("==", assign, t.nullLiteral()),
-                t.nullLiteral(),
-                t.callExpression(
-                  t.memberExpression(
-                    t.memberExpression(
-                      t.identifier("Function"),
-                      t.identifier("call"),
-                    ),
-                    t.identifier("bind"),
-                  ),
-                  [t.callExpression(id, [node.object]), t.cloneNode(ctx)],
-                ),
-              ),
+              template.expression.ast`
+                ${check} ? null : Function.call.bind(${id}(${thisArg}), ${thisArg2})
+              `,
+            );
+          } else if (t.isOptionalMemberExpression(node)) {
+            const check = extractOptionalCheck(
+              path.scope,
+              node as t.OptionalMemberExpression,
             );
+            callMethod(path, id, true, c =>
+              t.conditionalExpression(check, t.nullLiteral(), c),
+            );
+          } else {
+            callMethod(path, id, true);
           }
         } else if (t.isCallExpression(parent) && parent.callee === node) {
           callMethod(path, id, false);
         } else if (t.isOptionalMemberExpression(node)) {
-          let optionalNode = node;
-          while (
-            !optionalNode.optional &&
-            t.isOptionalMemberExpression(optionalNode.object)
-          ) {
-            optionalNode = optionalNode.object;
-          }
-          optionalNode.optional = false;
-
-          const ctx = path.scope.generateDeclaredUidIdentifier("context");
-          const assign = t.assignmentExpression("=", ctx, optionalNode.object);
-          optionalNode.object = t.cloneNode(ctx);
+          const check = extractOptionalCheck(path.scope, node);
 
           path.replaceWith(
             t.conditionalExpression(
-              t.binaryExpression("==", assign, t.nullLiteral()),
+              check,
               t.nullLiteral(),
               t.callExpression(id, [node.object]),
             ),

diff --git a/packages/babel-plugin-polyfill-corejs3/src/utils.ts b/packages/babel-plugin-polyfill-corejs3/src/utils.ts
@@ -1,4 +1,4 @@
-import { types as t } from "@babel/core";
+import { types as t, type NodePath } from "@babel/core";
 import corejsEntries from "../core-js-compat/entries.js";
 
 export const BABEL_RUNTIME = "@babel/runtime-corejs3";
@@ -7,19 +7,13 @@ export function callMethod(
   path: any,
   id: t.Identifier,
   optionalCall?: boolean,
+  wrapCallee?: (callee: t.Expression) => t.Expression,
 ) {
-  const { object } = path.node;
+  const [context1, context2] = maybeMemoizeContext(path.node, path.scope);
 
-  let context1, context2;
-  if (t.isIdentifier(object)) {
-    context1 = object;
-    context2 = t.cloneNode(object);
-  } else {
-    context1 = path.scope.generateDeclaredUidIdentifier("context");
-    context2 = t.assignmentExpression("=", t.cloneNode(context1), object);
-  }
+  let callee: t.Expression = t.callExpression(id, [context1]);
+  if (wrapCallee) callee = wrapCallee(callee);
 
-  const callee = t.callExpression(id, [context2]);
   const call = t.identifier("call");
 
   path.replaceWith(
@@ -28,7 +22,45 @@ export function callMethod(
       : t.memberExpression(callee, call),
   );
 
-  path.parentPath.unshiftContainer("arguments", context1);
+  path.parentPath.unshiftContainer("arguments", context2);
+}
+
+export function maybeMemoizeContext(
+  node: t.MemberExpression | t.OptionalMemberExpression,
+  scope: NodePath["scope"],
+) {
+  const { object } = node;
+
+  let context1, context2;
+  if (t.isIdentifier(object)) {
+    context2 = object;
+    context1 = t.cloneNode(object);
+  } else {
+    context2 = scope.generateDeclaredUidIdentifier("context");
+    context1 = t.assignmentExpression("=", t.cloneNode(context2), object);
+  }
+
+  return [context1, context2];
+}
+
+export function extractOptionalCheck(
+  scope: NodePath["scope"],
+  node: t.OptionalMemberExpression,
+) {
+  let optionalNode = node;
+  while (
+    !optionalNode.optional &&
+    t.isOptionalMemberExpression(optionalNode.object)
+  ) {
+    optionalNode = optionalNode.object;
+  }
+  optionalNode.optional = false;
+
+  const ctx = scope.generateDeclaredUidIdentifier("context");
+  const assign = t.assignmentExpression("=", ctx, optionalNode.object);
+  optionalNode.object = t.cloneNode(ctx);
+
+  return t.binaryExpression("==", assign, t.nullLiteral());
 }
 
 export function isCoreJSSource(source: string) {

diff --git a/packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/input.mjs b/packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/input.mjs
@@ -9,6 +9,7 @@ a?.b.toSorted(x).c;
 a.b?.toSorted(x).c;
 a?.b.toSorted(x)?.c;
 a.b?.toSorted(x)?.c;
+a?.b.toSorted?.(x).c;
 
 a.b.c.toSorted?.d.e;
 a.b?.c.toSorted.d.e;

diff --git a/packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/output.mjs b/packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/output.mjs
@@ -1,4 +1,4 @@
-var _context, _context2, _context3, _context4, _context5, _context6, _context7, _context8, _context9, _context10;
+var _context, _context2, _context3, _context4, _context5, _context6, _context7, _context8, _context9, _context10, _context11, _context12, _context13, _context14;
 import _toSortedInstanceProperty from "core-js-pure/stable/instance/to-sorted.js";
 import _Array$from from "core-js-pure/stable/array/from.js";
 _toSortedInstanceProperty(a).call(a, x);
@@ -7,13 +7,14 @@ _toSortedInstanceProperty(_context = a.b)?.call(_context, x);
 _toSortedInstanceProperty(_context2 = a.b)?.call(_context2, x).c;
 _toSortedInstanceProperty(_context3 = a.b)?.call(_context3, x)?.c;
 ((_context4 = a) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context4), _context4))?.(x);
-((_context5 = a) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context5.b), _context5))?.(x).c;
-((_context6 = a.b) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context6), _context6))?.(x).c;
-((_context7 = a) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context7.b), _context7))?.(x)?.c;
-((_context8 = a.b) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context8), _context8))?.(x)?.c;
+((_context5 = a) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context6 = _context5.b), _context6))?.(x).c;
+((_context7 = a.b) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context7), _context7))?.(x).c;
+((_context8 = a) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context9 = _context8.b), _context9))?.(x)?.c;
+((_context10 = a.b) == null ? null : Function.call.bind(_toSortedInstanceProperty(_context10), _context10))?.(x)?.c;
+((_context11 = a) == null ? null : _toSortedInstanceProperty(_context12 = _context11.b))?.call(_context12, x).c;
 _toSortedInstanceProperty(a.b.c)?.d.e;
-((_context9 = a.b) == null ? null : _toSortedInstanceProperty(_context9.c))?.d.e;
-((_context10 = a.b) == null ? null : _toSortedInstanceProperty(_context10.c))?.d.e;
+((_context13 = a.b) == null ? null : _toSortedInstanceProperty(_context13.c))?.d.e;
+((_context14 = a.b) == null ? null : _toSortedInstanceProperty(_context14.c))?.d.e;
 _Array$from(x);
 _Array$from?.(x);
 _Array$from?.(x).c;