use serde::{Deserialize, Serialize};
use std::fs;
use std::path::Path;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct DiscoveredBenchmark {
pub name: String,
pub framework: String,
pub command: String,
pub path: String,
pub language: String,
pub confidence: String,
}
pub fn discover_all(root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
results.extend(scan_rust_criterion(root));
results.extend(scan_go_benchmarks(root));
results.extend(scan_python_pytest_benchmark(root));
results.extend(scan_javascript_benchmark(root));
results.extend(scan_custom_directories(root));
results.sort_by(|a, b| a.name.cmp(&b.name));
results
}
fn scan_rust_criterion(root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
let cargo_toml = root.join("Cargo.toml");
if cargo_toml.is_file()
&& let Ok(content) = fs::read_to_string(&cargo_toml)
{
results.extend(parse_cargo_bench_targets(&content, root));
}
let benches_dir = root.join("benches");
if benches_dir.is_dir() {
results.extend(scan_dir_for_criterion(&benches_dir, root));
}
let mut seen = std::collections::HashSet::new();
results.retain(|b| seen.insert(b.name.clone()));
results
}
fn parse_cargo_bench_targets(content: &str, root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
let lines: Vec<&str> = content.lines().collect();
let mut i = 0;
while i < lines.len() {
let trimmed = lines[i].trim();
if trimmed == "[[bench]]" {
let mut name = None;
let mut harness = true;
let mut j = i + 1;
while j < lines.len() {
let ltrimmed = lines[j].trim();
if ltrimmed.starts_with('[') {
break;
}
if let Some(val) = extract_toml_string_value(ltrimmed, "name") {
name = Some(val);
}
if ltrimmed.starts_with("harness") && ltrimmed.contains("false") {
harness = false;
}
j += 1;
}
if let Some(bench_name) = name {
let confidence = if harness { "medium" } else { "high" };
let framework = if harness { "rust-bench" } else { "criterion" };
let command = format!("cargo bench --bench {bench_name}");
let bench_path = root.join("benches").join(format!("{bench_name}.rs"));
let rel_path = if bench_path.exists() {
format!("benches/{bench_name}.rs")
} else {
"Cargo.toml".to_string()
};
results.push(DiscoveredBenchmark {
name: bench_name,
framework: framework.to_string(),
command,
path: rel_path,
language: "rust".to_string(),
confidence: confidence.to_string(),
});
}
i = j;
} else {
i += 1;
}
}
results
}
fn extract_toml_string_value(line: &str, key: &str) -> Option<String> {
let trimmed = line.trim();
if !trimmed.starts_with(key) {
return None;
}
let rest = trimmed[key.len()..].trim();
let rest = rest.strip_prefix('=')?;
let rest = rest.trim();
let rest = rest.strip_prefix('"')?;
let rest = rest.strip_suffix('"')?;
Some(rest.to_string())
}
fn scan_dir_for_criterion(dir: &Path, root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
let entries = match fs::read_dir(dir) {
Ok(e) => e,
Err(_) => return results,
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
results.extend(scan_dir_for_criterion(&path, root));
} else if path.extension().is_some_and(|ext| ext == "rs")
&& let Ok(content) = fs::read_to_string(&path)
&& (content.contains("criterion_group!") || content.contains("criterion_main!"))
{
let rel_path = path
.strip_prefix(root)
.unwrap_or(&path)
.to_string_lossy()
.replace('\\', "/");
let bench_name = path
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string());
results.push(DiscoveredBenchmark {
name: bench_name.clone(),
framework: "criterion".to_string(),
command: format!("cargo bench --bench {bench_name}"),
path: rel_path,
language: "rust".to_string(),
confidence: "high".to_string(),
});
}
}
results
}
fn scan_go_benchmarks(root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
walk_files(root, &mut |path| {
let name = path.file_name().unwrap_or_default().to_string_lossy();
if name.ends_with("_test.go")
&& let Ok(content) = fs::read_to_string(path)
{
for line in content.lines() {
let trimmed = line.trim();
if let Some(rest) = trimmed.strip_prefix("func Benchmark")
&& let Some(paren_pos) = rest.find('(')
{
let func_name = rest[..paren_pos].trim();
if !func_name.is_empty()
&& func_name.chars().next().is_some_and(|c| c.is_uppercase())
{
let full_name = format!("Benchmark{func_name}");
let rel_path = path
.strip_prefix(root)
.unwrap_or(path)
.to_string_lossy()
.replace('\\', "/");
let pkg_dir = path
.parent()
.and_then(|p| p.strip_prefix(root).ok())
.map(|p| p.to_string_lossy().replace('\\', "/"))
.unwrap_or_else(|| ".".to_string());
results.push(DiscoveredBenchmark {
name: full_name.clone(),
framework: "go-bench".to_string(),
command: format!("go test -bench=^{full_name}$ -benchmem ./{pkg_dir}"),
path: rel_path,
language: "go".to_string(),
confidence: "high".to_string(),
});
}
}
}
}
});
results
}
fn scan_python_pytest_benchmark(root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
walk_files(root, &mut |path| {
let name = path.file_name().unwrap_or_default().to_string_lossy();
if (name.starts_with("test_") || name.ends_with("_test.py") || name.starts_with("bench_"))
&& name.ends_with(".py")
&& let Ok(content) = fs::read_to_string(path)
&& content.contains("benchmark")
&& content.contains("def ")
{
for line in content.lines() {
let trimmed = line.trim();
if let Some(rest) = trimmed.strip_prefix("def ")
&& let Some(paren_pos) = rest.find('(')
&& rest[paren_pos..].contains("benchmark")
{
let func_name = &rest[..paren_pos];
let rel_path = path
.strip_prefix(root)
.unwrap_or(path)
.to_string_lossy()
.replace('\\', "/");
results.push(DiscoveredBenchmark {
name: func_name.to_string(),
framework: "pytest-benchmark".to_string(),
command: format!("pytest --benchmark-only {rel_path}::{func_name}"),
path: rel_path,
language: "python".to_string(),
confidence: "high".to_string(),
});
}
}
}
});
results
}
fn scan_javascript_benchmark(root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
walk_files(root, &mut |path| {
let name = path.file_name().unwrap_or_default().to_string_lossy();
if (name.ends_with(".js") || name.ends_with(".mjs"))
&& (name.contains("bench") || name.contains("perf"))
&& let Ok(content) = fs::read_to_string(path)
&& (content.contains("suite.add") || content.contains("Suite"))
{
let rel_path = path
.strip_prefix(root)
.unwrap_or(path)
.to_string_lossy()
.replace('\\', "/");
let bench_name = path
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string());
results.push(DiscoveredBenchmark {
name: bench_name,
framework: "benchmark.js".to_string(),
command: format!("node {rel_path}"),
path: rel_path,
language: "javascript".to_string(),
confidence: "medium".to_string(),
});
}
});
results
}
fn scan_custom_directories(root: &Path) -> Vec<DiscoveredBenchmark> {
let mut results = Vec::new();
let dirs = ["benchmarks", "bench", "perf"];
for dir_name in &dirs {
let dir = root.join(dir_name);
if dir.is_dir() {
let entries = match fs::read_dir(&dir) {
Ok(e) => e,
Err(_) => continue,
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_file() && is_likely_executable(&path) {
let rel_path = path
.strip_prefix(root)
.unwrap_or(&path)
.to_string_lossy()
.replace('\\', "/");
let bench_name = path
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string());
results.push(DiscoveredBenchmark {
name: bench_name,
framework: "custom".to_string(),
command: format!("./{rel_path}"),
path: rel_path,
language: "unknown".to_string(),
confidence: "low".to_string(),
});
}
}
}
}
results
}
fn is_likely_executable(path: &Path) -> bool {
if let Some(ext) = path.extension() {
let ext = ext.to_string_lossy().to_lowercase();
if matches!(
ext.as_str(),
"sh" | "bash" | "zsh" | "py" | "rb" | "pl" | "exe" | "bat" | "cmd" | "ps1"
) {
return true;
}
}
if let Ok(content) = fs::read(path)
&& content.starts_with(b"#!")
{
return true;
}
path.extension().is_none()
}
fn walk_files(root: &Path, callback: &mut dyn FnMut(&Path)) {
walk_files_inner(root, root, callback);
}
fn walk_files_inner(dir: &Path, root: &Path, callback: &mut dyn FnMut(&Path)) {
let entries = match fs::read_dir(dir) {
Ok(e) => e,
Err(_) => return,
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
let name = path.file_name().unwrap_or_default().to_string_lossy();
if matches!(
name.as_ref(),
"target"
| "node_modules"
| ".git"
| ".hg"
| ".svn"
| "__pycache__"
| "vendor"
| "dist"
| "build"
| ".tox"
| ".venv"
| "venv"
) {
continue;
}
let depth = path
.strip_prefix(root)
.map(|p| p.components().count())
.unwrap_or(0);
if depth < 5 {
walk_files_inner(&path, root, callback);
}
} else {
callback(&path);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
fn setup_temp_dir() -> tempfile::TempDir {
tempfile::tempdir().expect("failed to create temp dir")
}
#[test]
fn test_discover_empty_dir() {
let tmp = setup_temp_dir();
let results = discover_all(tmp.path());
assert!(results.is_empty());
}
#[test]
fn test_extract_toml_string_value() {
assert_eq!(
extract_toml_string_value(r#"name = "my_bench""#, "name"),
Some("my_bench".to_string())
);
assert_eq!(
extract_toml_string_value(r#" name = "spaced" "#, "name"),
Some("spaced".to_string())
);
assert_eq!(
extract_toml_string_value(r#"harness = false"#, "name"),
None
);
assert_eq!(
extract_toml_string_value(r#"name = "has_underscore""#, "name"),
Some("has_underscore".to_string())
);
}
#[test]
fn test_parse_cargo_bench_targets() {
let content = r#"
[package]
name = "myproject"
[[bench]]
name = "my_bench"
harness = false
[[bench]]
name = "another"
"#;
let tmp = setup_temp_dir();
let results = parse_cargo_bench_targets(content, tmp.path());
assert_eq!(results.len(), 2);
assert_eq!(results[0].name, "my_bench");
assert_eq!(results[0].framework, "criterion");
assert_eq!(results[0].confidence, "high");
assert_eq!(results[1].name, "another");
assert_eq!(results[1].framework, "rust-bench");
assert_eq!(results[1].confidence, "medium");
}
#[test]
fn test_scan_criterion_file() {
let tmp = setup_temp_dir();
let benches_dir = tmp.path().join("benches");
fs::create_dir(&benches_dir).unwrap();
fs::write(
benches_dir.join("sort_bench.rs"),
r#"
use criterion::{criterion_group, criterion_main, Criterion};
fn bench_sort(c: &mut Criterion) {
c.bench_function("sort_1000", |b| {
b.iter(|| {
let mut v: Vec<i32> = (0..1000).rev().collect();
v.sort();
})
});
}
criterion_group!(benches, bench_sort);
criterion_main!(benches);
"#,
)
.unwrap();
let results = scan_rust_criterion(tmp.path());
assert_eq!(results.len(), 1);
assert_eq!(results[0].name, "sort_bench");
assert_eq!(results[0].framework, "criterion");
assert_eq!(results[0].language, "rust");
assert_eq!(results[0].confidence, "high");
assert!(
results[0]
.command
.contains("cargo bench --bench sort_bench")
);
}
#[test]
fn test_scan_go_benchmarks() {
let tmp = setup_temp_dir();
fs::write(
tmp.path().join("sort_test.go"),
r#"
package main
import "testing"
func BenchmarkSort(b *testing.B) {
for i := 0; i < b.N; i++ {
// sort something
}
}
func BenchmarkSearch(b *testing.B) {
for i := 0; i < b.N; i++ {
// search something
}
}
func TestNotABenchmark(t *testing.T) {}
"#,
)
.unwrap();
let results = scan_go_benchmarks(tmp.path());
assert_eq!(results.len(), 2);
let names: Vec<&str> = results.iter().map(|b| b.name.as_str()).collect();
assert!(names.contains(&"BenchmarkSort"));
assert!(names.contains(&"BenchmarkSearch"));
for b in &results {
assert_eq!(b.framework, "go-bench");
assert_eq!(b.language, "go");
assert_eq!(b.confidence, "high");
assert!(b.command.contains("go test -bench="));
}
}
#[test]
fn test_scan_python_pytest_benchmark() {
let tmp = setup_temp_dir();
fs::write(
tmp.path().join("test_perf.py"),
r#"
def test_sort_speed(benchmark):
benchmark(sorted, list(range(1000, 0, -1)))
def test_not_a_benchmark():
assert True
"#,
)
.unwrap();
let results = scan_python_pytest_benchmark(tmp.path());
assert_eq!(results.len(), 1);
assert_eq!(results[0].name, "test_sort_speed");
assert_eq!(results[0].framework, "pytest-benchmark");
assert_eq!(results[0].language, "python");
assert_eq!(results[0].confidence, "high");
}
#[test]
fn test_scan_javascript_benchmark() {
let tmp = setup_temp_dir();
fs::write(
tmp.path().join("bench.js"),
r#"
const Benchmark = require('benchmark');
const suite = new Benchmark.Suite;
suite.add('sort', function() {
[3,1,2].sort();
})
.run();
"#,
)
.unwrap();
let results = scan_javascript_benchmark(tmp.path());
assert_eq!(results.len(), 1);
assert_eq!(results[0].name, "bench");
assert_eq!(results[0].framework, "benchmark.js");
assert_eq!(results[0].language, "javascript");
assert_eq!(results[0].confidence, "medium");
}
#[test]
fn test_scan_custom_directories() {
let tmp = setup_temp_dir();
let bench_dir = tmp.path().join("benchmarks");
fs::create_dir(&bench_dir).unwrap();
fs::write(bench_dir.join("run_perf.sh"), "#!/bin/bash\necho hello").unwrap();
fs::write(bench_dir.join("README.md"), "# Benchmarks").unwrap();
let results = scan_custom_directories(tmp.path());
assert_eq!(results.len(), 1);
assert_eq!(results[0].name, "run_perf");
assert_eq!(results[0].framework, "custom");
assert_eq!(results[0].confidence, "low");
}
#[test]
fn test_is_likely_executable() {
let tmp = setup_temp_dir();
let sh_file = tmp.path().join("run.sh");
fs::write(&sh_file, "#!/bin/bash\necho hello").unwrap();
assert!(is_likely_executable(&sh_file));
let py_file = tmp.path().join("bench.py");
fs::write(&py_file, "print('hello')").unwrap();
assert!(is_likely_executable(&py_file));
let md_file = tmp.path().join("README.md");
fs::write(&md_file, "# Hello").unwrap();
assert!(!is_likely_executable(&md_file));
let shebang_file = tmp.path().join("my_bench");
fs::write(&shebang_file, "#!/usr/bin/env python\nprint('bench')").unwrap();
assert!(is_likely_executable(&shebang_file));
}
#[test]
fn test_discover_all_mixed() {
let tmp = setup_temp_dir();
fs::write(
tmp.path().join("Cargo.toml"),
r#"
[package]
name = "test"
[[bench]]
name = "perf_test"
harness = false
"#,
)
.unwrap();
fs::write(
tmp.path().join("algo_test.go"),
r#"
package algo
import "testing"
func BenchmarkAlgo(b *testing.B) {}
"#,
)
.unwrap();
let results = discover_all(tmp.path());
assert!(results.len() >= 2);
let names: Vec<&str> = results.iter().map(|b| b.name.as_str()).collect();
assert!(names.contains(&"perf_test"));
assert!(names.contains(&"BenchmarkAlgo"));
}
#[test]
fn test_discover_all_sorted_by_name() {
let tmp = setup_temp_dir();
let bench_dir = tmp.path().join("benchmarks");
fs::create_dir(&bench_dir).unwrap();
fs::write(bench_dir.join("zebra.sh"), "#!/bin/bash").unwrap();
fs::write(bench_dir.join("alpha.sh"), "#!/bin/bash").unwrap();
let results = discover_all(tmp.path());
if results.len() >= 2 {
assert!(results[0].name <= results[1].name);
}
}
#[test]
fn test_cargo_bench_targets_dedup() {
let tmp = setup_temp_dir();
fs::write(
tmp.path().join("Cargo.toml"),
r#"
[package]
name = "test"
[[bench]]
name = "my_bench"
harness = false
"#,
)
.unwrap();
let benches_dir = tmp.path().join("benches");
fs::create_dir(&benches_dir).unwrap();
fs::write(
benches_dir.join("my_bench.rs"),
"criterion_group!(benches, f);\ncriterion_main!(benches);",
)
.unwrap();
let results = scan_rust_criterion(tmp.path());
let count = results.iter().filter(|b| b.name == "my_bench").count();
assert_eq!(count, 1);
}
#[test]
fn test_walk_files_skips_git_dir() {
let tmp = setup_temp_dir();
let git_dir = tmp.path().join(".git");
fs::create_dir(&git_dir).unwrap();
fs::write(git_dir.join("config"), "core").unwrap();
let mut visited = Vec::new();
walk_files(tmp.path(), &mut |path| {
visited.push(path.to_path_buf());
});
assert!(
visited.is_empty(),
"should not visit files inside .git directory"
);
}
}