Add bench test (#20)

Author: danieleades

Cargo.toml

@@ -21,5 +21,9 @@ bitstream-io = "1.2.0"
 thiserror = "1.0.30"
 
 [dev-dependencies]
-test-case = "2.0.0"
 fenwick-model = { path = "./fenwick-model" }
+criterion = "0.3.5"
+
+[[bench]]
+name = "sherlock"
+harness = false

benches/common/mod.rs

@@ -0,0 +1,45 @@
+use arithmetic_coding::{Decoder, Encoder, Model};
+use bitstream_io::{BigEndian, BitReader, BitWrite, BitWriter};
+
+pub fn round_trip<M>(model: M, input: &[M::Symbol])
+where
+    M: Model + Clone,
+    M::Symbol: Copy + std::fmt::Debug + PartialEq,
+{
+    let buffer = encode(model.clone(), input.iter().copied());
+
+    let mut output = Vec::with_capacity(input.len());
+    for symbol in decode(model, &buffer) {
+        output.push(symbol);
+    }
+
+    assert_eq!(input, output.as_slice());
+}
+
+pub fn encode<M, I>(model: M, input: I) -> Vec<u8>
+where
+    M: Model,
+    I: IntoIterator<Item = M::Symbol>,
+{
+    let mut bitwriter = BitWriter::endian(Vec::new(), BigEndian);
+    let mut encoder = Encoder::<M>::new(model);
+
+    encoder.encode_all(input, &mut bitwriter).unwrap();
+    bitwriter.byte_align().unwrap();
+
+    bitwriter.into_writer()
+}
+
+pub fn decode<M>(model: M, buffer: &[u8]) -> Vec<M::Symbol>
+where
+    M: Model,
+{
+    let bitreader = BitReader::endian(buffer, BigEndian);
+    let mut decoder = Decoder::new(model, bitreader).unwrap();
+    let mut output = Vec::new();
+
+    while let Some(symbol) = decoder.decode_symbol().unwrap() {
+        output.push(symbol);
+    }
+    output
+}

benches/sherlock.rs

@@ -0,0 +1,73 @@
+use std::{fs::File, io::Read, ops::Range};
+
+use arithmetic_coding::Model;
+use fenwick_model::{simple::FenwickModel, ValueError};
+
+mod common;
+
+#[derive(Debug, Clone)]
+pub struct StringModel {
+    fenwick_model: FenwickModel,
+}
+
+impl StringModel {
+    #[must_use]
+    pub fn new(symbols: usize) -> Self {
+        let fenwick_model = FenwickModel::with_symbols(symbols, 1 << 20);
+        Self { fenwick_model }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("invalid character: {0}")]
+pub struct Error(char);
+
+impl Model for StringModel {
+    type B = u64;
+    type Symbol = u8;
+    type ValueError = ValueError;
+
+    fn probability(
+        &self,
+        symbol: Option<&Self::Symbol>,
+    ) -> Result<Range<Self::B>, Self::ValueError> {
+        let fenwick_symbol = symbol.map(|c| *c as usize);
+        self.fenwick_model.probability(fenwick_symbol.as_ref())
+    }
+
+    fn symbol(&self, value: Self::B) -> Option<Self::Symbol> {
+        self.fenwick_model.symbol(value).map(|x| x as u8)
+    }
+
+    fn max_denominator(&self) -> Self::B {
+        self.fenwick_model.max_denominator()
+    }
+
+    fn denominator(&self) -> Self::B {
+        self.fenwick_model.denominator()
+    }
+}
+
+fn round_trip(input: &[u8]) {
+    let model = StringModel::new(256);
+
+    common::round_trip(model, input);
+}
+
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+pub fn criterion_benchmark(c: &mut Criterion) {
+    let mut input_string = String::new();
+    File::open("./resources/sherlock.txt")
+        .unwrap()
+        .read_to_string(&mut input_string)
+        .unwrap();
+
+    let truncated: String = input_string.chars().take(3428).collect();
+    let input = truncated.as_bytes();
+
+    c.bench_function("round trip", |b| b.iter(|| round_trip(black_box(input))));
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

examples/fenwick_adaptive.rs

@@ -4,7 +4,7 @@ use arithmetic_coding::Model;
 
 mod common;
 
-use fenwick_model::simple::{FenwickModel, ValueError};
+use fenwick_model::{simple::FenwickModel, ValueError};
 
 const ALPHABET: &str =
     "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,\n-':()[]#*;\"!?*&é/àâè%@$";
@@ -18,7 +18,7 @@ pub struct StringModel {
 impl StringModel {
     #[must_use]
     pub fn new(alphabet: Vec<char>) -> Self {
-        let fenwick_model = FenwickModel::with_symbols(alphabet.len());
+        let fenwick_model = FenwickModel::with_symbols(alphabet.len(), 1 << 20);
         Self {
             alphabet,
             fenwick_model,

examples/fenwick_context_switching.rs

@@ -4,7 +4,7 @@ use arithmetic_coding::Model;
 
 mod common;
 
-use fenwick_model::context_switching::{FenwickModel, ValueError};
+use fenwick_model::{context_switching::FenwickModel, ValueError};
 
 const ALPHABET: &str =
     "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,\n-':()[]#*;\"!?*&é/àâè%@$";
@@ -18,7 +18,7 @@ pub struct StringModel {
 impl StringModel {
     #[must_use]
     pub fn new(alphabet: Vec<char>) -> Self {
-        let fenwick_model = FenwickModel::with_symbols(alphabet.len());
+        let fenwick_model = FenwickModel::with_symbols(alphabet.len(), 1 << 17);
         Self {
             alphabet,
             fenwick_model,

fenwick-model/src/context_switching.rs

@@ -1,37 +1,29 @@
-#![allow(missing_docs, unused)]
 //! Fenwick tree based context-switching model
 
 use arithmetic_coding_core::Model;
 
 use super::Weights;
+use crate::ValueError;
 
 #[derive(Debug, Clone)]
 pub struct FenwickModel {
     contexts: Vec<Weights>,
-    previous_context: usize,
     current_context: usize,
-    denominator: u64,
     max_denominator: u64,
 }
 
 impl FenwickModel {
     #[must_use]
-    pub fn with_symbols(symbols: usize) -> Self {
+    pub fn with_symbols(symbols: usize, max_denominator: u64) -> Self {
         let mut contexts = Vec::with_capacity(symbols + 1);
-        let mut denominator = 0;
-        let max_denominator = 1 << 17;
 
         for _ in 0..=symbols {
-            let weight = Weights::new(symbols);
-            denominator = denominator.max(weight.total());
             contexts.push(Weights::new(symbols));
         }
 
         Self {
             contexts,
-            previous_context: 1,
             current_context: 1,
-            denominator,
             max_denominator,
         }
     }
@@ -45,42 +37,34 @@ impl FenwickModel {
     }
 }
 
-#[derive(Debug, thiserror::Error)]
-#[error("invalid symbol received: {0}")]
-pub struct ValueError(usize);
-
 impl Model for FenwickModel {
     type B = u64;
     type Symbol = usize;
     type ValueError = ValueError;
 
-    fn probability(
-        &self,
-        symbol: Option<&Self::Symbol>,
-    ) -> Result<std::ops::Range<Self::B>, Self::ValueError> {
+    fn probability(&self, symbol: Option<&usize>) -> Result<std::ops::Range<u64>, ValueError> {
         Ok(self.context().range(symbol.copied()))
     }
 
-    fn max_denominator(&self) -> Self::B {
-        self.max_denominator
+    fn denominator(&self) -> u64 {
+        self.context().total
     }
 
-    fn symbol(&self, value: Self::B) -> Option<Self::Symbol> {
-        self.context().symbol(value)
+    fn max_denominator(&self) -> u64 {
+        self.max_denominator
     }
 
-    fn denominator(&self) -> Self::B {
-        self.context().total
+    fn symbol(&self, value: u64) -> Option<usize> {
+        self.context().symbol(value)
     }
 
-    fn update(&mut self, symbol: Option<&Self::Symbol>) {
+    fn update(&mut self, symbol: Option<&usize>) {
         debug_assert!(
             self.denominator() < self.max_denominator,
             "hit max denominator!"
         );
         if self.denominator() < self.max_denominator {
             self.context_mut().update(symbol.copied(), 1);
-            self.denominator = self.denominator.max(self.context().total());
         }
         self.current_context = symbol.map(|x| x + 1).unwrap_or_default();
     }

fenwick-model/src/lib.rs

@@ -75,3 +75,7 @@ impl Weights {
         self.total
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+#[error("invalid symbol received: {0}")]
+pub struct ValueError(pub usize);

fenwick-model/src/simple.rs

@@ -4,6 +4,7 @@
 use arithmetic_coding_core::Model;
 
 use super::Weights;
+use crate::ValueError;
 
 #[derive(Debug, Clone)]
 pub struct FenwickModel {
@@ -13,20 +14,16 @@ pub struct FenwickModel {
 
 impl FenwickModel {
     #[must_use]
-    pub fn with_symbols(symbols: usize) -> Self {
+    pub fn with_symbols(symbols: usize, max_denominator: u64) -> Self {
         let weights = Weights::new(symbols);
 
         Self {
             weights,
-            max_denominator: 1 << 17,
+            max_denominator,
         }
     }
 }
 
-#[derive(Debug, thiserror::Error)]
-#[error("invalid symbol received: {0}")]
-pub struct ValueError(pub usize);
-
 impl Model for FenwickModel {
     type B = u64;
     type Symbol = usize;