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

yog_cli/
bench.rs
