1use std::path::Path;
29
30use crate::manifest::YogToml;
31
32pub const BENCH_DIR: &str = "benches";
34pub const GENERATED_DIR: &str = ".yog-build/generated/benches";
36const CRITERION: &str = r#"criterion = { version = "0.8", default-features = false, features = ["cargo_bench_support"] }"#;
38const ATTRIBUTE: &str = "#[yog_api::bench]";
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum Takes {
44 Bencher,
46 Criterion,
48}
49
50#[derive(Debug, Clone, PartialEq, Eq)]
52pub struct BenchFile {
53 pub name: String,
54 pub functions: Vec<(String, Takes)>,
55}
56
57pub fn discover(root: &Path) -> Result<Vec<BenchFile>, String> {
59 let dir = root.join(BENCH_DIR);
60 let Ok(entries) = std::fs::read_dir(&dir) else {
61 return Ok(Vec::new());
62 };
63 let mut files = Vec::new();
64 for entry in entries.flatten() {
65 let path = entry.path();
66 if path.extension().and_then(|e| e.to_str()) != Some("rs") {
67 continue;
68 }
69 let Some(name) = path.file_stem().and_then(|s| s.to_str()) else {
70 continue;
71 };
72 let source =
73 std::fs::read_to_string(&path).map_err(|e| format!("{}: {e}", path.display()))?;
74 let functions = functions(&source).map_err(|e| format!("{BENCH_DIR}/{name}.rs: {e}"))?;
75 if !functions.is_empty() {
76 files.push(BenchFile {
77 name: name.to_string(),
78 functions,
79 });
80 }
81 }
82 files.sort_by(|a, b| a.name.cmp(&b.name));
83 Ok(files)
84}
85
86pub fn functions(source: &str) -> Result<Vec<(String, Takes)>, String> {
88 let mut out = Vec::new();
89 let mut rest = source;
90 while let Some(at) = rest.find(ATTRIBUTE) {
91 let line = source[..source.len() - rest.len() + at]
92 .matches('\n')
93 .count()
94 + 1;
95 rest = &rest[at + ATTRIBUTE.len()..];
96 let Some(fn_at) = rest.find("fn ") else {
97 return Err(format!("line {line}: {ATTRIBUTE} must be on a function"));
98 };
99 let after = &rest[fn_at + 3..];
100 let name: String = after
101 .trim_start()
102 .chars()
103 .take_while(|c| c.is_alphanumeric() || *c == '_')
104 .collect();
105 let params = after
106 .split_once('(')
107 .and_then(|(_, p)| p.split_once(')'))
108 .map(|(p, _)| p)
109 .unwrap_or_default();
110 let takes = if params.contains("Criterion") {
111 Takes::Criterion
112 } else if params.contains("Bencher") {
113 Takes::Bencher
114 } else {
115 return Err(format!(
116 "line {line}: `{name}` must take `&mut Bencher` or `&mut Criterion`"
117 ));
118 };
119 out.push((name, takes));
120 }
121 Ok(out)
122}
123
124pub fn main_source(file: &BenchFile) -> String {
126 let mut calls = String::new();
127 for (function, takes) in &file.functions {
128 calls.push_str(&match takes {
129 Takes::Bencher => {
130 format!(
131 " c.bench_function(\"{}/{function}\", bench::{function});\n",
132 file.name
133 )
134 }
135 Takes::Criterion => format!(" bench::{function}(&mut c);\n"),
136 });
137 }
138 format!(
139 "// Generated by yog from {BENCH_DIR}/{name}.rs — do not edit.\n\
140 #[path = \"../../../{BENCH_DIR}/{name}.rs\"]\n\
141 mod bench;\n\n\
142 fn main() {{\n\
143 \x20 let mut c = ::criterion::Criterion::default().configure_from_args();\n\
144 {calls}\
145 \x20 c.final_summary();\n\
146 }}\n",
147 name = file.name,
148 )
149}
150
151pub fn write_sources(root: &Path, files: &[BenchFile]) -> Result<(), String> {
153 let dir = root.join(GENERATED_DIR);
154 if let Ok(entries) = std::fs::read_dir(&dir) {
155 for entry in entries.flatten() {
156 let path = entry.path();
157 let stale = path
158 .file_stem()
159 .and_then(|s| s.to_str())
160 .is_none_or(|stem| !files.iter().any(|f| f.name == stem));
161 if stale {
162 let _ = std::fs::remove_file(path);
163 }
164 }
165 }
166 for file in files {
167 crate::write_if_changed(&dir.join(format!("{}.rs", file.name)), &main_source(file))?;
168 }
169 Ok(())
170}
171
172pub fn cargo_sections(meta: &YogToml, files: &[BenchFile]) -> String {
175 if files.is_empty() {
176 return String::new();
177 }
178 let mut out = String::from("\n[dev-dependencies]\n");
179 for (name, spec) in &meta.dev_dependencies {
180 out.push_str(&format!("{name} = {spec}\n"));
181 }
182 if !meta
183 .dev_dependencies
184 .iter()
185 .any(|(name, _)| name == "criterion")
186 {
187 out.push_str(CRITERION);
188 out.push('\n');
189 }
190 for file in files {
191 out.push_str(&format!(
192 "\n[[bench]]\nname = \"{name}\"\npath = \"generated/benches/{name}.rs\"\nharness = false\n",
193 name = file.name
194 ));
195 }
196 out
197}
198
199#[cfg(test)]
200mod tests {
201 use super::*;
202
203 #[test]
204 fn finds_marked_functions_in_order() {
205 let src = "use criterion::{Bencher, Criterion};\n\
206 fn helper() {}\n\
207 #[yog_api::bench]\nfn parse(b: &mut Bencher) { b.iter(|| 1) }\n\
208 #[yog_api::bench]\npub fn groups(c: &mut Criterion) {}\n";
209 assert_eq!(
210 functions(src).unwrap(),
211 [
212 ("parse".to_string(), Takes::Bencher),
213 ("groups".to_string(), Takes::Criterion)
214 ]
215 );
216 }
217
218 #[test]
219 fn rejects_a_function_without_a_bencher() {
220 let e = functions("\n#[yog_api::bench]\nfn nope(x: u32) {}\n").unwrap_err();
221 assert!(e.starts_with("line 2:"), "{e}");
222 }
223
224 #[test]
225 fn main_runs_each_function() {
226 let file = BenchFile {
227 name: "casting".into(),
228 functions: vec![
229 ("parse".into(), Takes::Bencher),
230 ("vm".into(), Takes::Criterion),
231 ],
232 };
233 let main = main_source(&file);
234 assert!(main.contains("#[path = \"../../../benches/casting.rs\"]"));
235 assert!(main.contains("c.bench_function(\"casting/parse\", bench::parse);"));
236 assert!(main.contains("bench::vm(&mut c);"));
237 }
238}