Split transformer implementation

Author: timokoessler

instrumentation-wasm/src/js_transformer/helpers/insert_single_statement_into_func.rs

@@ -0,0 +1,19 @@
+use super::parse_js_code_to_statements::parse_js_code_to_statements;
+use oxc_allocator::{Allocator, Box};
+use oxc_ast::ast::FunctionBody;
+use oxc_span::SourceType;
+
+pub fn insert_single_statement_into_func<'a>(
+    allocator: &'a Allocator,
+    body: &mut Box<'a, FunctionBody<'a>>,
+    pos: usize,
+    source_text: &'a str,
+) {
+    body.statements.insert(
+        pos,
+        parse_js_code_to_statements(&allocator, &source_text, SourceType::mjs())
+            .into_iter()
+            .next()
+            .unwrap(),
+    );
+}

instrumentation-wasm/src/js_transformer/helpers/mod.rs

@@ -1,4 +1,5 @@
 pub mod get_import_code_str;
 pub mod get_method_arg_names;
+pub mod insert_single_statement_into_func;
 pub mod parse_js_code_to_statements;
 pub mod select_sourcetype_based_on_enum;

instrumentation-wasm/src/js_transformer/mod.rs

@@ -1,3 +1,4 @@
 pub mod helpers;
 pub mod instructions;
 pub mod transformer;
+mod transformer_impl;

instrumentation-wasm/src/js_transformer/transformer.rs

@@ -1,18 +1,18 @@
-use oxc_allocator::{Allocator, Box};
-use oxc_ast::ast::{FunctionBody, MethodDefinition};
+use oxc_allocator::Allocator;
 use oxc_codegen::{Codegen, CodegenOptions};
 use oxc_parser::{ParseOptions, Parser};
 use oxc_semantic::SemanticBuilder;
 use oxc_span::SourceType;
-use oxc_traverse::{traverse_mut, Traverse, TraverseCtx};
+use oxc_traverse::traverse_mut;
 
 use super::{
     helpers::{
-        get_import_code_str::get_import_code_str, get_method_arg_names::get_method_arg_names,
+        get_import_code_str::get_import_code_str,
         parse_js_code_to_statements::parse_js_code_to_statements,
         select_sourcetype_based_on_enum::select_sourcetype_based_on_enum,
     },
     instructions::FileInstructions,
+    transformer_impl::Transformer,
 };
 
 pub fn transform_code_str(
@@ -84,90 +84,3 @@ pub fn transform_code_str(
 
     js.code
 }
-
-struct Transformer<'a> {
-    allocator: &'a Allocator,
-    file_instructions: &'a FileInstructions,
-    pkg_name: &'a str,
-    pkg_version: &'a str,
-}
-
-impl<'a> Traverse<'a> for Transformer<'a> {
-    fn enter_method_definition(
-        &mut self,
-        node: &mut MethodDefinition<'a>,
-        _ctx: &mut TraverseCtx<'a>,
-    ) {
-        if !node.key.is_identifier() || !node.value.body.is_some() || !node.key.name().is_some() {
-            return;
-        }
-
-        if !node.kind.is_method() {
-            // Ignore constructor, getters and setters for now
-            return;
-        }
-
-        let method_name = node.key.name().unwrap().to_string();
-
-        let found_instruction = self
-            .file_instructions
-            .functions
-            .iter()
-            .find(|f| f.node_type == "MethodDefinition" && f.name == method_name);
-
-        if found_instruction.is_none() {
-            return;
-        }
-
-        let instruction = found_instruction.unwrap();
-
-        // We need to collect the arg names before we make the body mutable
-        let arg_names = if instruction.modify_args {
-            get_method_arg_names(node)
-        } else {
-            Vec::new()
-        };
-
-        let body = node.value.body.as_mut().unwrap();
-
-        if instruction.modify_args && !arg_names.is_empty() {
-            // Modify the arguments by adding a statement to the beginning of the function
-            // [arg1, arg2, ...] = __instrumentModifyArgs('function_identifier', [arg1, arg2, ...]);
-
-            let arg_names_str = arg_names.join(", ");
-            let source_text: &'a str = self.allocator.alloc_str(&format!(
-                "[{}] = __instrumentModifyArgs('{}', [{}]);",
-                arg_names_str, instruction.identifier, arg_names_str
-            ));
-
-            insert_single_statement_into_func(self.allocator, body, 0, &source_text);
-        }
-
-        if instruction.inspect_args {
-            // Add a statement to the beginning of the function: __instrumentInspectArgs('function_identifier', arguments);
-            let source_text: &'a str = self.allocator.alloc_str(&format!(
-                "__instrumentInspectArgs('{}', arguments, '{}', '{}', '{}', this);",
-                instruction.identifier, self.pkg_name, self.pkg_version, instruction.name
-            ));
-
-            insert_single_statement_into_func(self.allocator, body, 0, source_text);
-        }
-
-        // Todo support return value modification
-    }
-}
-
-fn insert_single_statement_into_func<'a>(
-    allocator: &'a Allocator,
-    body: &mut Box<'a, FunctionBody<'a>>,
-    pos: usize,
-    source_text: &'a str,
-) {
-    body.statements.insert(
-        pos,
-        parse_js_code_to_statements(&allocator, &source_text, SourceType::mjs())
-            .into_iter()
-            .next()
-            .unwrap(),
-    );
-}

instrumentation-wasm/src/js_transformer/transformer_impl.rs

@@ -0,0 +1,79 @@
+use super::helpers::get_method_arg_names;
+use super::instructions::FileInstructions;
+use oxc_allocator::Allocator;
+use oxc_ast::ast::MethodDefinition;
+use oxc_traverse::{Traverse, TraverseCtx};
+
+use super::helpers::insert_single_statement_into_func::insert_single_statement_into_func;
+
+pub struct Transformer<'a> {
+    pub allocator: &'a Allocator,
+    pub file_instructions: &'a FileInstructions,
+    pub pkg_name: &'a str,
+    pub pkg_version: &'a str,
+}
+
+impl<'a> Traverse<'a> for Transformer<'a> {
+    fn enter_method_definition(
+        &mut self,
+        node: &mut MethodDefinition<'a>,
+        _ctx: &mut TraverseCtx<'a>,
+    ) {
+        if !node.key.is_identifier() || !node.value.body.is_some() || !node.key.name().is_some() {
+            return;
+        }
+
+        if !node.kind.is_method() {
+            // Ignore constructor, getters and setters for now
+            return;
+        }
+
+        let method_name = node.key.name().unwrap().to_string();
+
+        let found_instruction = self
+            .file_instructions
+            .functions
+            .iter()
+            .find(|f| f.node_type == "MethodDefinition" && f.name == method_name);
+
+        if found_instruction.is_none() {
+            return;
+        }
+
+        let instruction = found_instruction.unwrap();
+
+        // We need to collect the arg names before we make the body mutable
+        let arg_names = if instruction.modify_args {
+            get_method_arg_names::get_method_arg_names(node)
+        } else {
+            Vec::new()
+        };
+
+        let body = node.value.body.as_mut().unwrap();
+
+        if instruction.modify_args && !arg_names.is_empty() {
+            // Modify the arguments by adding a statement to the beginning of the function
+            // [arg1, arg2, ...] = __instrumentModifyArgs('function_identifier', [arg1, arg2, ...]);
+
+            let arg_names_str = arg_names.join(", ");
+            let source_text: &'a str = self.allocator.alloc_str(&format!(
+                "[{}] = __instrumentModifyArgs('{}', [{}]);",
+                arg_names_str, instruction.identifier, arg_names_str
+            ));
+
+            insert_single_statement_into_func(self.allocator, body, 0, &source_text);
+        }
+
+        if instruction.inspect_args {
+            // Add a statement to the beginning of the function: __instrumentInspectArgs('function_identifier', arguments);
+            let source_text: &'a str = self.allocator.alloc_str(&format!(
+                "__instrumentInspectArgs('{}', arguments, '{}', '{}', '{}', this);",
+                instruction.identifier, self.pkg_name, self.pkg_version, instruction.name
+            ));
+
+            insert_single_statement_into_func(self.allocator, body, 0, source_text);
+        }
+
+        // Todo support return value modification
+    }
+}