use std::io::Cursor;
use std::sync::OnceLock;
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use tempfile::TempDir;
use linesmith::data_context::DataDep;
use linesmith::theme::{self, Capability};
use linesmith::{
run_with_context,
segments::{Segment, BUILT_IN_SEGMENT_IDS, DEFAULT_SEGMENT_IDS},
RunContext,
};
const MINIMAL_PAYLOAD: &[u8] = include_bytes!("../tests/fixtures/claude_minimal.json");
const WORKTREE_PAYLOAD: &[u8] = include_bytes!("../tests/fixtures/claude_worktree.json");
const TERMINAL_WIDTH: u16 = 200;
const MINIMAL_IDS: &[&str] = &["model", "workspace"];
fn all_no_usage_ids() -> Vec<&'static str> {
BUILT_IN_SEGMENT_IDS
.iter()
.copied()
.filter(|id| !id.starts_with("rate_limit_") && *id != "extra_usage")
.collect()
}
fn fixture_repo() -> &'static std::path::Path {
static REPO: OnceLock<TempDir> = OnceLock::new();
REPO.get_or_init(|| {
let tmp = TempDir::new().expect("bench setup: TempDir");
gix::init(tmp.path()).expect("bench setup: gix init");
tmp
})
.path()
}
fn build_named_segments(ids: &[&str], name: &str, forbid_usage: bool) -> Vec<Box<dyn Segment>> {
let cfg = build_config(ids, name);
let mut warnings: Vec<String> = Vec::new();
let segs = linesmith::build_segments(Some(&cfg), None, |w| warnings.push(w.to_string()));
if !warnings.is_empty() {
panic!("bench '{name}' build_segments emitted warnings: {warnings:?}");
}
if segs.len() != ids.len() {
panic!(
"bench '{name}' built {} segments but requested {}: ids={:?}",
segs.len(),
ids.len(),
ids
);
}
if forbid_usage {
for (id, seg) in ids.iter().zip(segs.iter()) {
if seg.data_deps().contains(&DataDep::Usage) {
panic!(
"bench '{name}' includes segment '{id}' which declares DataDep::Usage; \
update all_no_usage_ids' filter to exclude it"
);
}
}
}
segs
}
fn build_config(ids: &[&str], name: &str) -> linesmith::config::Config {
let toml_src = format!(
"[line]\nsegments = [{}]\n",
ids.iter()
.map(|id| format!("\"{id}\""))
.collect::<Vec<_>>()
.join(", ")
);
toml_src
.parse()
.unwrap_or_else(|e| panic!("bench '{name}' synthetic config did not parse: {e}"))
}
fn run_context(cwd: std::path::PathBuf) -> RunContext<'static> {
RunContext::new(
theme::default_theme(),
Capability::None,
TERMINAL_WIDTH,
Some(cwd),
false,
)
}
fn render_once(
segments: &[Box<dyn Segment>],
payload: &[u8],
ctx: &RunContext<'_>,
sink: &mut Vec<u8>,
name: &str,
) {
sink.clear();
let mut stderr_sink = std::io::sink();
run_with_context(
Cursor::new(payload),
&mut *sink,
&mut stderr_sink,
segments,
ctx,
)
.unwrap_or_else(|e| {
panic!(
"bench '{name}' render failed over {} bytes: {e}",
payload.len()
)
});
criterion::black_box(&sink[..]);
}
fn cold_bench(
c: &mut Criterion,
name: &'static str,
ids: &[&str],
payload: &'static [u8],
forbid_usage: bool,
) {
let cwd = fixture_repo().to_path_buf();
c.bench_function(name, |b| {
b.iter_batched_ref(
|| {
(
build_named_segments(ids, name, forbid_usage),
Vec::with_capacity(4096),
)
},
|(segs, sink)| render_once(segs, payload, &run_context(cwd.clone()), sink, name),
BatchSize::PerIteration,
);
});
}
fn warm_bench(
c: &mut Criterion,
name: &'static str,
ids: &[&str],
payload: &'static [u8],
forbid_usage: bool,
) {
let cwd = fixture_repo().to_path_buf();
let segs = build_named_segments(ids, name, forbid_usage);
let ctx = run_context(cwd);
let mut sink = Vec::with_capacity(4096);
c.bench_function(name, |b| {
b.iter(|| render_once(&segs, payload, &ctx, &mut sink, name));
});
}
fn bench_minimal_cold(c: &mut Criterion) {
cold_bench(
c,
"render_minimal_cold",
MINIMAL_IDS,
MINIMAL_PAYLOAD,
false,
);
}
fn bench_minimal_warm(c: &mut Criterion) {
warm_bench(
c,
"render_minimal_warm",
MINIMAL_IDS,
MINIMAL_PAYLOAD,
false,
);
}
fn bench_default_cold(c: &mut Criterion) {
cold_bench(
c,
"render_default_cold",
DEFAULT_SEGMENT_IDS,
WORKTREE_PAYLOAD,
false,
);
}
fn bench_default_warm(c: &mut Criterion) {
warm_bench(
c,
"render_default_warm",
DEFAULT_SEGMENT_IDS,
WORKTREE_PAYLOAD,
false,
);
}
fn bench_all_no_usage_cold(c: &mut Criterion) {
let ids = all_no_usage_ids();
cold_bench(c, "render_all_no_usage_cold", &ids, WORKTREE_PAYLOAD, true);
}
fn bench_all_no_usage_warm(c: &mut Criterion) {
let ids = all_no_usage_ids();
warm_bench(c, "render_all_no_usage_warm", &ids, WORKTREE_PAYLOAD, true);
}
criterion_group!(
render,
bench_minimal_cold,
bench_minimal_warm,
bench_default_cold,
bench_default_warm,
bench_all_no_usage_cold,
bench_all_no_usage_warm,
);
criterion_main!(render);