feat: Combine config enums (#11)

* feat: Combine config enums

* Simplify test module matching

* reorder `InstrumentationConfig::new` args

Author: timfish

src/config.rs

@@ -3,37 +3,60 @@
  * This product includes software developed at Datadog (<https://www.datadoghq.com>/). Copyright 2025 Datadog, Inc.
  **/
 use crate::function_query::FunctionQuery;
-use nodejs_semver::{Range, Version};
+use nodejs_semver::{Range, SemverError, Version};
 use std::path::PathBuf;
 
-#[derive(Clone, Debug)]
-pub enum InstrumentationOperator {
-    Callback,
-    Promise,
-    Sync,
-    Async,
+#[derive(Debug, Clone)]
+pub struct ModuleMatcher {
+    pub name: String,
+    pub version_range: Range,
+    pub file_path: PathBuf,
 }
 
-impl InstrumentationOperator {
+impl ModuleMatcher {
+    /// Creates a new `ModuleMatcher` instance.
+    /// # Errors
+    /// Returns an error if the version range cannot be parsed.
+    pub fn new(name: &str, version_range: &str, file_path: &str) -> Result<Self, SemverError> {
+        Ok(Self {
+            name: name.to_string(),
+            version_range: Range::parse(version_range)?,
+            file_path: PathBuf::from(file_path),
+        })
+    }
+
     #[must_use]
-    pub fn as_str(&self) -> &'static str {
-        match self {
-            InstrumentationOperator::Callback => "traceCallback",
-            InstrumentationOperator::Promise => "tracePromise",
-            InstrumentationOperator::Sync => "traceSync",
-            InstrumentationOperator::Async => "traceAsync",
-        }
+    pub fn matches(&self, module_name: &str, version: &str, file_path: &PathBuf) -> bool {
+        let version: Version = match version.parse() {
+            Ok(v) => v,
+            Err(e) => {
+                println!("Failed to parse version {version}: {e}");
+                return false;
+            }
+        };
+
+        self.name == module_name
+            && version.satisfies(&self.version_range)
+            && self.file_path == *file_path
     }
 }
 
 #[derive(Debug, Clone)]
 pub struct InstrumentationConfig {
-    pub module_name: String,
-    pub version_range: Range,
-    pub file_path: PathBuf,
-    pub function_query: FunctionQuery,
-    pub operator: InstrumentationOperator,
     pub channel_name: String,
+    pub module: ModuleMatcher,
+    pub function_query: FunctionQuery,
+}
+
+impl InstrumentationConfig {
+    #[must_use]
+    pub fn new(channel_name: &str, module: ModuleMatcher, function_query: FunctionQuery) -> Self {
+        Self {
+            channel_name: channel_name.to_string(),
+            module,
+            function_query,
+        }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -44,35 +67,22 @@ pub struct Config {
 
 impl Config {
     #[must_use]
-    pub fn new(instrumentations: Vec<InstrumentationConfig>, dc_module: String) -> Self {
+    pub fn new(instrumentations: Vec<InstrumentationConfig>, dc_module: Option<String>) -> Self {
         Self {
             instrumentations,
-            dc_module,
+            dc_module: dc_module.unwrap_or_else(|| "diagnostics_channel".to_string()),
         }
     }
 
     #[must_use]
-    pub fn new_single_with_default_dc_module(instrumentation: InstrumentationConfig) -> Self {
-        Self {
-            instrumentations: vec![instrumentation],
-            dc_module: "diagnostics_channel".to_string(),
-        }
+    pub fn new_single(instrumentation: InstrumentationConfig) -> Self {
+        Self::new(vec![instrumentation], None)
     }
 }
 
 impl InstrumentationConfig {
     #[must_use]
     pub fn matches(&self, module_name: &str, version: &str, file_path: &PathBuf) -> bool {
-        let version: Version = match version.parse() {
-            Ok(v) => v,
-            Err(e) => {
-                println!("Failed to parse version {version}: {e}");
-                return false;
-            }
-        };
-
-        self.module_name == module_name
-            && version.satisfies(&self.version_range)
-            && self.file_path == *file_path
+        self.module.matches(module_name, version, file_path)
     }
 }

src/function_query.rs

@@ -5,7 +5,7 @@
 use swc_core::ecma::ast::{FnDecl, FnExpr, Function};
 
 #[derive(Debug, Clone)]
-pub enum FunctionType {
+pub(crate) enum FunctionType {
     FunctionDeclaration,
     FunctionExpression,
     Method,
@@ -15,45 +15,161 @@ pub enum FunctionType {
 pub enum FunctionKind {
     Sync,
     Async,
-    Generator,
-    AsyncGenerator,
 }
 
 impl FunctionKind {
     #[must_use]
     pub fn is_async(&self) -> bool {
-        matches!(self, FunctionKind::Async | FunctionKind::AsyncGenerator)
-    }
-
-    #[must_use]
-    pub fn is_generator(&self) -> bool {
-        matches!(self, FunctionKind::Generator | FunctionKind::AsyncGenerator)
+        matches!(self, FunctionKind::Async)
     }
 
     #[must_use]
     pub fn matches(&self, func: &Function) -> bool {
         match self {
             FunctionKind::Sync => !func.is_async && !func.is_generator,
             FunctionKind::Async => func.is_async && !func.is_generator,
-            FunctionKind::Generator => !func.is_async && func.is_generator,
-            FunctionKind::AsyncGenerator => func.is_async && func.is_generator,
+        }
+    }
+
+    #[must_use]
+    pub fn tracing_operator(&self) -> &'static str {
+        match self {
+            FunctionKind::Sync => "traceSync",
+            FunctionKind::Async => "tracePromise",
         }
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct FunctionQuery {
-    pub name: String,
-    pub class: Option<String>,
-    pub typ: FunctionType,
-    pub kind: FunctionKind,
-    pub index: usize,
+pub enum FunctionQuery {
+    ClassConstructor {
+        class_name: String,
+        index: usize,
+    },
+    ClassMethod {
+        class_name: String,
+        method_name: String,
+        kind: FunctionKind,
+        index: usize,
+    },
+    ObjectMethod {
+        method_name: String,
+        kind: FunctionKind,
+        index: usize,
+    },
+    FunctionDeclaration {
+        function_name: String,
+        kind: FunctionKind,
+        index: usize,
+    },
+    FunctionExpression {
+        expression_name: String,
+        kind: FunctionKind,
+        index: usize,
+    },
 }
 
 impl FunctionQuery {
+    #[must_use]
+    pub fn class_constructor(class_name: &str) -> Self {
+        FunctionQuery::ClassConstructor {
+            class_name: class_name.to_string(),
+            index: 0,
+        }
+    }
+
+    #[must_use]
+    pub fn class_method(class_name: &str, method_name: &str, kind: FunctionKind) -> Self {
+        FunctionQuery::ClassMethod {
+            class_name: class_name.to_string(),
+            method_name: method_name.to_string(),
+            kind,
+            index: 0,
+        }
+    }
+
+    #[must_use]
+    pub fn object_method(method_name: &str, kind: FunctionKind) -> Self {
+        FunctionQuery::ObjectMethod {
+            method_name: method_name.to_string(),
+            kind,
+            index: 0,
+        }
+    }
+
+    #[must_use]
+    pub fn function_declaration(function_name: &str, kind: FunctionKind) -> Self {
+        FunctionQuery::FunctionDeclaration {
+            function_name: function_name.to_string(),
+            kind,
+            index: 0,
+        }
+    }
+
+    #[must_use]
+    pub fn function_expression(expression_name: &str, kind: FunctionKind) -> Self {
+        FunctionQuery::FunctionExpression {
+            expression_name: expression_name.to_string(),
+            kind,
+            index: 0,
+        }
+    }
+
+    pub(crate) fn kind(&self) -> &FunctionKind {
+        match self {
+            FunctionQuery::ClassConstructor { .. } => &FunctionKind::Sync,
+            FunctionQuery::ClassMethod { kind, .. }
+            | FunctionQuery::ObjectMethod { kind, .. }
+            | FunctionQuery::FunctionDeclaration { kind, .. }
+            | FunctionQuery::FunctionExpression { kind, .. } => kind,
+        }
+    }
+
+    pub(crate) fn name(&self) -> &str {
+        match self {
+            FunctionQuery::ClassConstructor { .. } => "constructor",
+            FunctionQuery::ClassMethod { method_name, .. }
+            | FunctionQuery::ObjectMethod { method_name, .. } => method_name,
+            FunctionQuery::FunctionDeclaration { function_name, .. } => function_name,
+            FunctionQuery::FunctionExpression {
+                expression_name, ..
+            } => expression_name,
+        }
+    }
+
+    pub(crate) fn typ(&self) -> FunctionType {
+        match self {
+            FunctionQuery::ClassConstructor { .. }
+            | FunctionQuery::ClassMethod { .. }
+            | FunctionQuery::ObjectMethod { .. } => FunctionType::Method,
+            FunctionQuery::FunctionDeclaration { .. } => FunctionType::FunctionDeclaration,
+            FunctionQuery::FunctionExpression { .. } => FunctionType::FunctionExpression,
+        }
+    }
+
+    #[must_use]
+    pub(crate) fn index(&self) -> usize {
+        match self {
+            FunctionQuery::ClassConstructor { index, .. }
+            | FunctionQuery::ClassMethod { index, .. }
+            | FunctionQuery::ObjectMethod { index, .. }
+            | FunctionQuery::FunctionDeclaration { index, .. }
+            | FunctionQuery::FunctionExpression { index, .. } => *index,
+        }
+    }
+
+    #[must_use]
+    pub(crate) fn class_name(&self) -> Option<&str> {
+        match self {
+            FunctionQuery::ClassConstructor { class_name, .. }
+            | FunctionQuery::ClassMethod { class_name, .. } => Some(class_name),
+            _ => None,
+        }
+    }
+
     fn maybe_increment_count(&self, matches_except_count: bool, count: &mut usize) -> bool {
         if matches_except_count {
-            if self.index == *count {
+            if self.index() == *count {
                 true
             } else {
                 *count += 1;
@@ -65,23 +181,23 @@ impl FunctionQuery {
     }
 
     pub fn matches_decl(&self, func: &FnDecl, count: &mut usize) -> bool {
-        let matches_except_count = matches!(self.typ, FunctionType::FunctionDeclaration)
-            && self.kind.matches(&func.function)
-            && func.ident.sym == self.name;
+        let matches_except_count = matches!(self.typ(), FunctionType::FunctionDeclaration)
+            && self.kind().matches(&func.function)
+            && func.ident.sym == self.name();
         self.maybe_increment_count(matches_except_count, count)
     }
 
     pub fn matches_expr(&self, func: &FnExpr, count: &mut usize, name: &str) -> bool {
-        let matches_except_count = matches!(self.typ, FunctionType::FunctionExpression)
-            && self.kind.matches(&func.function)
-            && name == self.name;
+        let matches_except_count = matches!(self.typ(), FunctionType::FunctionExpression)
+            && self.kind().matches(&func.function)
+            && name == self.name();
         self.maybe_increment_count(matches_except_count, count)
     }
 
     pub fn matches_method(&self, func: &Function, count: &mut usize, name: &str) -> bool {
-        let matches_except_count = matches!(self.typ, FunctionType::Method)
-            && self.kind.matches(func)
-            && name == self.name;
+        let matches_except_count = matches!(self.typ(), FunctionType::Method)
+            && self.kind().matches(func)
+            && name == self.name();
         self.maybe_increment_count(matches_except_count, count)
     }
 }

src/instrumentation.rs

@@ -53,8 +53,8 @@ impl Instrumentation {
         ArrowExpr {
             params: vec![],
             body: Box::new(body.into()),
-            is_async: self.config.function_query.kind.is_async(),
-            is_generator: self.config.function_query.kind.is_generator(),
+            is_async: self.config.function_query.kind().is_async(),
+            is_generator: false,
             type_params: None,
             return_type: None,
             span: Span::default(),
@@ -68,7 +68,7 @@ impl Instrumentation {
             span: Span::default(),
             value: format!(
                 "orchestrion:{}:{}",
-                self.config.module_name, self.config.channel_name
+                self.config.module.name, self.config.channel_name
             )
             .into(),
             raw: None,
@@ -99,7 +99,7 @@ impl Instrumentation {
         let trace_ident = ident!(format!(
             "tr_ch_apm${}.{}",
             &self.config.channel_name,
-            self.config.operator.as_str()
+            self.config.function_query.kind().tracing_operator()
         ));
 
         body.stmts = vec![
@@ -230,8 +230,7 @@ impl Instrumentation {
         self.is_correct_class = self
             .config
             .function_query
-            .class
-            .as_ref()
+            .class_name()
             .is_none_or(|class| node.ident.sym.as_ref() == class);
         true
     }
@@ -261,11 +260,11 @@ impl Instrumentation {
     }
 
     pub fn visit_mut_constructor(&mut self, node: &mut Constructor) -> bool {
-        if !self.is_correct_class || self.config.function_query.name != "constructor" {
+        if !self.is_correct_class || self.config.function_query.name() != "constructor" {
             return false;
         }
 
-        if self.count == self.config.function_query.index && node.body.is_some() {
+        if self.count == self.config.function_query.index() && node.body.is_some() {
             if let Some(body) = node.body.as_mut() {
                 self.insert_constructor_tracing(body);
             }