arcweight 0.3.0

A high-performance, modular library for weighted finite state transducers with comprehensive examples and benchmarks
Documentation
use arcweight::prelude::*;
use criterion::{criterion_group, criterion_main, Criterion};
#[cfg(feature = "parallel")]
use rayon::prelude::*;
use std::hint::black_box;
use std::time::Duration;

fn create_large_fst(n: usize) -> VectorFst<TropicalWeight> {
    let mut fst = VectorFst::new();
    let mut states = Vec::new();

    for _ in 0..=n {
        states.push(fst.add_state());
    }

    fst.set_start(states[0]);
    fst.set_final(states[n], TropicalWeight::one());

    for i in 0..n {
        for j in 1..=3 {
            fst.add_arc(
                states[i],
                Arc::new(
                    j as u32,
                    j as u32,
                    TropicalWeight::new(j as f32),
                    states[i + 1],
                ),
            );
        }
    }

    fst
}

pub fn bench_arc_count(c: &mut Criterion) {
    let mut group = c.benchmark_group("arc_count");
    group.measurement_time(Duration::from_secs(10));

    let fst = create_large_fst(1000);

    group.bench_function("arc_count_sequential", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let count = fst.num_arcs_total();
            black_box(count);
        })
    });

    #[cfg(feature = "parallel")]
    group.bench_function("arc_count_parallel", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let states: Vec<_> = fst.states().collect();
            let count: usize = states.par_iter().map(|state| fst.num_arcs(*state)).sum();
            black_box(count);
        })
    });
    #[cfg(not(feature = "parallel"))]
    group.bench_function("arc_count_parallel", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let states: Vec<_> = fst.states().collect();
            let count: usize = states.iter().map(|state| fst.num_arcs(*state)).sum();
            black_box(count);
        })
    });

    group.finish();
}

pub fn bench_arc_iteration(c: &mut Criterion) {
    let mut group = c.benchmark_group("arc_iteration");
    group.measurement_time(Duration::from_secs(10));

    let fst = create_large_fst(1000);

    group.bench_function("arc_iter_sequential", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let mut count = 0;
            for state in fst.states() {
                for _arc in fst.arcs(state) {
                    count += 1;
                }
            }
            black_box(count);
        })
    });

    #[cfg(feature = "parallel")]
    group.bench_function("arc_iter_parallel", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let states: Vec<_> = fst.states().collect();
            let count: usize = states
                .par_iter()
                .map(|state| fst.arcs(*state).count())
                .sum();
            black_box(count);
        })
    });
    #[cfg(not(feature = "parallel"))]
    group.bench_function("arc_iter_parallel", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let states: Vec<_> = fst.states().collect();
            let count: usize = states.iter().map(|state| fst.arcs(*state).count()).sum();
            black_box(count);
        })
    });

    group.finish();
}

pub fn bench_arc_weight_sum(c: &mut Criterion) {
    let mut group = c.benchmark_group("arc_weight_sum");
    group.measurement_time(Duration::from_secs(10));

    let fst = create_large_fst(1000);

    group.bench_function("weight_sum_sequential", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let mut sum = 0.0f32;
            for state in fst.states() {
                for arc in fst.arcs(state) {
                    sum += arc.weight.value();
                }
            }
            black_box(sum);
        })
    });

    #[cfg(feature = "parallel")]
    group.bench_function("weight_sum_parallel", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let states: Vec<_> = fst.states().collect();
            let sum: f32 = states
                .par_iter()
                .flat_map(|state| fst.arcs(*state).collect::<Vec<_>>())
                .map(|arc| *arc.weight.value())
                .sum();
            black_box(sum);
        })
    });
    #[cfg(not(feature = "parallel"))]
    group.bench_function("weight_sum_parallel", |b| {
        b.iter(|| {
            let fst = black_box(&fst);
            let states: Vec<_> = fst.states().collect();
            let sum: f32 = states
                .iter()
                .flat_map(|state| fst.arcs(*state).collect::<Vec<_>>())
                .map(|arc| *arc.weight.value())
                .sum();
            black_box(sum);
        })
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_arc_count,
    bench_arc_iteration,
    bench_arc_weight_sum
);
criterion_main!(benches);