feat: Match more class declarations for export alias check

Author: AbhiPrasad

src/instrumentation.rs

@@ -272,7 +272,9 @@ impl Instrumentation {
             .matches_decl(node, &mut self.count)
             && node.function.body.is_some()
         {
-            if let Some(body) = node.function.body.as_mut() {
+            if node.function.is_generator {
+                self.has_injected = true;
+            } else if let Some(body) = node.function.body.as_mut() {
                 self.insert_tracing(body, &node.function.params, node.function.is_async);
             }
         }
@@ -287,6 +289,18 @@ impl Instrumentation {
                     if let Pat::Ident(name) = &decl.name {
                         traced = self.trace_expr_or_count(func_expr, &name.id.sym);
                     }
+                } else if let Some(class_expr) = init.as_mut_class() {
+                    if let Pat::Ident(name) = &decl.name {
+                        if class_expr.ident.is_none() {
+                            class_expr.ident = Some(name.id.clone());
+                        } else {
+                            self.is_correct_class = self
+                                .config
+                                .function_query
+                                .class_name()
+                                .is_none_or(|class| name.id.sym.as_ref() == class);
+                        }
+                    }
                 }
             }
         }
@@ -303,11 +317,15 @@ impl Instrumentation {
     }
 
     pub fn visit_mut_class_expr(&mut self, node: &mut ClassExpr) -> bool {
-        self.is_correct_class = self.config.function_query.class_name().is_none_or(|class| {
-            node.ident
-                .as_ref()
-                .is_some_and(|ident| ident.sym.as_ref() == class)
-        });
+        // Don't override is_correct_class if it was already set by
+        // visit_mut_var_decl for `var X = class Y { ... }` patterns.
+        if !self.is_correct_class {
+            self.is_correct_class = self.config.function_query.class_name().is_none_or(|class| {
+                node.ident
+                    .as_ref()
+                    .is_some_and(|ident| ident.sym.as_ref() == class)
+            });
+        }
         true
     }
 
@@ -328,7 +346,11 @@ impl Instrumentation {
             .matches_method(&mut self.count, name.as_ref())
             && node.function.body.is_some()
         {
-            if let Some(body) = node.function.body.as_mut() {
+            if node.function.is_generator {
+                // Skip generators — the tracing wrapper doesn't support yield*.
+                // Mark as injected so the transform doesn't report a failure.
+                self.has_injected = true;
+            } else if let Some(body) = node.function.body.as_mut() {
                 self.insert_tracing(body, &node.function.params, node.function.is_async);
             }
         }
@@ -382,7 +404,9 @@ impl Instrumentation {
             .matches_method(&mut self.count, name.as_ref())
             && node.function.body.is_some()
         {
-            if let Some(body) = node.function.body.as_mut() {
+            if node.function.is_generator {
+                self.has_injected = true;
+            } else if let Some(body) = node.function.body.as_mut() {
                 self.insert_tracing(body, &node.function.params, node.function.is_async);
             }
         }
@@ -416,6 +440,20 @@ impl Instrumentation {
                     _ => {}
                 }
             }
+        } else if let Some(class_expr) = node.right.as_mut_class() {
+            if let AssignTarget::Simple(target) = &node.left {
+                if let SimpleAssignTarget::Ident(name) = target {
+                    if class_expr.ident.is_none() {
+                        class_expr.ident = Some(name.id.clone());
+                    } else {
+                        self.is_correct_class = self
+                            .config
+                            .function_query
+                            .class_name()
+                            .is_none_or(|class| name.id.sym.as_ref() == class);
+                    }
+                }
+            }
         }
         !traced
     }

tests/const_class_export_alias_mjs/mod.mjs

@@ -0,0 +1,7 @@
+const J = class {
+  async fetch(url) {
+    return 42;
+  }
+};
+
+export { J as Undici };

tests/const_class_export_alias_mjs/mod.rs

@@ -0,0 +1,15 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn const_class_export_alias_mjs() {
+    transpile_and_test(
+        file!(),
+        true,
+        Config::new_single(InstrumentationConfig::new(
+            "Undici:fetch",
+            test_module_matcher(),
+            FunctionQuery::class_method("Undici", "fetch", FunctionKind::Async).as_export_alias(),
+        )),
+    );
+}

tests/const_class_export_alias_mjs/test.mjs

@@ -0,0 +1,12 @@
+import { Undici } from './instrumented.mjs';
+import { assert, getContext } from '../common/preamble.js';
+const context = getContext('orchestrion:undici:Undici:fetch');
+const undici = new Undici();
+const result = await undici.fetch('https://example.com');
+assert.strictEqual(result, 42);
+assert.deepStrictEqual(context, {
+  start: true,
+  end: true,
+  asyncStart: 42,
+  asyncEnd: 42
+});

tests/instrumentor_test.rs

@@ -7,6 +7,7 @@ mod common;
 mod arguments_mutation;
 mod class_expression_cjs;
 mod class_method_cjs;
+mod const_class_export_alias_mjs;
 mod constructor_cjs;
 mod constructor_mjs;
 mod decl_cjs;
@@ -18,10 +19,13 @@ mod expr_cjs;
 mod expr_mjs;
 mod index_cjs;
 mod injection_failure;
+mod let_class_export_alias_mjs;
 mod multiple_class_method_cjs;
 mod multiple_load_cjs;
 mod nested_functions;
 mod object_method_cjs;
 mod polyfill_cjs;
 mod polyfill_mjs;
 mod private_method_cjs;
+mod var_class_export_alias_mjs;
+mod var_named_class_export_alias_mjs;

tests/let_class_export_alias_mjs/mod.mjs

@@ -0,0 +1,7 @@
+let J = class {
+  async fetch(url) {
+    return 42;
+  }
+};
+
+export { J as Undici };

tests/let_class_export_alias_mjs/mod.rs

@@ -0,0 +1,15 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn let_class_export_alias_mjs() {
+    transpile_and_test(
+        file!(),
+        true,
+        Config::new_single(InstrumentationConfig::new(
+            "Undici:fetch",
+            test_module_matcher(),
+            FunctionQuery::class_method("Undici", "fetch", FunctionKind::Async).as_export_alias(),
+        )),
+    );
+}

tests/let_class_export_alias_mjs/test.mjs

@@ -0,0 +1,12 @@
+import { Undici } from './instrumented.mjs';
+import { assert, getContext } from '../common/preamble.js';
+const context = getContext('orchestrion:undici:Undici:fetch');
+const undici = new Undici();
+const result = await undici.fetch('https://example.com');
+assert.strictEqual(result, 42);
+assert.deepStrictEqual(context, {
+  start: true,
+  end: true,
+  asyncStart: 42,
+  asyncEnd: 42
+});

tests/var_class_export_alias_mjs/mod.mjs

@@ -0,0 +1,7 @@
+var J = class {
+  async fetch(url) {
+    return 42;
+  }
+};
+
+export { J as Undici };

tests/var_class_export_alias_mjs/mod.rs

@@ -0,0 +1,15 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn var_class_export_alias_mjs() {
+    transpile_and_test(
+        file!(),
+        true,
+        Config::new_single(InstrumentationConfig::new(
+            "Undici:fetch",
+            test_module_matcher(),
+            FunctionQuery::class_method("Undici", "fetch", FunctionKind::Async).as_export_alias(),
+        )),
+    );
+}