feat: assert the chosen precision is valid when calling 'chain' in debug builds (#80)

Co-authored-by: ⭐️NINIKA⭐️ <[email protected]>

Author: danieleades

src/common.rs

@@ -2,6 +2,8 @@ use std::ops::Range;
 
 use arithmetic_coding_core::BitStore;
 
+use crate::Model;
+
 #[derive(Debug)]
 pub struct State<B: BitStore> {
     pub precision: u32,
@@ -43,3 +45,15 @@ where
         self.low += (range * p.start) / denominator;
     }
 }
+
+pub fn assert_precision_sufficient<M: Model>(max_denominator: M::B, precision: u32) {
+    let frequency_bits = max_denominator.log2() + 1;
+    assert!(
+        (precision >= (frequency_bits + 2)),
+        "not enough bits of precision to prevent overflow/underflow",
+    );
+    assert!(
+        (frequency_bits + precision) <= M::B::BITS,
+        "not enough bits in BitStore to support the required precision",
+    );
+}

src/decoder.rs

@@ -4,7 +4,10 @@ use std::{io, ops::Range};
 
 use bitstream_io::BitRead;
 
-use crate::{common, BitStore, Model};
+use crate::{
+    common::{self, assert_precision_sufficient},
+    BitStore, Model,
+};
 
 // this algorithm is derived from this article - https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html
 
@@ -79,23 +82,15 @@ where
     /// If these constraints cannot be satisfied this method will panic in debug
     /// builds
     pub fn with_precision(model: M, input: R, precision: u32) -> Self {
-        let frequency_bits = model.max_denominator().log2() + 1;
-        debug_assert!(
-            (precision >= (frequency_bits + 2)),
-            "not enough bits of precision to prevent overflow/underflow",
-        );
-        debug_assert!(
-            (frequency_bits + precision) <= M::B::BITS,
-            "not enough bits in BitStore to support the required precision",
-        );
-
         let state = State::new(precision, input);
-
-        Self { model, state }
+        Self::with_state(state, model)
     }
 
     /// todo
-    pub const fn with_state(state: State<M::B, R>, model: M) -> Self {
+    pub fn with_state(state: State<M::B, R>, model: M) -> Self {
+        #[cfg(debug_assertions)]
+        assert_precision_sufficient::<M>(model.max_denominator(), state.state.precision);
+
         Self { model, state }
     }
 
@@ -144,10 +139,7 @@ where
     where
         X: Model<B = M::B>,
     {
-        Decoder {
-            model,
-            state: self.state,
-        }
+        Decoder::with_state(self.state, model)
     }
 
     /// todo

src/encoder.rs

@@ -4,7 +4,10 @@ use std::{io, ops::Range};
 
 use bitstream_io::BitWrite;
 
-use crate::{common, BitStore, Error, Model};
+use crate::{
+    common::{self, assert_precision_sufficient},
+    BitStore, Error, Model,
+};
 
 // this algorithm is derived from this article - https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html
 
@@ -65,27 +68,17 @@ where
     /// If these constraints cannot be satisfied this method will panic in debug
     /// builds
     pub fn with_precision(model: M, bitwriter: &'a mut W, precision: u32) -> Self {
-        let frequency_bits = model.max_denominator().log2() + 1;
-        debug_assert!(
-            (precision >= (frequency_bits + 2)),
-            "not enough bits of precision to prevent overflow/underflow",
-        );
-        debug_assert!(
-            (frequency_bits + precision) <= M::B::BITS,
-            "not enough bits in BitStore to support the required precision",
-        );
-
-        Self {
-            model,
-            state: State::new(precision, bitwriter),
-        }
+        let state = State::new(precision, bitwriter);
+        Self::with_state(state, model)
     }
 
     /// Create an encoder from an existing [`State`].
     ///
     /// This is useful for manually chaining a shared buffer through multiple
     /// encoders.
-    pub const fn with_state(state: State<'a, M::B, W>, model: M) -> Self {
+    pub fn with_state(state: State<'a, M::B, W>, model: M) -> Self {
+        #[cfg(debug_assertions)]
+        assert_precision_sufficient::<M>(model.max_denominator(), state.state.precision);
         Self { model, state }
     }
 
@@ -162,10 +155,7 @@ where
     where
         X: Model<B = M::B>,
     {
-        Encoder {
-            model,
-            state: self.state,
-        }
+        Encoder::with_state(self.state, model)
     }
 }
 

tests/precision_checking.rs

@@ -0,0 +1,132 @@
+// these tests check the asserts that are only present in debug configurations
+// so they won't pass in release mode
+#![cfg(debug_assertions)]
+
+use std::{convert::Infallible, io::Cursor, ops::Range};
+
+use arithmetic_coding::{decoder, encoder, Decoder, Encoder};
+use arithmetic_coding_core::one_shot;
+use bitstream_io::{BigEndian, BitReader, BitWriter};
+
+#[derive(Copy, Clone)]
+struct SmallModel;
+impl one_shot::Model for SmallModel {
+    type B = u64;
+    type Symbol = u64;
+    type ValueError = Infallible;
+
+    fn probability(&self, &value: &Self::Symbol) -> Result<Range<Self::B>, Self::ValueError> {
+        #[allow(clippy::range_plus_one)]
+        Ok(value..value + 1)
+    }
+
+    fn max_denominator(&self) -> Self::B {
+        2
+    }
+
+    fn symbol(&self, value: Self::B) -> Self::Symbol {
+        value
+    }
+}
+
+#[derive(Copy, Clone)]
+struct BigModel;
+impl one_shot::Model for BigModel {
+    type B = u64;
+    type Symbol = u64;
+    type ValueError = Infallible;
+
+    fn probability(&self, &value: &Self::Symbol) -> Result<Range<Self::B>, Self::ValueError> {
+        #[allow(clippy::range_plus_one)]
+        Ok(value..value + 1)
+    }
+
+    fn max_denominator(&self) -> Self::B {
+        u64::from(u32::MAX) / 2
+    }
+
+    fn symbol(&self, value: Self::B) -> Self::Symbol {
+        value
+    }
+}
+
+// this is one bit short of what it must be
+const PRECISION: u32 = 32;
+
+// Encoder::new should select the correct precision automagically, so we don't
+// expect it to panic
+#[test]
+fn encoder_new_doesnt_panic() {
+    Encoder::new(
+        one_shot::Wrapper::new(BigModel),
+        &mut BitWriter::endian(Vec::new(), BigEndian),
+    );
+}
+
+#[test]
+#[should_panic(expected = "not enough bits of precision to prevent overflow/underflow")]
+fn encoder_with_precision_panics() {
+    Encoder::with_precision(
+        one_shot::Wrapper::new(BigModel),
+        &mut BitWriter::endian(Vec::new(), BigEndian),
+        PRECISION,
+    );
+}
+
+#[test]
+#[should_panic(expected = "not enough bits of precision to prevent overflow/underflow")]
+fn encoder_with_state_panics() {
+    Encoder::with_state(
+        encoder::State::new(PRECISION, &mut BitWriter::endian(Vec::new(), BigEndian)),
+        one_shot::Wrapper::new(BigModel),
+    );
+}
+
+#[test]
+#[should_panic(expected = "not enough bits of precision to prevent overflow/underflow")]
+fn encoder_chain_panics() {
+    let mut writer = BitWriter::endian(Vec::new(), BigEndian);
+    let encoder =
+        Encoder::with_precision(one_shot::Wrapper::new(SmallModel), &mut writer, PRECISION);
+
+    encoder.chain(one_shot::Wrapper::new(BigModel));
+}
+
+#[test]
+fn decoder_new_doesnt_panic() {
+    Decoder::new(
+        one_shot::Wrapper::new(BigModel),
+        BitReader::endian(Cursor::new(&[]), BigEndian),
+    );
+}
+
+#[test]
+#[should_panic(expected = "not enough bits of precision to prevent overflow/underflow")]
+fn decoder_with_precision_panics() {
+    Decoder::with_precision(
+        one_shot::Wrapper::new(BigModel),
+        BitReader::endian(Cursor::new(&[]), BigEndian),
+        PRECISION,
+    );
+}
+
+#[test]
+#[should_panic(expected = "not enough bits of precision to prevent overflow/underflow")]
+fn decoder_with_state_panics() {
+    Decoder::with_state(
+        decoder::State::new(PRECISION, BitReader::endian(Cursor::new(&[]), BigEndian)),
+        one_shot::Wrapper::new(BigModel),
+    );
+}
+
+#[test]
+#[should_panic(expected = "not enough bits of precision to prevent overflow/underflow")]
+fn decoder_chain_panics() {
+    let decoder = Decoder::with_precision(
+        one_shot::Wrapper::new(SmallModel),
+        BitReader::endian(Cursor::new(&[]), BigEndian),
+        PRECISION,
+    );
+
+    decoder.chain(one_shot::Wrapper::new(BigModel));
+}