use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use context_forge::analysis::{
build_session_term_maps, classify_passages, compute_recurrence, score_passages,
ClassificationConfig, Lexicons, PassageContext, PrefilterConfig, RecurrenceConfig,
ScoringConfig, Tokenizer, TokenizerConfig,
};
const TRANSCRIPT_SIZES: &[(&str, usize)] = &[
("10kb", 10 * 1024),
("100kb", 100 * 1024),
("1mb", 1024 * 1024),
];
const SESSION_COUNTS: &[usize] = &[5, 50, 500];
const SENTENCE_TEMPLATES: &[&str] = &[
"We switched from Redis to Memcached because latency dropped significantly.",
"Cache mode was set to writeback after the incident review.",
"We should not enable verbose logging in production environments.",
"The deploy pipeline timeout changed to 30 seconds yesterday.",
"Primary database is now PostgreSQL 16 after the migration completed.",
"Yes always run cargo test before committing any changes.",
"General context discussion about implementation details and tradeoffs.",
"Confirmed the worker pool causes memory growth under sustained load.",
"The retry policy was updated to use exponential backoff with jitter.",
"Server IP changed to 10.0.0.2 after the network reconfiguration.",
];
fn build_session_content(target_bytes: usize, session_index: usize) -> String {
let mut content = String::with_capacity(target_bytes + 256);
let mut template_index = session_index % SENTENCE_TEMPLATES.len();
while content.len() < target_bytes {
content.push_str(SENTENCE_TEMPLATES[template_index]);
content.push(' ');
template_index = (template_index + 1) % SENTENCE_TEMPLATES.len();
}
content
}
fn build_corpus(transcript_bytes: usize, session_count: usize) -> Vec<String> {
let per_session = (transcript_bytes / session_count).max(64);
(0..session_count)
.map(|index| build_session_content(per_session, index))
.collect()
}
fn bench_build_session_term_maps(c: &mut Criterion) {
let tokenizer = Tokenizer::new(&TokenizerConfig::default());
let prefilter_config = PrefilterConfig::default();
let mut group = c.benchmark_group("build_session_term_maps");
for &(size_label, size_bytes) in TRANSCRIPT_SIZES {
for &session_count in SESSION_COUNTS {
let corpus = build_corpus(size_bytes, session_count);
let session_contents: Vec<Vec<&str>> = corpus
.iter()
.map(|content| vec![content.as_str()])
.collect();
let id = BenchmarkId::from_parameter(format!("{size_label}/{session_count}sess"));
group.bench_with_input(id, &session_contents, |b, session_contents| {
b.iter(|| build_session_term_maps(session_contents, &tokenizer, &prefilter_config));
});
}
}
group.finish();
}
fn build_passages(transcript_bytes: usize, session_count: usize) -> Vec<PassageContext> {
let per_session = (transcript_bytes / session_count).max(64);
let mut passages = Vec::new();
for session_index in 0..session_count {
let mut written = 0usize;
let mut template_index = session_index % SENTENCE_TEMPLATES.len();
let mut local_timestamp = 1_700_000_000_i64 + session_index as i64;
while written < per_session {
let text = SENTENCE_TEMPLATES[template_index];
written += text.len();
passages.push(PassageContext {
passage_text: text.to_string(),
triggering_terms: vec!["cache".to_string(), "redis".to_string()],
session_id: format!("session-{session_index}"),
timestamp: local_timestamp,
});
template_index = (template_index + 1) % SENTENCE_TEMPLATES.len();
local_timestamp += 1;
}
}
passages
}
fn bench_classify_passages(c: &mut Criterion) {
let lexicons = Lexicons::default();
let config = ClassificationConfig::default();
let mut group = c.benchmark_group("classify_passages");
for &(size_label, size_bytes) in TRANSCRIPT_SIZES {
for &session_count in SESSION_COUNTS {
let passages = build_passages(size_bytes, session_count);
let id = BenchmarkId::from_parameter(format!("{size_label}/{session_count}sess"));
group.bench_with_input(id, &passages, |b, passages| {
b.iter(|| classify_passages(passages, &lexicons, &config));
});
}
}
group.finish();
}
fn bench_score_passages(c: &mut Criterion) {
let lexicons = Lexicons::default();
let classify_config = ClassificationConfig::default();
let scoring_config = ScoringConfig::default();
let recurrence_config = RecurrenceConfig::default();
let tokenizer = Tokenizer::new(&TokenizerConfig::default());
let prefilter_config = PrefilterConfig::default();
let now_timestamp = 1_700_100_000_i64;
let mut group = c.benchmark_group("score_passages");
for &(size_label, size_bytes) in TRANSCRIPT_SIZES {
for &session_count in SESSION_COUNTS {
let passages = build_passages(size_bytes, session_count);
let classified = classify_passages(&passages, &lexicons, &classify_config);
let corpus = build_corpus(size_bytes, session_count);
let session_contents: Vec<Vec<&str>> = corpus
.iter()
.map(|content| vec![content.as_str()])
.collect();
let session_term_maps =
build_session_term_maps(&session_contents, &tokenizer, &prefilter_config);
let recurrence = compute_recurrence(&session_term_maps, &recurrence_config);
let recurrence_map = recurrence
.into_iter()
.map(|result| (result.term.clone(), result))
.collect();
let id = BenchmarkId::from_parameter(format!("{size_label}/{session_count}sess"));
group.bench_with_input(id, &classified, |b, classified| {
b.iter(|| {
score_passages(classified, &recurrence_map, &scoring_config, now_timestamp)
});
});
}
}
group.finish();
}
criterion_group!(
benches,
bench_build_session_term_maps,
bench_classify_passages,
bench_score_passages
);
criterion_main!(benches);