io/intcode: add variable-length code writing

diff --git a/nihav-core/src/io/intcode.rs b/nihav-core/src/io/intcode.rs
index bc6fb0ad59ede68c3c707943bf03d6a960028e04..55c8f859dd9a2fbc43df7dc8fb891207659aaeea 100644 (file)
--- a/nihav-core/src/io/intcode.rs
+++ b/nihav-core/src/io/intcode.rs
@@ -1,5 +1,6 @@
 //! Some universal integer codes support for bitstream reader.
 use crate::io::bitreader::{BitReader, BitReaderError, BitReaderResult};
+use crate::io::bitwriter::BitWriter;
 
 /// Unsigned integer code types.
 #[derive(Debug)]
@@ -185,10 +186,151 @@ impl<'a> IntCodeReader for BitReader<'a> {
     }
 }
 
+/// Universal integer code writer trait for bitstream writer.
+///
+/// # Examples
+///
+/// Write an unsigned Golomb code:
+/// ````
+/// use nihav_core::io::bitwriter::*;
+/// use nihav_core::io::intcode::{IntCodeWriter,UintCodeType};
+///
+/// let mut bw = BitWriter::new(Vec::new(), BitWriterMode::BE);
+/// bw.write_code(UintCodeType::Golomb(3), 42);
+/// ````
+///
+/// Write signed Elias code:
+/// ````
+/// use nihav_core::io::bitwriter::*;
+/// use nihav_core::io::intcode::{IntCodeWriter,IntCodeType};
+///
+/// let mut bw = BitWriter::new(Vec::new(), BitWriterMode::BE);
+/// bw.write_code_signed(IntCodeType::Gamma, 42)?;
+/// ````
+pub trait IntCodeWriter {
+    /// Writes an unsigned integer code of requested type.
+    fn write_code(&mut self, t: UintCodeType, val: u32);
+    /// Writes signed integer code of requested type.
+    fn write_code_signed(&mut self, t: IntCodeType, val: i32);
+}
+
+impl IntCodeWriter for BitWriter {
+    #[inline(always)]
+    fn write_code(&mut self, t: UintCodeType, val: u32) {
+        match t {
+            UintCodeType::UnaryOnes               => write_unary(self, val, 0),
+            UintCodeType::UnaryZeroes             => write_unary(self, val, 1),
+            UintCodeType::LimitedZeroes(len)      => write_unary_lim(self, val, len, 1),
+            UintCodeType::LimitedOnes(len)        => write_unary_lim(self, val, len, 0),
+            UintCodeType::LimitedUnary(len, term) => write_unary_lim(self, val, len, term),
+            UintCodeType::Unary012                => write_unary_lim(self, val, 2, 0),
+            UintCodeType::Unary210                => write_unary210(self, val),
+            UintCodeType::Golomb(m)               => write_golomb(self, val, m),
+            UintCodeType::Rice(k)                 => write_rice(self, val, k),
+            UintCodeType::Gamma                   => write_gamma(self, val),
+            UintCodeType::GammaP                  => write_gammap(self, val),
+        };
+    }
+    fn write_code_signed(&mut self, t: IntCodeType, val: i32) {
+        match t {
+            IntCodeType::Golomb(m)  => write_golomb(self, sval0mp_to_uval(val), m),
+            IntCodeType::Rice(k)    => write_rice(self, sval0mp_to_uval(val), k),
+            IntCodeType::Gamma      => write_gamma(self, sval0pm_to_uval(val)),
+            IntCodeType::GammaP     => write_gammap(self, sval0pm_to_uval(val)),
+        };
+    }
+}
+
+fn sval0mp_to_uval(val: i32) -> u32 {
+    if val < 0  { (-val as u32) * 2 - 1 }
+    else        { (val as u32) * 2 }
+}
+
+fn sval0pm_to_uval(val: i32) -> u32 {
+    if val >= 0 { (val as u32) * 2 + 1 }
+    else        { (-val as u32) * 2 }
+}
+
+fn write_unary210(bw: &mut BitWriter, val: u32) {
+    bw.write_bit(val == 0);
+    if val != 0 {
+        bw.write_bit(val == 1);
+    }
+}
+
+fn write_unary(bw: &mut BitWriter, val: u32, term: u32) {
+    let term = term != 0;
+    for _ in 0..val {
+        bw.write_bit(!term);
+    }
+    bw.write_bit(term);
+}
+
+fn write_unary_lim(bw: &mut BitWriter, val: u32, maxval: u32, term: u32) {
+    let term = term != 0;
+    for _ in 0..val {
+        bw.write_bit(!term);
+    }
+    if val < maxval {
+        bw.write_bit(term);
+    }
+}
+
+fn write_rice(bw: &mut BitWriter, val: u32, k: u8) {
+    let mut exp = val >> k;
+    while exp >= 16 {
+        bw.write(0, 16);
+        exp -= 16
+    }
+    if exp > 0 {
+        bw.write(0, exp as u8);
+    }
+    bw.write1();
+    if k > 0 {
+        let mant = val & ((1 << k) - 1);
+        bw.write(mant, k);
+    }
+}
+
+fn write_golomb(bw: &mut BitWriter, val: u32, m: u8) {
+    if m == 0 { return; }
+    let nbits = (8 - m.leading_zeros()) as u8;
+    if (m & (m - 1)) == 0 { return write_rice(bw, val, nbits); }
+    let q = val / u32::from(m);
+    let r = val % u32::from(m);
+    let cutoff = u32::from((1 << nbits) - m);
+
+    write_unary(bw, q, 0);
+    if r < cutoff {
+        bw.write(r, nbits - 1);
+    } else {
+        bw.write(r + cutoff, nbits);
+    }
+}
+
+fn write_gamma(bw: &mut BitWriter, val: u32) {
+    let val = val + 1;
+    let bits = (32 - val.leading_zeros()) as u8;
+    let mut mask = 1 << bits >> 2;
+    while mask != 0 {
+        bw.write0();
+        bw.write_bit((val & mask) != 0);
+        mask >>= 1;
+    }
+    bw.write1();
+}
+
+fn write_gammap(bw: &mut BitWriter, val: u32) {
+    let bits = 31 - val.leading_zeros();
+    write_unary(bw, bits, 1);
+    bw.write(val - (1 << bits), bits as u8);
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
     use crate::io::bitreader::*;
+    use crate::io::bitwriter::*;
 
     #[test]
     fn int_codes() {
@@ -199,4 +341,18 @@ mod test {
             assert_eq!(br.read_code(UintCodeType::Golomb(5)).unwrap(), i);
         }
     }
+    #[test]
+    fn rw_codes() {
+        let mut bw = BitWriter::new(Vec::new(), BitWriterMode::BE);
+        bw.write_code(UintCodeType::Golomb(5), 42);
+        bw.write_code(UintCodeType::Gamma, 42);
+        bw.write_code(UintCodeType::GammaP, 42);
+        let data = bw.end();
+        let mut br = BitReader::new(&data, BitReaderMode::BE);
+
+        let mut br = BitReader::new(&data, BitReaderMode::BE);
+        assert_eq!(br.read_code(UintCodeType::Golomb(5)).unwrap(), 42);
+        assert_eq!(br.read_code(UintCodeType::Gamma).unwrap(), 42);
+        assert_eq!(br.read_code(UintCodeType::GammaP).unwrap(), 42);
+    }
 }