From 84ee581980c2a215ddd5247b2e20d87b3cf148bf Mon Sep 17 00:00:00 2001
From: overlookmotel <557937+overlookmotel@users.noreply.github.com>
Date: Wed, 6 Nov 2024 16:06:51 +0000
Subject: [PATCH] refactor(transformer/async-generator-functions): simplify
 identifying whether within an async generator function (#7170)

Follow-on after stack up to #7148.

Split `is_inside_async_generator_function` into 2 functions `yield_is_inside_async_generator_function` and `async_is_inside_async_generator_function`.

There are a couple of differences between `yield` and `await`, which we can leverage to make both these functions simpler:

* `yield` cannot appear at top level (there's no such thing as "top level yield").
* `yield` cannot appear in an arrow function (there's no such thing as a "generator arrow function").

In addition:

* Because each function now handles a specific case, no need to check if function containing `async` is an async function. It must be. Ditto for `yield`.
* Remove the `if ctx.current_scope_flags().is_top()` check. Probably this check costs more than it saves because top level `await` is far less common than `await` within async functions. So this check was optimizing for an uncommon case, at a cost to the common case.
* Add more comments to explain the logic.

This is all quite small optimizations, but I was having difficulty understanding the logic, and found that splitting it up made it clearer (at least for me).
---
 .../es2018/async_generator_functions/mod.rs   | 55 ++++++++++++-------
 1 file changed, 34 insertions(+), 21 deletions(-)

diff --git a/crates/oxc_transformer/src/es2018/async_generator_functions/mod.rs b/crates/oxc_transformer/src/es2018/async_generator_functions/mod.rs
index d36d4763e2a32..6b2f03a237880 100644
--- a/crates/oxc_transformer/src/es2018/async_generator_functions/mod.rs
+++ b/crates/oxc_transformer/src/es2018/async_generator_functions/mod.rs
@@ -157,32 +157,13 @@ impl<'a, 'ctx> Traverse<'a> for AsyncGeneratorFunctions<'a, 'ctx> {
 }
 
 impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
-    /// Check whether the current node is inside an async generator function.
-    fn is_inside_async_generator_function(ctx: &mut TraverseCtx<'a>) -> bool {
-        // Early return if current scope is top because we don't need to transform top-level await expression.
-        if ctx.current_scope_flags().is_top() {
-            return false;
-        }
-
-        for ancestor in ctx.ancestors() {
-            match ancestor {
-                Ancestor::FunctionBody(func) => return *func.r#async() && *func.generator(),
-                Ancestor::ArrowFunctionExpressionBody(_) => {
-                    return false;
-                }
-                _ => {}
-            }
-        }
-        false
-    }
-
     /// Transform `yield * argument` expression to `yield asyncGeneratorDelegate(asyncIterator(argument))`.
     fn transform_yield_expression(
         &self,
         expr: &mut YieldExpression<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) -> Option<Expression<'a>> {
-        if !expr.delegate || !Self::is_inside_async_generator_function(ctx) {
+        if !expr.delegate || !Self::yield_is_inside_async_generator_function(ctx) {
             return None;
         }
 
@@ -196,6 +177,19 @@ impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
         })
     }
 
+    /// Check whether `yield` is inside an async generator function.
+    fn yield_is_inside_async_generator_function(ctx: &TraverseCtx<'a>) -> bool {
+        for ancestor in ctx.ancestors() {
+            // Note: `yield` cannot appear within function params.
+            // Also cannot appear in arrow functions because no such thing as a generator arrow function.
+            if let Ancestor::FunctionBody(func) = ancestor {
+                return *func.r#async();
+            }
+        }
+        // `yield` can only appear in a function
+        unreachable!();
+    }
+
     /// Transforms `await expr` expression to `yield awaitAsyncGenerator(expr)`.
     /// Ignores top-level await expression.
     fn transform_await_expression(
@@ -203,7 +197,7 @@ impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
         expr: &mut AwaitExpression<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) -> Option<Expression<'a>> {
-        if !Self::is_inside_async_generator_function(ctx) {
+        if !Self::async_is_inside_async_generator_function(ctx) {
             return None;
         }
 
@@ -213,4 +207,23 @@ impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
 
         Some(ctx.ast.expression_yield(SPAN, false, Some(argument)))
     }
+
+    /// Check whether `await` is inside an async generator function.
+    fn async_is_inside_async_generator_function(ctx: &TraverseCtx<'a>) -> bool {
+        for ancestor in ctx.ancestors() {
+            match ancestor {
+                // Note: `await` cannot appear within function params
+                Ancestor::FunctionBody(func) => {
+                    return *func.generator();
+                }
+                Ancestor::ArrowFunctionExpressionBody(_) => {
+                    // Arrow functions can't be generator functions
+                    return false;
+                }
+                _ => {}
+            }
+        }
+        // Top level await
+        false
+    }
 }