Support function assignments for instrumentation

Author: timokoessler

…nsformer/helpers/get_method_arg_names.rs …/js_transformer/helpers/get_arg_names.rsinstrumentation-wasm/src/js_transformer/helpers/get_method_arg_names.rs renamed to instrumentation-wasm/src/js_transformer/helpers/get_arg_names.rs instrumentation-wasm/src/js_transformer/helpers/get_method_arg_names.rs renamed to instrumentation-wasm/src/js_transformer/helpers/get_arg_names.rs

@@ -1,4 +1,4 @@
-use oxc_ast::ast::MethodDefinition;
+use oxc_ast::ast::{Function, MethodDefinition};
 
 pub fn get_method_arg_names(method_definition: &MethodDefinition) -> Vec<String> {
     let mut arg_names = Vec::new();
@@ -14,3 +14,13 @@ pub fn get_method_arg_names(method_definition: &MethodDefinition) -> Vec<String>
 
     arg_names
 }
+
+pub fn get_function_arg_names(method_definition: &Function) -> Vec<String> {
+    let mut arg_names = Vec::new();
+
+    method_definition.params.items.iter().for_each(|param| {
+        arg_names.push(param.pattern.get_identifier_name().unwrap().to_string());
+    });
+
+    arg_names
+}

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

@@ -0,0 +1,51 @@
+use oxc_allocator::Allocator;
+use oxc_ast::ast::{Expression, MemberExpression};
+
+pub fn get_name_str_for_member_expr<'a>(
+    allocator: &'a Allocator,
+    member_expr: &MemberExpression,
+) -> Option<&'a str> {
+    let mut obj_name_str: &str = "";
+    let mut prop_name_str: &str = "";
+
+    match member_expr {
+        MemberExpression::StaticMemberExpression(static_member_expr) => {
+            prop_name_str = static_member_expr.property.name.as_str()
+        }
+        MemberExpression::ComputedMemberExpression(computed_member_expr) => {
+            match &computed_member_expr.expression {
+                Expression::Identifier(identifier_ref) => {
+                    prop_name_str = identifier_ref.name.as_str();
+                }
+                _ => {
+                    // Unsupported AST type
+                    return None;
+                }
+            }
+        }
+        _ => {
+            // Unsupported AST type
+            return None;
+        }
+    }
+
+    match member_expr.object() {
+        Expression::Identifier(identifier_ref) => {
+            obj_name_str = identifier_ref.name.as_str();
+        }
+        _ => {
+            // Unsupported AST type
+            return None;
+        }
+    }
+
+    if obj_name_str.is_empty() || prop_name_str.is_empty() {
+        return None;
+    }
+
+    if member_expr.is_computed() {
+        return Some(allocator.alloc_str(&format!("{}[{}]", obj_name_str, prop_name_str)));
+    }
+
+    Some(&allocator.alloc_str(&format!("{}.{}", obj_name_str, prop_name_str)))
+}

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

@@ -0,0 +1,52 @@
+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(),
+    );
+}
+
+// Add a statement to the beginning of the function: __instrumentInspectArgs('function_identifier', arguments);
+pub fn insert_inspect_args<'a>(
+    allocator: &'a Allocator,
+    identifier: &str,
+    pkg_name: &'a str,
+    pkg_version: &'a str,
+    instruction_name: &str,
+    body: &mut Box<'a, FunctionBody<'a>>,
+) {
+    let source_text: &'a str = allocator.alloc_str(&format!(
+        "__instrumentInspectArgs('{}', arguments, '{}', '{}', '{}', this);",
+        identifier, pkg_name, pkg_version, instruction_name
+    ));
+
+    insert_single_statement_into_func(allocator, body, 0, source_text);
+}
+
+// Modify the arguments by adding a statement to the beginning of the function
+// [arg1, arg2, ...] = __instrumentModifyArgs('function_identifier', [arg1, arg2, ...]);
+pub fn insert_modify_args<'a>(
+    allocator: &'a Allocator,
+    identifier: &str,
+    arg_names_str: &str,
+    body: &mut Box<'a, FunctionBody<'a>>,
+) {
+    let source_text: &'a str = allocator.alloc_str(&format!(
+        "[{}] = __instrumentModifyArgs('{}', [{}]);",
+        arg_names_str, identifier, arg_names_str
+    ));
+
+    insert_single_statement_into_func(allocator, body, 0, source_text);
+}

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

@@ -1,5 +1,6 @@
+pub mod get_arg_names;
 pub mod get_import_code_str;
-pub mod get_method_arg_names;
-pub mod insert_single_statement_into_func;
+pub mod get_name_str_for_member_expr;
+pub mod insert_code;
 pub mod parse_js_code_to_statements;
 pub mod select_sourcetype_based_on_enum;

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

@@ -10,7 +10,7 @@ pub fn parse_js_code_to_statements<'a>(
 ) -> Vec<'a, Statement<'a>> {
     let parser_result = Parser::new(&allocator, source_text, source_type)
         .with_options(ParseOptions {
-            allow_return_outside_function: true,
+            allow_return_outside_function: true, // Could be partical code string that should be parsed
             ..ParseOptions::default()
         })
         .parse();

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

@@ -29,12 +29,7 @@ pub fn transform_code_str(
 
     let source_type = select_sourcetype_based_on_enum(src_type);
 
-    let parser_result = Parser::new(&allocator, &code, source_type)
-        .with_options(ParseOptions {
-            allow_return_outside_function: true, // Maybe some run
-            ..ParseOptions::default()
-        })
-        .parse();
+    let parser_result = Parser::new(&allocator, &code, source_type).parse();
 
     if parser_result.panicked || parser_result.errors.len() > 0 {
         return format!("#ERR: {:?}", parser_result.errors);

instrumentation-wasm/src/js_transformer/transformer.rs

@@ -1,10 +1,14 @@
-use super::helpers::get_method_arg_names;
-use super::instructions::FileInstructions;
 use oxc_allocator::Allocator;
-use oxc_ast::ast::MethodDefinition;
+use oxc_ast::ast::{AssignmentOperator, Expression, MethodDefinition};
 use oxc_traverse::{Traverse, TraverseCtx};
 
-use super::helpers::insert_single_statement_into_func::insert_single_statement_into_func;
+use super::helpers::{
+    get_arg_names::get_function_arg_names,
+    get_arg_names::get_method_arg_names,
+    get_name_str_for_member_expr::get_name_str_for_member_expr,
+    insert_code::{insert_inspect_args, insert_modify_args},
+};
+use super::instructions::FileInstructions;
 
 pub struct Transformer<'a> {
     pub allocator: &'a Allocator,
@@ -44,34 +48,107 @@ impl<'a> Traverse<'a> for Transformer<'a> {
 
         // 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)
+            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);
+            insert_modify_args(
+                self.allocator,
+                &instruction.identifier,
+                &arg_names_str,
+                body,
+            );
         }
 
         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_inspect_args(
+                self.allocator,
+                &instruction.identifier,
+                self.pkg_name,
+                self.pkg_version,
+                &instruction.name,
+                body,
+            );
+        }
+
+        // Todo support return value modification
+    }
+
+    fn enter_assignment_expression(
+        &mut self,
+        node: &mut oxc_ast::ast::AssignmentExpression<'a>,
+        _ctx: &mut TraverseCtx<'a>,
+    ) {
+        if node.operator != AssignmentOperator::Assign {
+            // Return if operator is not =
+            return;
+        }
+
+        if !node.left.is_member_expression() || !node.right.is_function() {
+            return;
+        }
 
-            insert_single_statement_into_func(self.allocator, body, 0, source_text);
+        let member_expression = node.left.as_member_expression().unwrap();
+
+        let name_str_opt = get_name_str_for_member_expr(self.allocator, member_expression);
+        if name_str_opt.is_none() {
+            // Could not determine the name of the assignment
+            return;
+        }
+        let name_str = name_str_opt.unwrap();
+
+        let matching_instruction = self
+            .file_instructions
+            .functions
+            .iter()
+            .find(|f| f.node_type == "FunctionAssignment" && f.name == name_str);
+
+        if matching_instruction.is_none() {
+            // This function assignment should not be instrumented
+            return;
+        }
+
+        let instruction = matching_instruction.unwrap();
+
+        // We need to modify the function expression
+        let function_expression = match &mut node.right {
+            Expression::FunctionExpression(func_expr) => func_expr,
+            _ => return,
+        };
+
+        // We need to collect the arg names before we make the body mutable
+        let arg_names = if instruction.modify_args {
+            get_function_arg_names(function_expression)
+        } else {
+            Vec::new()
+        };
+
+        let body = function_expression.body.as_mut().unwrap();
+
+        if instruction.modify_args && !arg_names.is_empty() {
+            let arg_names_str = arg_names.join(", ");
+            insert_modify_args(
+                self.allocator,
+                &instruction.identifier,
+                &arg_names_str,
+                body,
+            );
+        }
+
+        if instruction.inspect_args {
+            insert_inspect_args(
+                self.allocator,
+                &instruction.identifier,
+                self.pkg_name,
+                self.pkg_version,
+                &instruction.name,
+                body,
+            );
         }
 
         // Todo support return value modification