feat: avoid parsing source code multiple times when patching imports

Author: not-matthias

go-runner/src/builder/patcher.rs

@@ -1,6 +1,7 @@
 //! Patches the imports to use codspeed rather than the official "testing" package.
 
 use crate::prelude::*;
+use itertools::Itertools;
 use rayon::prelude::*;
 use std::fs;
 use std::path::Path;
@@ -71,19 +72,22 @@ pub fn patch_imports<P: AsRef<Path>>(folder: P) -> anyhow::Result<()> {
 
 /// Internal function to apply import patterns to Go source code
 pub fn patch_imports_for_source(source: &str) -> String {
-    let replace_import = |mut source: String, import_path: &str, replacement: &str| -> String {
+    let mut source = source.to_string();
+
+    // If we can't parse the source, skip this replacement
+    // This can happen with template files or malformed Go code
+    let parsed = match gosyn::parse_source(&source) {
+        Ok(p) => p,
+        Err(_) => return source,
+    };
+
+    let mut replacements = vec![];
+    let mut find_replace_range = |import_path: &str, replacement: &str| {
         // Optimization: check if the import path exists in the source before parsing
         if !source.contains(import_path) {
-            return source;
+            return;
         }
 
-        // If we can't parse the source, skip this replacement
-        // This can happen with template files or malformed Go code
-        let parsed = match gosyn::parse_source(&source) {
-            Ok(p) => p,
-            Err(_) => return source,
-        };
-
         if let Some(import) = parsed
             .imports
             .iter()
@@ -92,39 +96,43 @@ pub fn patch_imports_for_source(source: &str) -> String {
             let start_pos = import.path.pos;
             let end_pos = start_pos + import.path.value.len();
 
-            source.replace_range(start_pos..end_pos, replacement);
+            replacements.push((start_pos..end_pos, replacement.to_string()));
         }
-
-        source
     };
 
-    let mut source = replace_import(
-        source.to_string(),
-        "testing",
-        "testing \"github.com/CodSpeedHQ/codspeed-go/testing/testing\"",
-    );
-
     // Then replace sub-packages like "testing/synctest"
     for testing_pkg in &["fstest", "iotest", "quick", "slogtest", "synctest"] {
-        source = replace_import(
-            source.to_string(),
+        find_replace_range(
             &format!("testing/{}", testing_pkg),
             &format!(
                 "{testing_pkg} \"github.com/CodSpeedHQ/codspeed-go/testing/testing/{testing_pkg}\""
             ),
         );
     }
 
-    let source = replace_import(
-        source,
+    find_replace_range(
+        "testing",
+        "testing \"github.com/CodSpeedHQ/codspeed-go/testing/testing\"",
+    );
+    find_replace_range(
         "github.com/thejerf/slogassert",
         "\"github.com/CodSpeedHQ/codspeed-go/pkg/slogassert\"",
     );
-    replace_import(
-        source,
+    find_replace_range(
         "github.com/frankban/quicktest",
         "\"github.com/CodSpeedHQ/codspeed-go/pkg/quicktest\"",
-    )
+    );
+
+    // Apply replacements in reverse order to avoid shifting positions
+    for (range, replacement) in replacements
+        .into_iter()
+        .sorted_by_key(|(range, _)| range.start)
+        .rev()
+    {
+        source.replace_range(range, &replacement);
+    }
+
+    source
 }
 
 /// Patches imports and package in specific test files
@@ -390,6 +398,17 @@ func BenchmarkExample(b *testing.B) {
 func TestExample(t *testing.T) {
     s := "package main"
 }
+"#;
+
+    const MANY_TESTING_IMPORTS: &str = r#"package subpackages
+import (
+	"bytes"
+	"io"
+	"testing"
+	"testing/fstest"
+	"testing/iotest"
+	"testing/synctest"
+)
 "#;
 
     #[rstest]
@@ -408,6 +427,7 @@ func TestExample(t *testing.T) {
         MULTILINE_IMPORT_WITH_TESTING_STRING
     )]
     #[case("package_main", PACKAGE_MAIN)]
+    #[case("many_testing_imports", MANY_TESTING_IMPORTS)]
     fn test_patch_go_source(#[case] test_name: &str, #[case] source: &str) {
         let result = patch_imports_for_source(source);
         let result = patch_package_for_source(result).unwrap();

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__many_testing_imports.snap

@@ -0,0 +1,13 @@
+---
+source: go-runner/src/builder/patcher.rs
+expression: result
+---
+package subpackages
+import (
+	"bytes"
+	"io"
+	testing "github.com/CodSpeedHQ/codspeed-go/testing/testing"
+	fstest "github.com/CodSpeedHQ/codspeed-go/testing/testing/fstest"
+	iotest "github.com/CodSpeedHQ/codspeed-go/testing/testing/iotest"
+	synctest "github.com/CodSpeedHQ/codspeed-go/testing/testing/synctest"
+)