use arcweight::prelude::*;
use criterion::{criterion_group, criterion_main, Criterion};
use std::hint::black_box;
fn create_linear_fst(size: usize) -> VectorFst<TropicalWeight> {
let mut fst = VectorFst::new();
if size == 0 {
return fst;
}
let states: Vec<_> = (0..size).map(|_| fst.add_state()).collect();
fst.set_start(states[0]);
fst.set_final(states[size - 1], TropicalWeight::one());
for i in 0..size - 1 {
fst.add_arc(
states[i],
Arc::new(
(i % 26 + 97) as u32,
(i % 26 + 97) as u32,
TropicalWeight::new(1.0),
states[i + 1],
),
);
}
fst
}
fn create_branching_fst(states: usize, branches: usize) -> VectorFst<TropicalWeight> {
let mut fst = VectorFst::new();
if states == 0 {
return fst;
}
let state_ids: Vec<_> = (0..states).map(|_| fst.add_state()).collect();
fst.set_start(state_ids[0]);
fst.set_final(state_ids[states - 1], TropicalWeight::one());
for i in 0..states - 1 {
for j in 0..branches.min(states - i - 1) {
let target_state = state_ids[i + j + 1];
fst.add_arc(
state_ids[i],
Arc::new(
((i + j) % 26 + 97) as u32,
((i + j) % 26 + 97) as u32,
TropicalWeight::new((j + 1) as f32),
target_state,
),
);
}
}
fst
}
fn bench_state_sort_linear(c: &mut Criterion) {
let mut group = c.benchmark_group("state_sort_linear");
for size in [10, 100, 1000, 5000].iter() {
let fst = create_linear_fst(*size);
group.bench_function(format!("bfs_{size}"), |b| {
b.iter(|| {
let mut fst_clone = black_box(&fst).clone();
state_sort(&mut fst_clone, StateSortType::Bfs).unwrap();
black_box(fst_clone)
})
});
group.bench_function(format!("dfs_{size}"), |b| {
b.iter(|| {
let mut fst_clone = black_box(&fst).clone();
state_sort(&mut fst_clone, StateSortType::Dfs).unwrap();
black_box(fst_clone)
})
});
}
group.finish();
}
fn bench_arc_sort_linear(c: &mut Criterion) {
let mut group = c.benchmark_group("arc_sort_linear");
for size in [10, 100, 1000, 5000].iter() {
let fst = create_linear_fst(*size);
group.bench_function(format!("ilabel_{size}"), |b| {
b.iter(|| {
let mut fst_clone = black_box(&fst).clone();
arc_sort(&mut fst_clone, ArcSortType::Ilabel).unwrap();
black_box(fst_clone)
})
});
group.bench_function(format!("olabel_{size}"), |b| {
b.iter(|| {
let mut fst_clone = black_box(&fst).clone();
arc_sort(&mut fst_clone, ArcSortType::Olabel).unwrap();
black_box(fst_clone)
})
});
}
group.finish();
}
fn bench_arc_sort_branching(c: &mut Criterion) {
let mut group = c.benchmark_group("arc_sort_branching");
for size in [10, 100, 500, 1000].iter() {
let fst = create_branching_fst(*size, 3);
group.bench_function(format!("ilabel_{size}"), |b| {
b.iter(|| {
let mut fst_clone = black_box(&fst).clone();
arc_sort(&mut fst_clone, ArcSortType::Ilabel).unwrap();
black_box(fst_clone)
})
});
}
group.finish();
}
criterion_group!(
benches,
bench_state_sort_linear,
bench_arc_sort_linear,
bench_arc_sort_branching
);
criterion_main!(benches);