Build injected functions using AstBuilder

Author: timokoessler

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

@@ -1,67 +1,197 @@
-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(),
-    );
-}
+use oxc_allocator::{Allocator, Box, Vec as OxcVec};
+use oxc_ast::{
+    ast::{
+        Argument, ArrayExpressionElement, AssignmentOperator, AssignmentTarget, Expression,
+        FunctionBody,
+    },
+    AstBuilder, NONE,
+};
+use oxc_span::SPAN;
 
 // Add a statement to the beginning of the function: __instrumentInspectArgs('function_identifier', arguments);
 pub fn insert_inspect_args<'a>(
     allocator: &'a Allocator,
+    builder: &'a AstBuilder,
     identifier: &str,
     pkg_version: &'a str,
     body: &mut Box<'a, FunctionBody<'a>>,
 ) {
-    let source_text: &'a str = allocator.alloc_str(&format!(
-        "__instrumentInspectArgs('{}', arguments, '{}', this);",
-        identifier, pkg_version
-    ));
+    let mut inspect_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(4);
+
+    // Add the identifier to the arguments
+    inspect_args.push(Argument::StringLiteral(builder.alloc_string_literal(
+        SPAN,
+        allocator.alloc_str(identifier),
+        None,
+    )));
+
+    // Add the arguments object as the second argument
+    inspect_args.push(builder.expression_identifier(SPAN, "arguments").into());
+
+    // Add the package version as the third argument
+    inspect_args.push(Argument::StringLiteral(builder.alloc_string_literal(
+        SPAN,
+        allocator.alloc_str(pkg_version),
+        None,
+    )));
+
+    // Add the `this` context as the fourth argument
+    inspect_args.push(builder.expression_identifier(SPAN, "this").into());
 
-    insert_single_statement_into_func(allocator, body, 0, source_text);
+    // Build and add a call expression
+    let call_expr = builder.expression_call(
+        SPAN,
+        builder.expression_identifier(SPAN, "__instrumentInspectArgs"),
+        NONE,
+        inspect_args,
+        false,
+    );
+
+    let stmt_expression = builder.statement_expression(SPAN, call_expr);
+
+    body.statements.insert(0, stmt_expression);
 }
 
 // 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,
+    builder: &'a AstBuilder,
     identifier: &str,
-    arg_names_str: &str,
+    arg_names: &Vec<String>,
     body: &mut Box<'a, FunctionBody<'a>>,
     modify_arguments_object: bool,
 ) {
     if modify_arguments_object {
         // If we are modifying the arguments object, we need to use the arguments object directly
         // instead of the individual arguments
-        let source_text: &'a str = allocator.alloc_str(&format!(
-            "Object.assign(arguments, __instrumentModifyArgs('{}', Array.from(arguments)));",
-            identifier
-        ));
-        insert_single_statement_into_func(allocator, body, 0, source_text);
+        // Object.assign(arguments, __instrumentModifyArgs('id', Array.from(arguments)));
+
+        let mut obj_assign_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(2);
+
+        // First argument is the arguments object
+        obj_assign_args.push(builder.expression_identifier(SPAN, "arguments").into());
+
+        // Second argument is the call to __instrumentModifyArgs
+
+        let mut instrument_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(2);
+        // Add the identifier to the arguments
+        instrument_args.push(Argument::StringLiteral(builder.alloc_string_literal(
+            SPAN,
+            allocator.alloc_str(identifier),
+            None,
+        )));
+
+        let mut array_from_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(1);
+        // Add the arguments object as the first argument to Array.from
+        array_from_args.push(builder.expression_identifier(SPAN, "arguments").into());
+
+        let array_from_call = builder.expression_call(
+            SPAN,
+            Expression::StaticMemberExpression(builder.alloc_static_member_expression(
+                SPAN,
+                builder.expression_identifier(SPAN, "Array"),
+                builder.identifier_name(SPAN, "from"),
+                false,
+            )),
+            NONE,
+            array_from_args,
+            false,
+        );
+
+        instrument_args.push(array_from_call.into());
+
+        let instrument_modify_args_call = builder.expression_call(
+            SPAN,
+            builder.expression_identifier(SPAN, "__instrumentModifyArgs"),
+            NONE,
+            instrument_args,
+            false,
+        );
+
+        obj_assign_args.push(instrument_modify_args_call.into());
+
+        let obj_assign_call_expr = builder.expression_call(
+            SPAN,
+            Expression::StaticMemberExpression(builder.alloc_static_member_expression(
+                SPAN,
+                builder.expression_identifier(SPAN, "Object"),
+                builder.identifier_name(SPAN, "assign"),
+                false,
+            )),
+            NONE,
+            obj_assign_args,
+            false,
+        );
+
+        let stmt_expression = builder.statement_expression(SPAN, obj_assign_call_expr);
+
+        body.statements.insert(0, stmt_expression);
+
         return;
     }
 
-    if arg_names_str.is_empty() {
-        // If there are no arguments to modify, we can skip this step
+    if arg_names.is_empty() {
+        // If there are no arguments to modify, we can skip the modification
         return;
     }
 
-    let source_text: &'a str = allocator.alloc_str(&format!(
-        "[{}] = __instrumentModifyArgs('{}', [{}]);",
-        arg_names_str, identifier, arg_names_str
+    use oxc_ast::ast::AssignmentTargetMaybeDefault;
+
+    let mut array_assignment_target_identifiers: OxcVec<
+        'a,
+        Option<AssignmentTargetMaybeDefault<'a>>,
+    > = builder.vec_with_capacity(arg_names.len());
+
+    for name in arg_names {
+        array_assignment_target_identifiers.push(Some(
+            AssignmentTargetMaybeDefault::AssignmentTargetIdentifier(
+                builder.alloc_identifier_reference(SPAN, allocator.alloc_str(name)),
+            ),
+        ));
+    }
+
+    let mut instrument_modify_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(2);
+    // Add the identifier to the arguments
+    instrument_modify_args.push(Argument::StringLiteral(builder.alloc_string_literal(
+        SPAN,
+        allocator.alloc_str(identifier),
+        None,
+    )));
+
+    let mut instrument_modify_args_array_elements: OxcVec<'a, ArrayExpressionElement<'a>> =
+        builder.vec_with_capacity(arg_names.len());
+
+    for name in arg_names {
+        instrument_modify_args_array_elements.push(ArrayExpressionElement::Identifier(
+            builder.alloc_identifier_reference(SPAN, allocator.alloc_str(name)),
+        ));
+    }
+
+    instrument_modify_args.push(Argument::ArrayExpression(
+        builder.alloc_array_expression(SPAN, instrument_modify_args_array_elements),
     ));
 
-    insert_single_statement_into_func(allocator, body, 0, source_text);
+    let instrument_modify_args_call = builder.expression_call(
+        SPAN,
+        builder.expression_identifier(SPAN, "__instrumentModifyArgs"),
+        NONE,
+        instrument_modify_args,
+        false,
+    );
+
+    let arr_assignment_expr = builder.expression_assignment(
+        SPAN,
+        AssignmentOperator::Assign,
+        AssignmentTarget::ArrayAssignmentTarget(builder.alloc_array_assignment_target(
+            SPAN,
+            array_assignment_target_identifiers,
+            None,
+        )),
+        instrument_modify_args_call,
+    );
+
+    let stmt_expression = builder.statement_expression(SPAN, arr_assignment_expr);
+
+    body.statements.insert(0, stmt_expression);
 }

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

@@ -1,4 +1,4 @@
-use oxc_allocator::{Allocator, Vec};
+use oxc_allocator::{Allocator, Vec as OxcVec};
 use oxc_ast::{
     ast::{
         Argument, BindingPatternKind, BindingProperty, ImportDeclarationSpecifier,
@@ -35,19 +35,19 @@ pub fn insert_import_statement<'a>(
     has_module_syntax: bool,
     allocator: &'a Allocator,
     builder: &'a AstBuilder,
-    body: &mut Vec<'a, Statement<'a>>,
+    body: &mut OxcVec<'a, Statement<'a>>,
     file_instructions: &FileInstructions,
 ) {
     // Common JS require() statement
     if source_type.is_script() || source_type.is_unambiguous() && !has_module_syntax {
-        let mut require_args: Vec<'a, Argument<'a>> = builder.vec_with_capacity(1);
+        let mut require_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(1);
         require_args.push(Argument::StringLiteral(builder.alloc_string_literal(
             SPAN,
             allocator.alloc_str(INSTRUMENT_IMPORT_SOURCE),
             None,
         )));
 
-        let mut binding_properties: Vec<'a, BindingProperty<'a>> = builder.vec_with_capacity(3);
+        let mut binding_properties: OxcVec<'a, BindingProperty<'a>> = builder.vec_with_capacity(3);
 
         for (method_name, predicate) in IMPORT_METHODS.iter() {
             // Only import the function if it is used in the file
@@ -66,6 +66,11 @@ pub fn insert_import_statement<'a>(
             }
         }
 
+        if binding_properties.is_empty() {
+            // If there are no methods to import, we can skip the require statement
+            return;
+        }
+
         let mut declarations = builder.vec_with_capacity(1);
         declarations.push(builder.variable_declarator(
             SPAN,
@@ -102,7 +107,7 @@ pub fn insert_import_statement<'a>(
     }
     // else: ESM import statement
 
-    let mut specifiers: Vec<'a, ImportDeclarationSpecifier<'a>> = builder.vec_with_capacity(3);
+    let mut specifiers: OxcVec<'a, ImportDeclarationSpecifier<'a>> = builder.vec_with_capacity(3);
     for (method_name, predicate) in IMPORT_METHODS.iter() {
         if file_instructions.functions.iter().any(|f| predicate(f)) {
             specifiers.push(builder.import_declaration_specifier_import_specifier(
@@ -114,6 +119,11 @@ pub fn insert_import_statement<'a>(
         }
     }
 
+    if specifiers.is_empty() {
+        // If there are no methods to import, we can skip the import statement
+        return;
+    }
+
     let import_stmt = Statement::ImportDeclaration(builder.alloc_import_declaration(
         SPAN,
         Some(specifiers),

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

@@ -2,6 +2,5 @@ pub mod get_arg_names;
 pub mod get_name_str_for_member_expr;
 pub mod insert_code;
 pub mod insert_import_statement;
-pub mod parse_js_code_to_statements;
 pub mod select_sourcetype_based_on_enum;
 pub mod transform_return_statements;

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

@@ -1,4 +1,4 @@
-use oxc_allocator::{Allocator, Box, Vec};
+use oxc_allocator::{Allocator, Box, Vec as OxcVec};
 use oxc_ast::{
     ast::{Argument, FunctionBody, Statement},
     AstBuilder, NONE,
@@ -15,7 +15,7 @@ pub fn transform_return_statements<'a>(
     for statement in &mut body.statements {
         match statement {
             Statement::ReturnStatement(return_stmt) => {
-                let mut instrument_args: Vec<'a, Argument<'a>> = builder.vec_with_capacity(2);
+                let mut instrument_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(2);
 
                 // Add the identifier to the arguments
                 instrument_args.push(Argument::StringLiteral(builder.alloc_string_literal(

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

@@ -57,11 +57,11 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         let body = node.value.body.as_mut().unwrap();
 
         if instruction.modify_args && !arg_names.is_empty() {
-            let arg_names_str = arg_names.join(", ");
             insert_modify_args(
                 self.allocator,
+                self.ast_builder,
                 &instruction.identifier,
-                &arg_names_str,
+                &arg_names,
                 body,
                 instruction.modify_arguments_object,
             );
@@ -70,6 +70,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         if instruction.inspect_args {
             insert_inspect_args(
                 self.allocator,
+                self.ast_builder,
                 &instruction.identifier,
                 self.pkg_version,
                 body,
@@ -138,11 +139,11 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         let body = function_expression.body.as_mut().unwrap();
 
         if instruction.modify_args {
-            let arg_names_str = arg_names.join(", ");
             insert_modify_args(
                 self.allocator,
+                self.ast_builder,
                 &instruction.identifier,
-                &arg_names_str,
+                &arg_names,
                 body,
                 instruction.modify_arguments_object,
             );
@@ -151,6 +152,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         if instruction.inspect_args {
             insert_inspect_args(
                 self.allocator,
+                self.ast_builder,
                 &instruction.identifier,
                 self.pkg_version,
                 body,