add adaptive model examples using fenwick trees (#16)

Author: danieleades

Cargo.toml

@@ -13,4 +13,5 @@ bitstream-io = "1.2.0"
 thiserror = "1.0.30"
 
 [dev-dependencies]
+fenwick = { version = "1.0.0" }
 test-case = "2.0.0"

examples/common/mod.rs

@@ -47,3 +47,30 @@ where
     }
     output
 }
+
+#[allow(unused)]
+pub fn round_trip_string<M>(model: M, input: String)
+where
+    M: Model<Symbol = char> + Clone,
+{
+    let input_bytes = input.bytes().len();
+
+    let buffer = encode(model.clone(), input.chars());
+
+    let output_bytes = buffer.len();
+
+    println!("input bytes: {}", input_bytes);
+    println!("output bytes: {}", output_bytes);
+
+    println!(
+        "compression ratio: {}",
+        input_bytes as f32 / output_bytes as f32
+    );
+
+    let output = decode(model, &buffer);
+
+    let mut prefix: String = output.into_iter().take(299).collect();
+    prefix.push_str("...");
+
+    println!("{}", prefix);
+}

examples/fenwick/context_switching.rs

@@ -0,0 +1,87 @@
+#![allow(missing_docs, unused)]
+//! Fenwick tree based context-switching model
+
+use arithmetic_coding::Model;
+
+use super::Weights;
+
+#[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 {
+        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,
+        }
+    }
+
+    fn context(&self) -> &Weights {
+        &self.contexts[self.current_context]
+    }
+
+    fn context_mut(&mut self) -> &mut Weights {
+        &mut self.contexts[self.current_context]
+    }
+}
+
+#[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> {
+        Ok(self.context().range(symbol.copied()))
+    }
+
+    fn max_denominator(&self) -> Self::B {
+        self.max_denominator
+    }
+
+    fn symbol(&self, value: Self::B) -> Option<Self::Symbol> {
+        self.context().symbol(value)
+    }
+
+    fn denominator(&self) -> Self::B {
+        self.context().total
+    }
+
+    fn update(&mut self, symbol: Option<&Self::Symbol>) {
+        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();
+    }
+}

examples/fenwick/mod.rs

@@ -0,0 +1,77 @@
+//! [`Models`](crate::Model) implemented using Fenwick trees
+
+use std::ops::Range;
+
+pub mod context_switching;
+pub mod simple;
+
+/// A wrapper around a vector of fenwick counts, with one additional weight for
+/// EOF.
+#[derive(Debug, Clone)]
+struct Weights {
+    fenwick_counts: Vec<u64>,
+    total: u64,
+}
+
+impl Weights {
+    fn new(n: usize) -> Self {
+        // we add one extra value here to account for the EOF
+        let mut fenwick_counts = vec![0; n + 1];
+
+        for i in 0..fenwick_counts.len() {
+            fenwick::array::update(&mut fenwick_counts, i, 1);
+        }
+
+        let total = fenwick_counts.len() as u64;
+        Self {
+            fenwick_counts,
+            total,
+        }
+    }
+
+    fn update(&mut self, i: Option<usize>, delta: u64) {
+        let index = i.map(|i| i + 1).unwrap_or_default();
+        fenwick::array::update(&mut self.fenwick_counts, index, delta);
+        self.total += delta;
+    }
+
+    fn prefix_sum(&self, i: Option<usize>) -> u64 {
+        let index = i.map(|i| i + 1).unwrap_or_default();
+        fenwick::array::prefix_sum(&self.fenwick_counts, index)
+    }
+
+    fn range(&self, i: Option<usize>) -> Range<u64> {
+        let index = i.map(|i| i + 1).unwrap_or_default();
+
+        let upper = fenwick::array::prefix_sum(&self.fenwick_counts, index);
+
+        let lower = if index == 0 {
+            0
+        } else {
+            fenwick::array::prefix_sum(&self.fenwick_counts, index - 1)
+        };
+        lower..upper
+    }
+
+    fn len(&self) -> usize {
+        self.fenwick_counts.len() - 1
+    }
+
+    fn symbol(&self, prefix_sum: u64) -> Option<usize> {
+        if prefix_sum < self.prefix_sum(None) {
+            return None;
+        }
+
+        for i in 0..self.len() {
+            if prefix_sum < self.prefix_sum(Some(i)) {
+                return Some(i);
+            }
+        }
+
+        unreachable!()
+    }
+
+    fn total(&self) -> u64 {
+        self.total
+    }
+}

examples/fenwick/simple.rs

@@ -0,0 +1,71 @@
+#![allow(missing_docs, unused)]
+//! simple adaptive model using a fenwick tree
+
+use arithmetic_coding::Model;
+
+use super::Weights;
+
+#[derive(Debug, Clone)]
+pub struct FenwickModel {
+    weights: Weights,
+    max_denominator: u64,
+}
+
+impl FenwickModel {
+    #[must_use]
+    pub fn with_symbols(symbols: usize) -> Self {
+        let weights = Weights::new(symbols);
+
+        Self {
+            weights,
+            max_denominator: 1 << 17,
+        }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("invalid symbol received: {0}")]
+pub struct ValueError(pub 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> {
+        if let Some(s) = symbol.copied() {
+            if s >= self.weights.len() {
+                Err(ValueError(s))
+            } else {
+                Ok(self.weights.range(Some(s)))
+            }
+        } else {
+            Ok(self.weights.range(None))
+        }
+    }
+
+    fn max_denominator(&self) -> Self::B {
+        self.max_denominator
+    }
+
+    fn symbol(&self, value: Self::B) -> Option<Self::Symbol> {
+        self.weights.symbol(value)
+    }
+
+    fn denominator(&self) -> Self::B {
+        self.weights.total()
+    }
+
+    fn update(&mut self, symbol: Option<&Self::Symbol>) {
+        debug_assert!(
+            self.denominator() < self.max_denominator,
+            "hit max denominator!"
+        );
+        if self.denominator() < self.max_denominator {
+            self.weights.update(symbol.copied(), 1);
+        }
+    }
+}

examples/fenwick_adaptive.rs

@@ -0,0 +1,76 @@
+use std::{fs::File, io::Read, ops::Range};
+
+use arithmetic_coding::Model;
+
+mod common;
+mod fenwick;
+
+use self::fenwick::simple::{FenwickModel, ValueError};
+
+const ALPHABET: &str =
+    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,\n-':()[]#*;\"!?*&é/àâè%@$";
+
+#[derive(Debug, Clone)]
+pub struct StringModel {
+    alphabet: Vec<char>,
+    fenwick_model: FenwickModel,
+}
+
+impl StringModel {
+    #[must_use]
+    pub fn new(alphabet: Vec<char>) -> Self {
+        let fenwick_model = FenwickModel::with_symbols(alphabet.len());
+        Self {
+            alphabet,
+            fenwick_model,
+        }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("invalid character: {0}")]
+pub struct Error(char);
+
+impl Model for StringModel {
+    type B = u64;
+    type Symbol = char;
+    type ValueError = ValueError;
+
+    fn probability(
+        &self,
+        symbol: Option<&Self::Symbol>,
+    ) -> Result<Range<Self::B>, Self::ValueError> {
+        let fenwick_symbol = symbol.map(|c| self.alphabet.iter().position(|x| x == c).unwrap());
+        self.fenwick_model.probability(fenwick_symbol.as_ref())
+    }
+
+    fn symbol(&self, value: Self::B) -> Option<Self::Symbol> {
+        let index = self.fenwick_model.symbol(value)?;
+        self.alphabet.get(index).copied()
+    }
+
+    fn max_denominator(&self) -> Self::B {
+        self.fenwick_model.max_denominator()
+    }
+
+    fn denominator(&self) -> Self::B {
+        self.fenwick_model.denominator()
+    }
+
+    fn update(&mut self, symbol: Option<&Self::Symbol>) {
+        let fenwick_symbol = symbol.map(|c| self.alphabet.iter().position(|x| x == c).unwrap());
+        self.fenwick_model.update(fenwick_symbol.as_ref())
+    }
+}
+
+fn main() {
+    let model = StringModel::new(ALPHABET.chars().collect());
+
+    let mut input = String::new();
+    File::open("./resources/sherlock.txt")
+        .unwrap()
+        .read_to_string(&mut input)
+        .unwrap();
+
+    common::round_trip_string(model, input);
+}