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