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