Skip to content

Do not split awaitAsyncGenerator in await yield#17362

Merged
nicolo-ribaudo merged 2 commits intobabel:mainfrom
nicolo-ribaudo:awrap-no-split
Feb 27, 2026
Merged

Do not split awaitAsyncGenerator in await yield#17362
nicolo-ribaudo merged 2 commits intobabel:mainfrom
nicolo-ribaudo:awrap-no-split

Conversation

@nicolo-ribaudo
Copy link
Copy Markdown
Member

Q                       A
Fixed Issues? Fixes #1, Fixes #2
Patch: Bug Fix?
Major: Breaking Change?
Minor: New Feature?
Tests Added + Pass? Yes
Documentation PR Link
Any Dependency Changes?
License MIT

I noticed that the screenshot in #17360 (comment) was weird, because we were storing regeneratorRuntime on a context variable.

See the diff between the first and second commit in the test.

@nicolo-ribaudo nicolo-ribaudo added the PR: Polish 💅 A type of pull request used for our changelog categories label Jun 3, 2025
@babel-bot
Copy link
Copy Markdown
Collaborator

babel-bot commented Jun 3, 2025

Build successful! You can test your changes in the REPL here: https://babeljs.io/repl/build/61030

leapManager: leap.LeapManager;
scope: Scope;
vars: t.VariableDeclarator[];
awrapNodes: Set<t.Node> | undefined;
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
awrapNodes: Set<t.Node> | undefined;
awrapNodes: Set<t.CallExpression> | undefined;

Is awrapNodes short for asyncWrapNodes?

@liuxingbaoyu
Copy link
Copy Markdown
Member

liuxingbaoyu commented Jun 7, 2025

I tried another approach to avoid an extra traversal and it seems to work. :)

diff --git a/packages/babel-plugin-transform-regenerator/src/regenerator/emit.ts b/packages/babel-plugin-transform-regenerator/src/regenerator/emit.ts
index 6907046f4b..d620616806 100644
--- a/packages/babel-plugin-transform-regenerator/src/regenerator/emit.ts
+++ b/packages/babel-plugin-transform-regenerator/src/regenerator/emit.ts
@@ -73,7 +73,6 @@ export class Emitter {
   leapManager: leap.LeapManager;
   scope: Scope;
   vars: t.VariableDeclarator[];
-  awrapNodes: Set<t.Node> | undefined;
 
   pluginPass: PluginPass;
 
@@ -81,13 +80,11 @@ export class Emitter {
     contextId: t.Identifier,
     scope: Scope,
     vars: t.VariableDeclarator[],
-    awrapNodes: Set<t.Node> | undefined,
     pluginPass: PluginPass,
   ) {
     this.pluginPass = pluginPass;
     this.scope = scope;
     this.vars = vars;
-    this.awrapNodes = awrapNodes;
 
     // Used to generate unique temporary names.
     this.nextTempId = 0;
@@ -1078,70 +1075,66 @@ export class Emitter {
 
         let injectFirstArg = null;
 
-        if (this.awrapNodes?.has(path.node)) {
-          newCallee = calleePath.node;
-        } else {
-          if (t.isMemberExpression(calleePath.node)) {
-            if (hasLeapingArgs) {
-              // If the arguments of the CallExpression contained any yield
-              // expressions, then we need to be sure to evaluate the callee
-              // before evaluating the arguments, but if the callee was a member
-              // expression, then we must be careful that the object of the
-              // member expression still gets bound to `this` for the call.
-
-              const newObject = self.explodeViaTempVar(
-                // Assign the exploded callee.object expression to a temporary
-                // variable so that we can use it twice without reevaluating it.
-                self.makeTempVar(),
-                calleePath.get("object") as NodePath<t.Expression>,
-                hasLeapingChildren,
-              );
-
-              const newProperty = calleePath.node.computed
-                ? self.explodeViaTempVar(
-                    null,
-                    calleePath.get("property") as NodePath<t.Expression>,
-                    hasLeapingChildren,
-                  )
-                : calleePath.node.property;
+        if (t.isMemberExpression(calleePath.node)) {
+          if (hasLeapingArgs) {
+            // If the arguments of the CallExpression contained any yield
+            // expressions, then we need to be sure to evaluate the callee
+            // before evaluating the arguments, but if the callee was a member
+            // expression, then we must be careful that the object of the
+            // member expression still gets bound to `this` for the call.
+
+            const newObject = self.explodeViaTempVar(
+              // Assign the exploded callee.object expression to a temporary
+              // variable so that we can use it twice without reevaluating it.
+              self.makeTempVar(),
+              calleePath.get("object") as NodePath<t.Expression>,
+              hasLeapingChildren,
+            );
 
-              injectFirstArg = newObject;
+            const newProperty = calleePath.node.computed
+              ? self.explodeViaTempVar(
+                  null,
+                  calleePath.get("property") as NodePath<t.Expression>,
+                  hasLeapingChildren,
+                )
+              : calleePath.node.property;
 
-              newCallee = t.memberExpression(
-                t.memberExpression(
-                  t.cloneNode(newObject),
-                  newProperty,
-                  calleePath.node.computed,
-                ),
-                t.identifier("call"),
-                false,
-              );
-            } else {
-              newCallee = self.explodeExpression(
-                calleePath as NodePath<t.Expression>,
-              );
-            }
+            injectFirstArg = newObject;
+
+            newCallee = t.memberExpression(
+              t.memberExpression(
+                t.cloneNode(newObject),
+                newProperty,
+                calleePath.node.computed,
+              ),
+              t.identifier("call"),
+              false,
+            );
           } else {
-            newCallee = self.explodeViaTempVar(
-              null,
+            newCallee = self.explodeExpression(
               calleePath as NodePath<t.Expression>,
-              hasLeapingChildren,
             );
+          }
+        } else {
+          newCallee = self.explodeViaTempVar(
+            null,
+            calleePath as NodePath<t.Expression>,
+            hasLeapingChildren,
+          );
 
-            if (t.isMemberExpression(newCallee)) {
-              // If the callee was not previously a MemberExpression, then the
-              // CallExpression was "unqualified," meaning its `this` object
-              // should be the global object. If the exploded expression has
-              // become a MemberExpression (e.g. a context property, probably a
-              // temporary variable), then we need to force it to be unqualified
-              // by using the (0, object.property)(...) trick; otherwise, it
-              // will receive the object of the MemberExpression as its `this`
-              // object.
-              newCallee = t.sequenceExpression([
-                t.numericLiteral(0),
-                t.cloneNode(newCallee),
-              ]);
-            }
+          if (t.isMemberExpression(newCallee)) {
+            // If the callee was not previously a MemberExpression, then the
+            // CallExpression was "unqualified," meaning its `this` object
+            // should be the global object. If the exploded expression has
+            // become a MemberExpression (e.g. a context property, probably a
+            // temporary variable), then we need to force it to be unqualified
+            // by using the (0, object.property)(...) trick; otherwise, it
+            // will receive the object of the MemberExpression as its `this`
+            // object.
+            newCallee = t.sequenceExpression([
+              t.numericLiteral(0),
+              t.cloneNode(newCallee),
+            ]);
           }
         }
 
@@ -1446,6 +1439,48 @@ export class Emitter {
             : "sent",
         );
 
+      case "AwaitExpression": {
+        after = this.loc();
+        const arg = self.explodeExpression(path.get("argument"));
+
+        self.emitAssign(
+          self.contextProperty(
+            process.env.BABEL_8_BREAKING ||
+              util.newHelpersAvailable(this.pluginPass)
+              ? "n"
+              : "next",
+          ),
+          after,
+        );
+
+        const helper =
+          // This is slightly tricky: newer versions of the `regeneratorRuntime`
+          // helper support using `awaitAsyncGenerator` as an alternative to
+          // `regeneratorRuntime().awrap`. There is no direct way to test if we
+          // have that part of the helper available, but we know that it has been
+          // introduced in the same version as `regeneratorKeys`.
+          process.env.BABEL_8_BREAKING ||
+          util.newHelpersAvailable(self.pluginPass)
+            ? self.pluginPass.addHelper("awaitAsyncGenerator")
+            : util.runtimeProperty(self.pluginPass, "awrap");
+
+        const ret = t.returnStatement(
+          t.cloneNode(t.callExpression(helper, [arg])) || null,
+        );
+        // Preserve the `yield` location so that source mappings for the statements
+        // link back to the yield properly.
+        ret.loc = expr.loc;
+        self.emit(ret);
+        self.mark(after);
+
+        return self.contextProperty(
+          process.env.BABEL_8_BREAKING ||
+            util.newHelpersAvailable(self.pluginPass)
+            ? "v"
+            : "sent",
+        );
+      }
+
       case "ClassExpression":
         return finish(self.explodeClass(path));
 
diff --git a/packages/babel-plugin-transform-regenerator/src/regenerator/meta.ts b/packages/babel-plugin-transform-regenerator/src/regenerator/meta.ts
index 8fa5325438..b3094da26f 100644
--- a/packages/babel-plugin-transform-regenerator/src/regenerator/meta.ts
+++ b/packages/babel-plugin-transform-regenerator/src/regenerator/meta.ts
@@ -85,6 +85,7 @@ const sideEffectTypes = {
 // These types are the direct cause of all leaps in control flow.
 const leapTypes = {
   YieldExpression: true,
+  AwaitExpression: true,
   BreakStatement: true,
   ContinueStatement: true,
   ReturnStatement: true,
diff --git a/packages/babel-plugin-transform-regenerator/src/regenerator/visit.ts b/packages/babel-plugin-transform-regenerator/src/regenerator/visit.ts
index d271f8f306..7eb3c04f51 100644
--- a/packages/babel-plugin-transform-regenerator/src/regenerator/visit.ts
+++ b/packages/babel-plugin-transform-regenerator/src/regenerator/visit.ts
@@ -49,14 +49,6 @@ export const getVisitor = (t: any): Visitor<PluginPass> => ({
       path.ensureBlock();
       const bodyBlockPath = path.get("body");
 
-      let awrapNodes: Set<t.Node> | undefined;
-
-      if (node.async) {
-        awrapNodes = new Set();
-        // eslint-disable-next-line @typescript-eslint/no-use-before-define
-        bodyBlockPath.traverse(awaitVisitor, { pluginPass: this, awrapNodes });
-      }
-
       // eslint-disable-next-line @typescript-eslint/no-use-before-define
       bodyBlockPath.traverse(functionSentVisitor, {
         context: contextId,
@@ -115,13 +107,7 @@ export const getVisitor = (t: any): Visitor<PluginPass> => ({
         );
       }
 
-      const emitter = new Emitter(
-        contextId,
-        path.scope,
-        vars,
-        awrapNodes,
-        this,
-      );
+      const emitter = new Emitter(contextId, path.scope, vars, this);
       emitter.explode(path.get("body"));
 
       if (vars.length > 0) {
@@ -376,37 +362,3 @@ const functionSentVisitor = {
     }
   },
 };
-
-const awaitVisitor: Visitor<{
-  pluginPass: PluginPass;
-  awrapNodes: Set<t.Node>;
-}> = {
-  Function: function (path) {
-    path.skip(); // Don't descend into nested function scopes.
-  },
-
-  AwaitExpression: function (path, { pluginPass, awrapNodes }) {
-    const t = util.getTypes();
-
-    // Convert await expressions to yield expressions.
-    const argument = path.node.argument;
-
-    const helper =
-      // This is slightly tricky: newer versions of the `regeneratorRuntime`
-      // helper support using `awaitAsyncGenerator` as an alternative to
-      // `regeneratorRuntime().awrap`. There is no direct way to test if we
-      // have that part of the helper available, but we know that it has been
-      // introduced in the same version as `regeneratorKeys`.
-      process.env.BABEL_8_BREAKING || util.newHelpersAvailable(pluginPass)
-        ? pluginPass.addHelper("awaitAsyncGenerator")
-        : util.runtimeProperty(pluginPass, "awrap");
-
-    // Transforming `await x` to `yield regeneratorRuntime.awrap(x)`
-    // causes the argument to be wrapped in such a way that the runtime
-    // can distinguish between awaited and merely yielded values.
-    const awrap = t.callExpression(helper, [argument]);
-    util.replaceWithOrRemove(path, t.yieldExpression(awrap, false));
-
-    awrapNodes.add(awrap);
-  },
-};
diff --git a/packages/babel-plugin-transform-regenerator/test/fixtures/misc/await-yield/output.js b/packages/babel-plugin-transform-regenerator/test/fixtures/misc/await-yield/output.js
index 9545190960..951ae6fed7 100644
--- a/packages/babel-plugin-transform-regenerator/test/fixtures/misc/await-yield/output.js
+++ b/packages/babel-plugin-transform-regenerator/test/fixtures/misc/await-yield/output.js
@@ -1,15 +1,14 @@
 var _marked = /*#__PURE__*/babelHelpers.regenerator().m(f);
 function f() {
-  var res, _t;
+  var res;
   return babelHelpers.regeneratorAsyncGen(function (_context) {
     while (1) switch (_context.n) {
       case 0:
         _context.n = 1;
         return 1;
       case 1:
-        _t = _context.v;
         _context.n = 2;
-        return babelHelpers.awaitAsyncGenerator(_t);
+        return babelHelpers.awaitAsyncGenerator(_context.v);
       case 2:
         res = _context.v;
       case 3:

Copy link
Copy Markdown
Member

@liuxingbaoyu liuxingbaoyu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I applied #17362 (comment) and rebased it.

@pkg-pr-new
Copy link
Copy Markdown

pkg-pr-new bot commented Feb 26, 2026

Open in StackBlitz

commit: d53d2cb

@nicolo-ribaudo nicolo-ribaudo merged commit 59c8b94 into babel:main Feb 27, 2026
55 checks passed
@nicolo-ribaudo nicolo-ribaudo deleted the awrap-no-split branch February 27, 2026 16:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

PR: Polish 💅 A type of pull request used for our changelog categories Spec: Async Generators

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants