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