From 9c4ee86474ac8298e5d2ef2b58a8e5f2160ef914 Mon Sep 17 00:00:00 2001
From: Zane Duffield <duffieldzane@gmail.com>
Date: Sun, 21 Apr 2024 23:03:06 +1000
Subject: [PATCH] Avoid iterating files when `--exact` is passed in

This fixes a quadratic performance issue in which the test directories are iterated for every test
case when run under nextest.
---
 src/runner.rs                   | 67 ++++++++++++++++++++++++++++-----
 src/utils.rs                    |  5 +++
 tests/files/::colon::dir/::.txt |  1 +
 tests/files/::colon::dir/a.txt  |  1 +
 4 files changed, 65 insertions(+), 9 deletions(-)
 create mode 100644 tests/files/::colon::dir/::.txt
 create mode 100644 tests/files/::colon::dir/a.txt

diff --git a/src/runner.rs b/src/runner.rs
index 1582cf0bb..81608a645 100644
--- a/src/runner.rs
+++ b/src/runner.rs
@@ -11,8 +11,7 @@ use libtest_mimic::{Arguments, Trial};
 pub fn runner(requirements: &[Requirements]) -> ExitCode {
     let args = Arguments::from_args();
 
-    let mut tests: Vec<_> = requirements.iter().flat_map(|req| req.expand()).collect();
-    tests.sort_unstable_by(|a, b| a.name().cmp(b.name()));
+    let tests = find_tests(&args, requirements);
 
     let conclusion = libtest_mimic::run(&args, tests);
 
@@ -25,6 +24,43 @@ pub fn runner(requirements: &[Requirements]) -> ExitCode {
     conclusion.exit_code()
 }
 
+fn find_tests(args: &Arguments, requirements: &[Requirements]) -> Vec<Trial> {
+    let tests: Vec<_> = if let Some(exact_filter) = exact_filter(args) {
+        let exact_tests: Vec<_> = requirements
+            .iter()
+            .flat_map(|req| req.exact(exact_filter))
+            .collect();
+
+        match exact_tests.len() {
+            0 if is_nextest() => {
+                panic!("Failed to find exact match for filter {exact_filter}");
+            }
+            len @ (2..) if is_nextest() => {
+                panic!("Only expected one but found {len} exact matches for filter {exact_filter}");
+            }
+            _ => {}
+        }
+        exact_tests
+    } else {
+        let mut tests: Vec<_> = requirements.iter().flat_map(|req| req.expand()).collect();
+        tests.sort_unstable_by(|a, b| a.name().cmp(b.name()));
+        tests
+    };
+    tests
+}
+
+fn is_nextest() -> bool {
+    std::env::var_os("NEXTEST").is_some_and(|val| val.eq("1"))
+}
+
+fn exact_filter(args: &Arguments) -> Option<&str> {
+    if args.exact && args.skip.is_empty() {
+        args.filter.as_deref()
+    } else {
+        None
+    }
+}
+
 #[doc(hidden)]
 pub struct Requirements {
     test: TestFn,
@@ -49,6 +85,25 @@ impl Requirements {
         }
     }
 
+    fn trial(&self, path: Utf8PathBuf) -> Trial {
+        let testfn = self.test;
+        let name = utils::derive_test_name(&self.root, &path, &self.test_name);
+        Trial::test(name, move || {
+            testfn
+                .call(&path)
+                .map_err(|err| format!("{:?}", err).into())
+        })
+    }
+
+    fn exact(&self, filter: &str) -> Option<Trial> {
+        let path = utils::derive_test_path(&self.root, filter, &self.test_name)?;
+        if path.exists() {
+            Some(self.trial(path))
+        } else {
+            None
+        }
+    }
+
     /// Scans all files in a given directory, finds matching ones and generates a test descriptor
     /// for each of them.
     fn expand(&self) -> Vec<Trial> {
@@ -66,13 +121,7 @@ impl Requirements {
                         error
                     )
                 }) {
-                    let testfn = self.test;
-                    let name = utils::derive_test_name(&self.root, &path, &self.test_name);
-                    Some(Trial::test(name, move || {
-                        testfn
-                            .call(&path)
-                            .map_err(|err| format!("{:?}", err).into())
-                    }))
+                    Some(self.trial(path))
                 } else {
                     None
                 }
diff --git a/src/utils.rs b/src/utils.rs
index bcd64789a..d17f13bb3 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -33,3 +33,8 @@ pub fn derive_test_name(root: &Utf8Path, path: &Utf8Path, test_name: &str) -> St
 
     format!("{}::{}", test_name, relative)
 }
+
+pub fn derive_test_path(root: &Utf8Path, filter: &str, test_name: &str) -> Option<Utf8PathBuf> {
+    let relative = filter.strip_prefix(test_name)?.strip_prefix("::")?;
+    Some(root.join(relative))
+}
diff --git a/tests/files/::colon::dir/::.txt b/tests/files/::colon::dir/::.txt
new file mode 100644
index 000000000..eaabfa772
--- /dev/null
+++ b/tests/files/::colon::dir/::.txt
@@ -0,0 +1 @@
+floop
\ No newline at end of file
diff --git a/tests/files/::colon::dir/a.txt b/tests/files/::colon::dir/a.txt
new file mode 100644
index 000000000..ca2fb7494
--- /dev/null
+++ b/tests/files/::colon::dir/a.txt
@@ -0,0 +1 @@
+flarp
\ No newline at end of file