Improve output when wrapping functions

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";
@@ -13,7 +11,9 @@ const {
  yieldExpression,
 } = t;
 
-const awaitVisitor = traverse.visitors.merge<{ wrapAwait: t.Expression }>([
+const awaitVisitor = traverse.visitors.merge<{
+ wrapAwait: t.Expression | (() => t.Expression);
+}>([
  {
  ArrowFunctionExpression(path) {
  path.skip();
@@ -25,7 +25,12 @@ const awaitVisitor = traverse.visitors.merge<{ wrapAwait: t.Expression }>([
  path.replaceWith(
  yieldExpression(
  wrapAwait
- ? callExpression(cloneNode(wrapAwait), [argument.node])
+ ? callExpression(
+ typeof wrapAwait === "function"
+ ? wrapAwait()
+ : cloneNode(wrapAwait),
+ [argument.node],
+ )
  : argument.node,
  ),
  );
@@ -37,8 +42,9 @@ const awaitVisitor = traverse.visitors.merge<{ wrapAwait: t.Expression }>([
 export default function (
  path: NodePath<t.Function>,
  helpers: {
- wrapAsync: t.Expression;
- wrapAwait?: t.Expression;
+ wrapAsync: t.Expression | (() => t.Expression);
+ wrapAwait?: t.Expression | (() => t.Expression);
+ callAsync?: () => t.Expression;
  },
  noNewArrows?: boolean,
  ignoreFunctionLength?: boolean,
@@ -54,9 +60,10 @@ export default function (
 
  wrapFunction(
  path,
- cloneNode(helpers.wrapAsync),
+ helpers.wrapAsync,
  noNewArrows,
  ignoreFunctionLength,
+ helpers.callAsync,
  );
 
  const isProperty =
@@ -65,7 +72,7 @@ export default function (
  path.parentPath.isObjectProperty() ||
  path.parentPath.isClassProperty();
 
- if (!isProperty && !isIIFE && path.isExpression()) {
+ if (!isProperty && !isIIFE && path.isCallExpression()) {
  annotateAsPure(path);
  }
 
@@ -99,3 +106,5 @@ export default function (
  return false;
  }
 }
+
+export { buildOnCallExpression } from "@babel/helper-wrap-function";

packages/babel-helper-transform-fixture-test-runner/src/index.ts

@@ -456,7 +456,7 @@ function validateFile(
  throw new Error(
  `Expected ${expectedLoc} to match transform output.\n` +
  `To autogenerate a passing version of this file, delete ` +
- ` the file and re-run the tests.\n\n` +
+ `the file and re-run the tests.\n\n` +
  `Diff:\n\n${diff(expectedCode, actualCode, { expand: false })}`,
  );
  }

packages/babel-helper-wrap-function/src/index.ts

@@ -10,8 +10,14 @@ import {
  isRestElement,
  returnStatement,
  isCallExpression,
+ cloneNode,
+ toExpression,
+ identifier,
+ assignmentExpression,
+ logicalExpression,
 } from "@babel/types";
 import type * as t from "@babel/types";
+import type { PluginPass } from "@babel/core";
 
 type ExpressionWrapperBuilder<ExtraBody extends t.Node[]> = (
  replacements?: Parameters<ReturnType<typeof template.expression>>[0],
@@ -61,6 +67,23 @@ const buildDeclarationWrapper = template.statements(`
  }
 `);
 
+const buildWrapper = template.statement(`
+ function NAME(PARAMS) {
+ return CALL;
+ }
+`) as (
+ replacements: Parameters<ReturnType<typeof template.expression>>[0],
+) => t.FunctionDeclaration;
+
+const wrappedFns = new WeakMap<t.CallExpression, t.Function>();
+
+function markCallWrapped(path: NodePath<t.Function>) {
+ wrappedFns.set(
+ (path.get("body.body.0.argument") as NodePath<t.CallExpression>).node,
+ (path.get("body.body.0.argument.callee") as NodePath<t.Function>).node,
+ );
+}
+
 function classOrObjectMethod(
  path: NodePath<t.ClassMethod | t.ClassPrivateMethod | t.ObjectMethod>,
  callId: t.Expression,
@@ -127,9 +150,7 @@ function plainFunction(
  functionId = node.id;
  node.id = null;
  node.type = "FunctionExpression";
- built = callExpression(callId, [
- node as Exclude<typeof node, t.FunctionDeclaration>,
- ]);
+ built = callExpression(callId, [node as t.FunctionExpression]);
  }
 
  const params: t.Identifier[] = [];
@@ -179,19 +200,158 @@ function plainFunction(
 
 export default function wrapFunction(
  path: NodePath<t.Function>,
- callId: t.Expression,
+ callId: t.Expression | (() => t.Expression),
  // TODO(Babel 8): Consider defaulting to false for spec compliance
  noNewArrows: boolean = true,
  ignoreFunctionLength: boolean = false,
+ callAsync?: () => t.Expression,
 ) {
+ if (callAsync) {
+ if (path.isMethod()) {
+ const node = path.node;
+ const body = node.body;
+
+ const container = functionExpression(
+ null,
+ [],
+ blockStatement(body.body),
+ true,
+ );
+ body.body = [returnStatement(callExpression(callAsync(), [container]))];
+
+ // Regardless of whether or not the wrapped function is an async method
+ // or generator the outer function should not be
+ node.async = false;
+ node.generator = false;
+
+ // Unwrap the wrapper IIFE's environment so super and this and such still work.
+ (
+ path.get("body.body.0.argument.arguments.0") as NodePath
+ ).unwrapFunctionEnvironment();
+ } else {
+ let node;
+ let functionId = null;
+ const nodeParams = path.node.params;
+
+ if (path.isArrowFunctionExpression()) {
+ let path2 = path.arrowFunctionToExpression({ noNewArrows });
+ if (!process.env.BABEL_8_BREAKING) {
+ // arrowFunctionToExpression returns undefined in @babel/traverse < 7.18.10
+ path2 ??= path as unknown as NodePath<
+ t.FunctionDeclaration | t.FunctionExpression | t.CallExpression
+ >;
+ }
+ // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
+ node = path2.node as
+ | t.FunctionDeclaration
+ | t.FunctionExpression
+ | t.CallExpression;
+ } else {
+ node = path.node;
+ }
+ const isDeclaration = path.isFunctionDeclaration();
+
+ let built = node;
+ if (!isCallExpression(node)) {
+ functionId = node.id;
+ built = callExpression(callAsync(), [
+ functionExpression(null, node.params, node.body, node.generator),
+ identifier("this"),
+ identifier("arguments"),
+ ]);
+ }
+
+ const params: t.Identifier[] = [];
+ for (const param of nodeParams) {
+ if (isAssignmentPattern(param) || isRestElement(param)) {
+ break;
+ }
+ params.push(path.scope.generateUidIdentifier("x"));
+ }
+
+ let wrapper: t.Function = buildWrapper({
+ NAME: functionId,
+ CALL: built,
+ PARAMS: params,
+ });
+
+ if (!isDeclaration) {
+ wrapper = toExpression(wrapper);
+ nameFunction({
+ node: wrapper,
+ parent: (path as NodePath<t.FunctionExpression>).parent,
+ scope: path.scope,
+ });
+ }
+
+ if (
+ isDeclaration ||
+ wrapper.id ||
+ (!ignoreFunctionLength && params.length)
+ ) {
+ path.replaceWith(wrapper);
+ markCallWrapped(path);
+ } else {
+ // we can omit this wrapper as the conditions it protects for do not apply
+ path.replaceWith(
+ callExpression((callId as () => t.Expression)(), [
+ node as t.FunctionExpression,
+ ]),
+ );
+ }
+ }
+ return;
+ }
+
  if (path.isMethod()) {
- classOrObjectMethod(path, callId);
+ classOrObjectMethod(path, callId as t.Expression);
  } else {
  plainFunction(
  path as NodePath<Exclude<t.Function, t.Method>>,
- callId,
+ callId as t.Expression,
  noNewArrows,
  ignoreFunctionLength,
  );
  }
 }
+
+export function buildOnCallExpression(helperName: string) {
+ return {
+ CallExpression: {
+ exit(path: NodePath<t.CallExpression>, state: PluginPass) {
+ if (!state.availableHelper(helperName)) {
+ return;
+ }
+ if (isCallExpression(path.parent)) {
+ const wrappedFn = wrappedFns.get(path.parent);
+
+ if (!wrappedFn || wrappedFn === path.node.callee) return;
+
+ const fnPath = path.findParent(p => p.isFunction());
+
+ if (
+ fnPath
+ .findParent(p => p.isLoop() || p.isFunction() || p.isClass())
+ ?.isLoop() !== true
+ ) {
+ const ref = path.scope.generateUidIdentifier("ref");
+ fnPath.parentPath.scope.push({
+ id: ref,
+ });
+ const oldNode = path.node;
+ const comments = path.node.leadingComments;
+ if (comments) path.node.leadingComments = undefined;
+ path.replaceWith(
+ assignmentExpression(
+ "=",
+ cloneNode(ref),
+ logicalExpression("||", cloneNode(ref), oldNode),
+ ),
+ );
+ if (comments) oldNode.leadingComments = comments;
+ }
+ }
+ },
+ },
+ };
+}

packages/babel-helpers/src/helpers-generated.ts

@@ -76,16 +76,26 @@ const helpers: Record<string, Helper> = {
  "7.15.9",
  'export default function _asyncIterator(r){var n,t,o,e=2;for("undefined"!=typeof Symbol&&(t=Symbol.asyncIterator,o=Symbol.iterator);e--;){if(t&&null!=(n=r[t]))return n.call(r);if(o&&null!=(n=r[o]))return new AsyncFromSyncIterator(n.call(r));t="@@asyncIterator",o="@@iterator"}throw new TypeError("Object is not async iterable")}function AsyncFromSyncIterator(r){function AsyncFromSyncIteratorContinuation(r){if(Object(r)!==r)return Promise.reject(new TypeError(r+" is not an object."));var n=r.done;return Promise.resolve(r.value).then((function(r){return{value:r,done:n}}))}return AsyncFromSyncIterator=function(r){this.s=r,this.n=r.next},AsyncFromSyncIterator.prototype={s:null,n:null,next:function(){return AsyncFromSyncIteratorContinuation(this.n.apply(this.s,arguments))},return:function(r){var n=this.s.return;return void 0===n?Promise.resolve({value:r,done:!0}):AsyncFromSyncIteratorContinuation(n.apply(this.s,arguments))},throw:function(r){var n=this.s.return;return void 0===n?Promise.reject(r):AsyncFromSyncIteratorContinuation(n.apply(this.s,arguments))}},new AsyncFromSyncIterator(r)}',
  ),
- // size: 429, gzip size: 251
+ // size: 132, gzip size: 115
  asyncToGenerator: helper(
  "7.0.0-beta.0",
- 'function asyncGeneratorStep(n,t,e,r,o,a,c){try{var i=n[a](c),u=i.value}catch(n){return void e(n)}i.done?t(u):Promise.resolve(u).then(r,o)}export default function _asyncToGenerator(n){return function(){var t=this,e=arguments;return new Promise((function(r,o){var a=n.apply(t,e);function _next(n){asyncGeneratorStep(a,r,o,_next,_throw,"next",n)}function _throw(n){asyncGeneratorStep(a,r,o,_next,_throw,"throw",n)}_next(void 0)}))}}',
+ 'import callAsync from"callAsync";export default function _asyncToGenerator(n){return function(){return callAsync(n,this,arguments)}}',
  ),
  // size: 119, gzip size: 113
  awaitAsyncGenerator: helper(
  "7.0.0-beta.0",
  'import OverloadYield from"OverloadYield";export default function _awaitAsyncGenerator(e){return new OverloadYield(e,0)}',
  ),
+ // size: 271, gzip size: 200
+ callAsync: helper(
+ "7.24.4",
+ "export default function _callAsync(t,e,n){return new Promise((function(r,i){function step(t,e){try{var n=t?o.next(e):o.throw(e),c=n.value}catch(t){return void i(t)}n.done?r(c):Promise.resolve(c).then(s,a)}var o=t.apply(e,n),s=step.bind(this,1),a=step.bind(this,0);s()}))}",
+ ),
+ // size: 96, gzip size: 107
+ callSkipFirstGeneratorNext: helper(
+ "7.20.5",
+ "export default function _callSkipFirstGeneratorNext(t,e,r){var a=t.apply(e,r);return a.next(),a}",
+ ),
  // size: 366, gzip size: 187
  callSuper: helper(
  "7.23.8",
@@ -298,6 +308,11 @@ const helpers: Record<string, Helper> = {
  "7.0.0-beta.0",
  'export default function _newArrowCheck(n,r){if(n!==r)throw new TypeError("Cannot instantiate an arrow function")}',
  ),
+ // size: 133, gzip size: 108
+ newAsyncGenerator: helper(
+ "7.24.4",
+ 'import AsyncGenerator from"AsyncGenerator";export default function _newAsyncGenerator(e,n,r){return new AsyncGenerator(e.apply(n,r))}',
+ ),
  // size: 204, gzip size: 171
  nonIterableRest: helper(
  "7.0.0-beta.0",

packages/babel-helpers/src/helpers/asyncToGenerator.js

@@ -1,35 +1,9 @@
 /* @minVersion 7.0.0-beta.0 */
 
-function asyncGeneratorStep(gen, resolve, reject, _next, _throw, key, arg) {
- try {
- var info = gen[key](arg);
- var value = info.value;
- } catch (error) {
- reject(error);
- return;
- }
-
- if (info.done) {
- resolve(value);
- } else {
- Promise.resolve(value).then(_next, _throw);
- }
-}
+import callAsync from "callAsync";
 
 export default function _asyncToGenerator(fn) {
  return function () {
- var self = this,
- args = arguments;
- return new Promise(function (resolve, reject) {
- var gen = fn.apply(self, args);
- function _next(value) {
- asyncGeneratorStep(gen, resolve, reject, _next, _throw, "next", value);
- }
- function _throw(err) {
- asyncGeneratorStep(gen, resolve, reject, _next, _throw, "throw", err);
- }
-
- _next(undefined);
- });
+ return callAsync(fn, this, arguments);
  };
 }

packages/babel-helpers/src/helpers/callAsync.js

@@ -0,0 +1,26 @@
+/* @minVersion 7.24.4 */
+
+export default function _callAsync(fn, self, args) {
+ return new Promise(function (resolve, reject) {
+ function step(happy, arg) {
+ try {
+ var info = happy ? gen.next(arg) : gen.throw(arg);
+ var value = info.value;
+ } catch (error) {
+ reject(error);
+ return;
+ }
+
+ if (info.done) {
+ resolve(value);
+ } else {
+ Promise.resolve(value).then(_next, _throw);
+ }
+ }
+ var gen = fn.apply(self, args),
+ _next = step.bind(this, 1),
+ _throw = step.bind(this, 0);
+
+ _next();
+ });
+}