Skip to main content

yog_cli/project/
bench.rs

1//! `yog bench` — a mod's benchmarks: `benches/*.rs` files whose
2//! `#[yog_api::bench]` functions run under criterion.
3//!
4//! ```ignore
5//! // benches/casting.rs
6//! use criterion::{Bencher, Criterion};
7//!
8//! #[yog_api::bench]
9//! fn parse(b: &mut Bencher) {
10//!     b.iter(|| hexcasting::math::HexPattern::from_angles("qaq", HexDir::East));
11//! }
12//!
13//! #[yog_api::bench]
14//! fn vm(c: &mut Criterion) {
15//!     let mut group = c.benchmark_group("vm");
16//!     // ...
17//! }
18//! ```
19//!
20//! No harness, no `criterion_main!`, no Cargo.toml: each file becomes a
21//! `[[bench]]` target of the generated crate whose `main`
22//! (`obj/generated/benches/<file>.rs`) runs its functions in source
23//! order. A function taking `&mut Bencher` is one benchmark named
24//! `<file>/<fn>`; one taking `&mut Criterion` sets up its own (groups,
25//! throughput, batched inputs). The mod is linked in as a library under
26//! its crate name; with no game, the bus's actions answer nothing
27//! (their replies' defaults).
28
29use std::path::Path;
30use yog_data::Diagnostic;
31
32use crate::project::manifest::YogToml;
33use yog_codegen::{Header, Regenerate};
34use yog_codegen_rust::Rust;
35
36/// Where the mod's benchmark files live, under the project root.
37pub const BENCH_DIR: &str = "benches";
38/// The generated `main`s, under the project root.
39pub const GENERATED_DIR: &str =
40    yog_ids::concatcp!(yog_ids::project::GENERATED, yog_ids::sep::PATH, "benches");
41/// criterion, added to `[dev-dependencies]` unless yog.toml lists its own.
42const CRITERION: &str = r#"criterion = { version = "0.8", default-features = false, features = ["cargo_bench_support"] }"#;
43/// The marker attribute, as written in a benchmark file.
44const ATTRIBUTE: &str = "#[yog_api::bench]";
45
46/// What a benchmark function takes.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub(crate) enum Takes {
49    /// `&mut Bencher` — one benchmark, `b.iter(...)`.
50    Bencher,
51    /// `&mut Criterion` — the function registers its own.
52    Criterion,
53}
54
55/// One `benches/<name>.rs` file and its `#[yog_api::bench]` functions.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub struct BenchFile {
58    pub(crate) name: String,
59    pub(crate) functions: Vec<(String, Takes)>,
60}
61
62/// A problem with a benchmark file.
63fn problem(message: impl std::fmt::Display) -> Diagnostic {
64    Diagnostic::error(yog_ids::build_problem::BENCH, message)
65}
66
67/// The mod's benchmark files, by name; none without a `benches/` folder.
68pub 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
96/// The `#[yog_api::bench]` functions of a benchmark file, in source order.
97pub(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
135/// The generated `main` of one benchmark file.
136pub(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
167/// Writes every benchmark file's `main`, dropping those of removed files.
168pub 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
188/// The generated Cargo.toml's `[dev-dependencies]` — `test_deps` (the
189/// mod's tests' own), and with benchmarks theirs — and `[[bench]]` targets.
190pub(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}