fix: Cannot polyfilled on optional chains (#221)

Co-authored-by: Nicolò Ribaudo <hello@nicr.dev>

packages/babel-helper-define-polyfill-provider/src/visitors/usage.ts

@@ -30,7 +30,9 @@ export default (callProvider: CallProvider) => {
     callProvider({ kind: "global", name }, path);
   }
 
-  function analyzeMemberExpression(path: NodePath<t.MemberExpression>) {
+  function analyzeMemberExpression(
+    path: NodePath<t.MemberExpression | t.OptionalMemberExpression>,
+  ) {
     const key = resolveKey(path.get("property"), path.node.computed);
     return { key, handleAsMemberExpression: !!key && key !== "prototype" };
   }
@@ -48,7 +50,9 @@ export default (callProvider: CallProvider) => {
       handleReferencedIdentifier(path);
     },
 
-    MemberExpression(path: NodePath<t.MemberExpression>) {
+    "MemberExpression|OptionalMemberExpression"(
+      path: NodePath<t.MemberExpression | t.OptionalMemberExpression>,
+    ) {
       const { key, handleAsMemberExpression } = analyzeMemberExpression(path);
       if (!handleAsMemberExpression) return;
 

packages/babel-helper-define-polyfill-provider/test/descriptors.js

@@ -152,4 +152,17 @@ describe("descriptors", () => {
     expect(path.type).toBe("BinaryExpression");
     expect(path.toString()).toBe("'values' in Object");
   });
+
+  it("optional chains", () => {
+    const [desc, path] = getDescriptor("a?.includes();", "property");
+
+    expect(desc).toEqual({
+      kind: "property",
+      object: "a",
+      key: "includes",
+      placement: "static",
+    });
+    expect(path.type).toBe("OptionalMemberExpression");
+    expect(path.toString()).toBe("a?.includes");
+  });
 });

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";
@@ -111,9 +113,9 @@ export default defineProvider<Options>(function (
 
   function maybeInjectPure(
     desc: CoreJSPolyfillDescriptor,
-    hint,
-    utils,
-    object?,
+    hint: string,
+    utils: ReturnType<typeof getUtils>,
+    object?: string,
   ) {
     if (
       desc.pure &&
@@ -247,7 +249,9 @@ export default defineProvider<Options>(function (
 
       if (meta.kind === "property") {
         // We can't compile destructuring and updateExpression.
-        if (!path.isMemberExpression()) return;
+        if (!path.isMemberExpression() && !path.isOptionalMemberExpression()) {
+          return;
+        }
         if (!path.isReferenced()) return;
         if (path.parentPath.isUpdateExpression()) return;
         if (t.isSuper(path.node.object)) {
@@ -326,7 +330,31 @@ export default defineProvider<Options>(function (
           // @ts-expect-error
           meta.object,
         );
-        if (id) path.replaceWith(id);
+        if (id) {
+          path.replaceWith(id);
+          let { parentPath } = path;
+          if (
+            parentPath.isOptionalMemberExpression() ||
+            parentPath.isOptionalCallExpression()
+          ) {
+            do {
+              const parentAsNotOptional = parentPath as NodePath as NodePath<
+                t.MemberExpression | t.CallExpression
+              >;
+              parentAsNotOptional.type = parentAsNotOptional.node.type =
+                parentPath.type === "OptionalMemberExpression"
+                  ? "MemberExpression"
+                  : "CallExpression";
+              delete parentAsNotOptional.node.optional;
+
+              ({ parentPath } = parentPath);
+            } while (
+              (parentPath.isOptionalMemberExpression() ||
+                parentPath.isOptionalCallExpression()) &&
+              !parentPath.node.optional
+            );
+          }
+        }
       } else if (resolved.kind === "instance") {
         const id = maybeInjectPure(
           resolved.desc,
@@ -337,9 +365,40 @@ export default defineProvider<Options>(function (
         );
         if (!id) return;
 
-        const { node } = path as NodePath<t.MemberExpression>;
-        if (t.isCallExpression(path.parent, { callee: node })) {
-          callMethod(path, id);
+        const { node, parent } = path as NodePath<
+          t.MemberExpression | t.OptionalMemberExpression
+        >;
+
+        if (t.isOptionalCallExpression(parent) && parent.callee === node) {
+          const wasOptional = parent.optional;
+          parent.optional = !wasOptional;
+
+          if (!wasOptional) {
+            const check = extractOptionalCheck(
+              path.scope,
+              node as t.OptionalMemberExpression,
+            );
+            const [thisArg, thisArg2] = maybeMemoizeContext(node, path.scope);
+
+            path.replaceWith(
+              check(
+                template.expression.ast`
+                  Function.call.bind(${id}(${thisArg}), ${thisArg2})
+                `,
+              ),
+            );
+          } else if (t.isOptionalMemberExpression(node)) {
+            const check = extractOptionalCheck(path.scope, node);
+            callMethod(path, id, true, check);
+          } else {
+            callMethod(path, id, true);
+          }
+        } else if (t.isCallExpression(parent) && parent.callee === node) {
+          callMethod(path, id, false);
+        } else if (t.isOptionalMemberExpression(node)) {
+          const check = extractOptionalCheck(path.scope, node);
+          path.replaceWith(check(t.callExpression(id, [node.object])));
+          if (t.isOptionalMemberExpression(parent)) parent.optional = true;
         } else {
           path.replaceWith(t.callExpression(id, [node.object]));
         }

packages/babel-plugin-polyfill-corejs3/src/utils.ts

@@ -1,25 +1,71 @@
-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";
 
-export function callMethod(path: any, id: t.Identifier) {
-  const { object } = path.node;
+export function callMethod(
+  path: any,
+  id: t.Identifier,
+  optionalCall?: boolean,
+  wrapCallee?: (callee: t.Expression) => t.Expression,
+) {
+  const [context1, context2] = maybeMemoizeContext(path.node, path.scope);
+
+  let callee: t.Expression = t.callExpression(id, [context1]);
+  if (wrapCallee) callee = wrapCallee(callee);
+
+  const call = t.identifier("call");
+
+  path.replaceWith(
+    optionalCall
+      ? t.optionalMemberExpression(callee, call, false, true)
+      : t.memberExpression(callee, call),
+  );
+
+  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)) {
-    context1 = object;
-    context2 = t.cloneNode(object);
+    context2 = object;
+    context1 = t.cloneNode(object);
   } else {
-    context1 = path.scope.generateDeclaredUidIdentifier("context");
-    context2 = t.assignmentExpression("=", t.cloneNode(context1), object);
+    context2 = scope.generateDeclaredUidIdentifier("context");
+    context1 = t.assignmentExpression("=", t.cloneNode(context2), object);
   }
 
-  path.replaceWith(
-    t.memberExpression(t.callExpression(id, [context2]), t.identifier("call")),
-  );
+  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);
 
-  path.parentPath.unshiftContainer("arguments", context1);
+  return ifNotNullish =>
+    t.conditionalExpression(
+      t.binaryExpression("==", assign, t.nullLiteral()),
+      t.unaryExpression("void", t.numericLiteral(0)),
+      ifNotNullish,
+    );
 }
 
 export function isCoreJSSource(source: string) {

packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-global/optional-chaining/input.mjs

@@ -0,0 +1,2 @@
+Array?.from();
+c?.toSorted();

packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-global/optional-chaining/options.json

@@ -0,0 +1,10 @@
+{
+  "plugins": [
+    [
+      "@@/polyfill-corejs3",
+      {
+        "method": "usage-global"
+      }
+    ]
+  ]
+}

packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-global/optional-chaining/output.mjs

@@ -0,0 +1,5 @@
+import "core-js/modules/es.array.from.js";
+import "core-js/modules/es.array.sort.js";
+import "core-js/modules/es.string.iterator.js";
+Array?.from();
+c?.toSorted();

packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/input.mjs

@@ -0,0 +1,35 @@
+a.toSorted(x);
+a.toSorted?.(x);
+a.b.toSorted?.(x);
+a.b.toSorted?.(x).c;
+a.b.toSorted?.(x)?.c;
+
+a?.toSorted(x);
+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;
+a.b?.c.toSorted?.d.e;
+
+Array.from(x);
+Array.from?.(x);
+Array.from?.(x).c;
+Array.from?.(x)?.c;
+
+Array?.from(x);
+Array?.from(x).c;
+Array?.from(x)?.c;
+Array?.from?.(x);
+Array?.from?.(x).c;
+Array?.from?.(x)?.c;
+
+Array?.from;
+Array?.from.x.y;
+Array.from?.x.y;
+Array?.from?.x.y;
+Array.from.x?.y;
+Array?.from.x().y?.();

packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/options.json

@@ -0,0 +1,14 @@
+{
+  "targets": {
+    "chrome": "30"
+  },
+  "plugins": [
+    [
+      "@@/polyfill-corejs3",
+      {
+        "method": "usage-pure",
+        "version": "3.38"
+      }
+    ]
+  ]
+}

packages/babel-plugin-polyfill-corejs3/test/fixtures/usage-pure/optional-chaining/output.mjs

@@ -0,0 +1,33 @@
+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);
+_toSortedInstanceProperty(a)?.call(a, x);
+_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 ? void 0 : Function.call.bind(_toSortedInstanceProperty(_context4), _context4))?.(x);
+((_context5 = a) == null ? void 0 : Function.call.bind(_toSortedInstanceProperty(_context6 = _context5.b), _context6))?.(x).c;
+((_context7 = a.b) == null ? void 0 : Function.call.bind(_toSortedInstanceProperty(_context7), _context7))?.(x).c;
+((_context8 = a) == null ? void 0 : Function.call.bind(_toSortedInstanceProperty(_context9 = _context8.b), _context9))?.(x)?.c;
+((_context10 = a.b) == null ? void 0 : Function.call.bind(_toSortedInstanceProperty(_context10), _context10))?.(x)?.c;
+((_context11 = a) == null ? void 0 : _toSortedInstanceProperty(_context12 = _context11.b))?.call(_context12, x).c;
+_toSortedInstanceProperty(a.b.c)?.d.e;
+((_context13 = a.b) == null ? void 0 : _toSortedInstanceProperty(_context13.c))?.d.e;
+((_context14 = a.b) == null ? void 0 : _toSortedInstanceProperty(_context14.c))?.d.e;
+_Array$from(x);
+_Array$from(x);
+_Array$from(x).c;
+_Array$from(x)?.c;
+_Array$from(x);
+_Array$from(x).c;
+_Array$from(x)?.c;
+_Array$from(x);
+_Array$from(x).c;
+_Array$from(x)?.c;
+_Array$from;
+_Array$from.x.y;
+_Array$from.x.y;
+_Array$from.x.y;
+_Array$from.x?.y;
+_Array$from.x().y?.();