Improve output when wrapping functions (e.g. async functions) (#15922)

Co-authored-by: Nicolò Ribaudo <nicolo.ribaudo@gmail.com>
babel · Sep 7, 2023 · b9a2244 · b9a2244
1 parent 1fb3b7a
commit b9a2244
Show file tree

Hide file tree

Showing 76 changed files with 396 additions and 700 deletions.
diff --git a/packages/babel-helper-remap-async-to-generator/src/index.ts b/packages/babel-helper-remap-async-to-generator/src/index.ts
@@ -1,5 +1,3 @@
-/* @noflow */
-
 import type { NodePath } from "@babel/traverse";
 import wrapFunction from "@babel/helper-wrap-function";
 import annotateAsPure from "@babel/helper-annotate-as-pure";
@@ -65,7 +63,7 @@ export default function (
     path.parentPath.isObjectProperty() ||
     path.parentPath.isClassProperty();
 
-  if (!isProperty && !isIIFE && path.isExpression()) {
+  if (!isProperty && !isIIFE && path.isCallExpression()) {
     annotateAsPure(path);
   }
 

diff --git a/packages/babel-helper-wrap-function/src/index.ts b/packages/babel-helper-wrap-function/src/index.ts
@@ -10,56 +10,18 @@ import {
   isRestElement,
   returnStatement,
   isCallExpression,
+  cloneNode,
+  toExpression,
 } from "@babel/types";
 import type * as t from "@babel/types";
 
-type ExpressionWrapperBuilder<ExtraBody extends t.Node[]> = (
-  replacements?: Parameters<ReturnType<typeof template.expression>>[0],
-) => t.CallExpression & {
-  callee: t.FunctionExpression & {
-    body: {
-      body: [
-        t.VariableDeclaration & {
-          declarations: [
-            { init: t.FunctionExpression | t.ArrowFunctionExpression },
-          ];
-        },
-        ...ExtraBody,
-      ];
-    };
-  };
-};
-
-const buildAnonymousExpressionWrapper = template.expression(`
-  (function () {
-    var REF = FUNCTION;
-    return function NAME(PARAMS) {
-      return REF.apply(this, arguments);
-    };
-  })()
-`) as ExpressionWrapperBuilder<
-  [t.ReturnStatement & { argument: t.FunctionExpression }]
->;
-
-const buildNamedExpressionWrapper = template.expression(`
-  (function () {
-    var REF = FUNCTION;
-    function NAME(PARAMS) {
-      return REF.apply(this, arguments);
-    }
-    return NAME;
-  })()
-`) as ExpressionWrapperBuilder<
-  [t.FunctionDeclaration, t.ReturnStatement & { argument: t.Identifier }]
->;
-
-const buildDeclarationWrapper = template.statements(`
-  function NAME(PARAMS) { return REF.apply(this, arguments); }
-  function REF() {
-    REF = FUNCTION;
-    return REF.apply(this, arguments);
+const buildWrapper = template.statement(`
+  function NAME(PARAMS) {
+    return (REF = REF || FUNCTION).apply(this, arguments);
   }
-`);
+`) as (
+  replacements: Parameters<ReturnType<typeof template.expression>>[0],
+) => t.FunctionDeclaration;
 
 function classOrObjectMethod(
   path: NodePath<t.ClassMethod | t.ClassPrivateMethod | t.ObjectMethod>,
@@ -140,40 +102,32 @@ function plainFunction(
     params.push(path.scope.generateUidIdentifier("x"));
   }
 
-  const wrapperArgs = {
-    NAME: functionId || null,
-    REF: path.scope.generateUidIdentifier(functionId ? functionId.name : "ref"),
+  const ref = path.scope.generateUidIdentifier(
+    functionId ? functionId.name : "ref",
+  );
+
+  let wrapper: t.Function = buildWrapper({
+    NAME: functionId,
+    REF: ref,
     FUNCTION: built,
     PARAMS: params,
-  };
+  });
+
+  if (!isDeclaration) {
+    wrapper = toExpression(wrapper);
+    nameFunction({
+      node: wrapper,
+      parent: (path as NodePath<t.FunctionExpression>).parent,
+      scope: path.scope,
+    });
+  }
 
-  if (isDeclaration) {
-    const container = buildDeclarationWrapper(wrapperArgs);
-    path.replaceWith(container[0]);
-    path.insertAfter(container[1]);
+  if (isDeclaration || wrapper.id || (!ignoreFunctionLength && params.length)) {
+    path.replaceWith(wrapper);
+    path.parentPath.scope.push({ id: cloneNode(ref) });
   } else {
-    let container;
-
-    if (functionId) {
-      container = buildNamedExpressionWrapper(wrapperArgs);
-    } else {
-      container = buildAnonymousExpressionWrapper(wrapperArgs);
-
-      const returnFn = container.callee.body.body[1].argument;
-      nameFunction({
-        node: returnFn,
-        parent: (path as NodePath<t.FunctionExpression>).parent,
-        scope: path.scope,
-      });
-      functionId = returnFn.id;
-    }
-
-    if (functionId || (!ignoreFunctionLength && params.length)) {
-      path.replaceWith(container);
-    } else {
-      // we can omit this wrapper as the conditions it protects for do not apply
-      path.replaceWith(built);
-    }
+    // we can omit this wrapper as the conditions it protects for do not apply
+    path.replaceWith(built);
   }
 }
 

diff --git a/...posal-explicit-resource-management/test/fixtures/integration/async-to-generator/output.js b/...posal-explicit-resource-management/test/fixtures/integration/async-to-generator/output.js
@@ -1,8 +1,6 @@
+var _fn;
 function fn() {
-  return _fn.apply(this, arguments);
-}
-function _fn() {
-  _fn = babelHelpers.asyncToGenerator(function* () {
+  return (_fn = _fn || babelHelpers.asyncToGenerator(function* () {
     yield 0;
     try {
       var _stack = [];
@@ -14,6 +12,5 @@ function _fn() {
     } finally {
       yield babelHelpers.dispose(_stack, _error, _hasError);
     }
-  });
-  return _fn.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/packages/babel-plugin-proposal-function-sent/test/fixtures/function-sent/basic/output.js b/packages/babel-plugin-proposal-function-sent/test/fixtures/function-sent/basic/output.js
@@ -1,10 +1,7 @@
+var _gen;
 function gen() {
-  return _gen.apply(this, arguments);
-}
-function _gen() {
-  _gen = babelHelpers.skipFirstGeneratorNext(function* () {
+  return (_gen = _gen || babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     let sent = _functionSent;
-  });
-  return _gen.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...bel-plugin-proposal-function-sent/test/fixtures/generator-kinds/async-generator/output.js b/...bel-plugin-proposal-function-sent/test/fixtures/generator-kinds/async-generator/output.js
@@ -1,10 +1,7 @@
+var _foo;
 function foo() {
-  return _foo.apply(this, arguments);
-}
-function _foo() {
-  _foo = babelHelpers.wrapAsyncGenerator(babelHelpers.skipFirstGeneratorNext(function* () {
+  return (_foo = _foo || babelHelpers.wrapAsyncGenerator(babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     _functionSent = yield babelHelpers.awaitAsyncGenerator(_functionSent);
-  }));
-  return _foo.apply(this, arguments);
+  }))).apply(this, arguments);
 }
diff --git a/...-proposal-function-sent/test/fixtures/generator-kinds/export-default-anonymous/output.mjs b/...-proposal-function-sent/test/fixtures/generator-kinds/export-default-anonymous/output.mjs
@@ -1,10 +1,7 @@
+var _ref;
 export default function () {
-  return _ref.apply(this, arguments);
-}
-function _ref() {
-  _ref = babelHelpers.skipFirstGeneratorNext(function* () {
+  return (_ref = _ref || babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     return _functionSent;
-  });
-  return _ref.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...ugin-proposal-function-sent/test/fixtures/generator-kinds/export-default-named/output.mjs b/...ugin-proposal-function-sent/test/fixtures/generator-kinds/export-default-named/output.mjs
@@ -1,10 +1,7 @@
+var _gen;
 export default function gen() {
-  return _gen.apply(this, arguments);
-}
-function _gen() {
-  _gen = babelHelpers.skipFirstGeneratorNext(function* () {
+  return (_gen = _gen || babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     return _functionSent;
-  });
-  return _gen.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/packages/babel-plugin-proposal-function-sent/test/fixtures/generator-kinds/export/output.mjs b/packages/babel-plugin-proposal-function-sent/test/fixtures/generator-kinds/export/output.mjs
@@ -1,10 +1,7 @@
+var _gen;
 export function gen() {
-  return _gen.apply(this, arguments);
-}
-function _gen() {
-  _gen = babelHelpers.skipFirstGeneratorNext(function* () {
+  return (_gen = _gen || babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     return _functionSent;
-  });
-  return _gen.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...el-plugin-proposal-function-sent/test/fixtures/generator-kinds/expression-named/output.js b/...el-plugin-proposal-function-sent/test/fixtures/generator-kinds/expression-named/output.js
@@ -1,10 +1,7 @@
-const foo = function () {
-  var _gen = babelHelpers.skipFirstGeneratorNext(function* () {
+var _gen;
+const foo = function gen() {
+  return (_gen = _gen || babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     return _functionSent;
-  });
-  function gen() {
-    return _gen.apply(this, arguments);
-  }
-  return gen;
-}();
+  })).apply(this, arguments);
+};
diff --git a/...ges/babel-plugin-proposal-function-sent/test/fixtures/generator-kinds/statement/output.js b/...ges/babel-plugin-proposal-function-sent/test/fixtures/generator-kinds/statement/output.js
@@ -1,10 +1,7 @@
+var _gen;
 function gen() {
-  return _gen.apply(this, arguments);
-}
-function _gen() {
-  _gen = babelHelpers.skipFirstGeneratorNext(function* () {
+  return (_gen = _gen || babelHelpers.skipFirstGeneratorNext(function* () {
     let _functionSent = yield;
     return _functionSent;
-  });
-  return _gen.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...-transform-async-generator-functions/test/fixtures/async-generators/declaration/output.js b/...-transform-async-generator-functions/test/fixtures/async-generators/declaration/output.js
@@ -1,12 +1,9 @@
+var _agf;
 function agf() {
-  return _agf.apply(this, arguments);
-}
-function _agf() {
-  _agf = babelHelpers.wrapAsyncGenerator(function* () {
+  return (_agf = _agf || babelHelpers.wrapAsyncGenerator(function* () {
     this;
     yield babelHelpers.awaitAsyncGenerator(1);
     yield 2;
     return 3;
-  });
-  return _agf.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...n-transform-async-generator-functions/test/fixtures/async-generators/expression/output.js b/...n-transform-async-generator-functions/test/fixtures/async-generators/expression/output.js
@@ -1,12 +1,9 @@
-/*#__PURE__*/(function () {
-  var _agf = babelHelpers.wrapAsyncGenerator(function* () {
+var _agf;
+(function agf() {
+  return (_agf = _agf || babelHelpers.wrapAsyncGenerator(function* () {
     this;
     yield babelHelpers.awaitAsyncGenerator(1);
     yield 2;
     return 3;
-  });
-  function agf() {
-    return _agf.apply(this, arguments);
-  }
-  return agf;
-})();
+  })).apply(this, arguments);
+});
diff --git a/...m-async-generator-functions/test/fixtures/async-generators/transform-class-keys/output.js b/...m-async-generator-functions/test/fixtures/async-generators/transform-class-keys/output.js
@@ -1,14 +1,11 @@
+var _fn;
 function fn() {
-  return _fn.apply(this, arguments);
-}
-function _fn() {
-  _fn = babelHelpers.wrapAsyncGenerator(function* () {
+  return (_fn = _fn || babelHelpers.wrapAsyncGenerator(function* () {
     class A {
       [yield 1]() {}
     }
     class B extends A {
       [yield babelHelpers.awaitAsyncGenerator(1)]() {}
     }
-  });
-  return _fn.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...orm-async-generator-functions/test/fixtures/async-generators/yield-star-babel-7/output.js b/...orm-async-generator-functions/test/fixtures/async-generators/yield-star-babel-7/output.js
@@ -1,10 +1,7 @@
+var _g;
 function g() {
-  return _g.apply(this, arguments);
-}
-function _g() {
-  _g = babelHelpers.wrapAsyncGenerator(function* () {
+  return (_g = _g || babelHelpers.wrapAsyncGenerator(function* () {
     yield* babelHelpers.asyncGeneratorDelegate(babelHelpers.asyncIterator([1, 2, 3]), babelHelpers.awaitAsyncGenerator);
     yield* babelHelpers.asyncGeneratorDelegate(babelHelpers.asyncIterator(iterable), babelHelpers.awaitAsyncGenerator);
-  });
-  return _g.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...n-transform-async-generator-functions/test/fixtures/async-generators/yield-star/output.js b/...n-transform-async-generator-functions/test/fixtures/async-generators/yield-star/output.js
@@ -1,10 +1,7 @@
+var _g;
 function g() {
-  return _g.apply(this, arguments);
-}
-function _g() {
-  _g = babelHelpers.wrapAsyncGenerator(function* () {
+  return (_g = _g || babelHelpers.wrapAsyncGenerator(function* () {
     yield* babelHelpers.asyncGeneratorDelegate(babelHelpers.asyncIterator([1, 2, 3]));
     yield* babelHelpers.asyncGeneratorDelegate(babelHelpers.asyncIterator(iterable));
-  });
-  return _g.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...ugin-transform-async-generator-functions/test/fixtures/for-await/async-function/output.js b/...ugin-transform-async-generator-functions/test/fixtures/for-await/async-function/output.js
@@ -1,8 +1,6 @@
+var _f;
 function f() {
-  return _f.apply(this, arguments);
-}
-function _f() {
-  _f = babelHelpers.asyncToGenerator(function* () {
+  return (_f = _f || babelHelpers.asyncToGenerator(function* () {
     var _iteratorAbruptCompletion = false;
     var _didIteratorError = false;
     var _iteratorError;
@@ -27,6 +25,5 @@ function _f() {
         }
       }
     }
-  });
-  return _f.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...gin-transform-async-generator-functions/test/fixtures/for-await/async-generator/output.js b/...gin-transform-async-generator-functions/test/fixtures/for-await/async-generator/output.js
@@ -1,8 +1,6 @@
+var _g;
 function g() {
-  return _g.apply(this, arguments);
-}
-function _g() {
-  _g = babelHelpers.wrapAsyncGenerator(function* () {
+  return (_g = _g || babelHelpers.wrapAsyncGenerator(function* () {
     var _iteratorAbruptCompletion = false;
     var _didIteratorError = false;
     var _iteratorError;
@@ -27,6 +25,5 @@ function _g() {
         }
       }
     }
-  });
-  return _g.apply(this, arguments);
+  })).apply(this, arguments);
 }
diff --git a/...ync-generator-functions/test/fixtures/for-await/create-async-from-sync-iterator/output.js b/...ync-generator-functions/test/fixtures/for-await/create-async-from-sync-iterator/output.js
@@ -1,8 +1,6 @@
+var _fn;
 function fn() {
-  return _fn.apply(this, arguments);
-}
-function _fn() {
-  _fn = babelHelpers.wrapAsyncGenerator(function* () {
+  return (_fn = _fn || babelHelpers.wrapAsyncGenerator(function* () {
     var _iteratorAbruptCompletion = false;
     var _didIteratorError = false;
     var _iteratorError;
@@ -29,6 +27,5 @@ function _fn() {
         }
       }
     }
-  });
-  return _fn.apply(this, arguments);
+  })).apply(this, arguments);
 }