nihav-duck/src/codecs/vpenc/coder.rs

   1 use nihav_core::io::byteio::*;
   2 use nihav_core::codecs::{EncoderResult, EncoderError};
   3 use super::models::*;
   4
   5 pub struct EncSeq {
   6     pub bit:    bool,
   7     pub idx:    u8,
   8 }
   9
  10 pub struct TokenSeq<T: PartialEq> {
  11     pub val:    T,
  12     pub seq:    &'static [EncSeq],
  13 }
  14
  15 #[macro_export]
  16 macro_rules! bit_entry {
  17     (T; $idx:expr) => {EncSeq {bit: true,  idx: $idx }};
  18     (F; $idx:expr) => {EncSeq {bit: false, idx: $idx }};
  19 }
  20
  21 #[macro_export]
  22 macro_rules! bit_seq {
  23     ($val: expr; $( $bit:tt),* ; $( $idx:expr),* ) => {
  24         TokenSeq {
  25             val: $val,
  26             seq:
  27                 &[
  28                 $(
  29                     bit_entry!($bit; $idx),
  30                 )*
  31                 ]
  32         }
  33     };
  34 }
  35
  36 pub struct BoolEncoder<'a, 'b> {
  37     bw:     &'a mut ByteWriter<'b>,
  38     val:    u32,
  39     range:  u32,
  40     bits:   u8,
  41     saved:  u8,
  42     run:    usize,
  43 }
  44
  45 impl<'a, 'b> BoolEncoder<'a, 'b> {
  46     pub fn new(bw: &'a mut ByteWriter<'b>) -> Self {
  47         Self {
  48             bw,
  49             val:    0,
  50             range:  255,
  51             bits:   0,
  52             saved:  0,
  53             run:    0,
  54         }
  55     }
  56     pub fn put_bool(&mut self, bit: bool, prob: u8) -> EncoderResult<()> {
  57         let split = 1 + (((self.range - 1) * u32::from(prob)) >> 8);
  58         if bit {
  59             self.range -= split;
  60             self.val   += split;
  61         } else {
  62             self.range = split;
  63         }
  64
  65         if self.range < 128 {
  66             self.renorm()?;
  67         }
  68         Ok(())
  69     }
  70     fn flush_run(&mut self, overflow: bool) -> EncoderResult<()> {
  71         if self.run > 0 {
  72             self.bw.write_byte(self.saved + (overflow as u8))?;
  73             if !overflow {
  74                 for _ in 1..self.run {
  75                     self.bw.write_byte(0xFF)?;
  76                 }
  77             } else {
  78                 for _ in 1..self.run {
  79                     self.bw.write_byte(0)?;
  80                 }
  81             }
  82             self.run = 0;
  83         }
  84         Ok(())
  85     }
  86     fn renorm(&mut self) -> EncoderResult<()> {
  87         let bits = (self.range.leading_zeros() & 7) as u8;
  88         self.range <<= bits;
  89         if self.bits + bits < 23 {
  90             self.bits += bits;
  91             self.val <<= bits;
  92         } else {
  93             for _ in 0..bits {
  94                 if (self.bits == 23) && ((self.val >> 31) != 0) {
  95                     self.flush_run(true)?;
  96                 }
  97                 self.val <<= 1;
  98                 self.bits += 1;
  99                 if self.bits == 24 {
 100                     let tbyte = (self.val >> 24) as u8;
 101                     let nbyte = (self.val >> 16) as u8;
 102                     if tbyte < 0xFF {
 103                         self.flush_run(false)?;
 104                         if nbyte < 0xFE {
 105                             self.bw.write_byte(tbyte)?;
 106                         } else {
 107                             self.saved = tbyte;
 108                             self.run = 1;
 109                         }
 110                     } else {
 111                         self.run += 1;
 112                     }
 113                     self.val &= 0xFFFFFF;
 114                     self.bits -= 8;
 115                 }
 116             }
 117         }
 118         Ok(())
 119     }
 120     pub fn flush(mut self) -> EncoderResult<()> {
 121         self.flush_run(false)?;
 122         self.val <<= 24 - self.bits;
 123         self.bw.write_u32be(self.val)?;
 124         Ok(())
 125     }
 126
 127     pub fn put_bits(&mut self, val: u32, len: u8) -> EncoderResult<()> {
 128         let mut mask = 1 << (len - 1);
 129         while mask != 0 {
 130             self.put_bool((val & mask) != 0, 128)?;
 131             mask >>= 1;
 132         }
 133         Ok(())
 134     }
 135     pub fn put_probability(&mut self, prob: u8) -> EncoderResult<()> {
 136         self.put_bits(u32::from(prob >> 1), 7)
 137     }
 138     pub fn encode_probability(&mut self, new: u8, old: u8, prob: u8) -> EncoderResult<()> {
 139         self.put_bool(new != old, prob)?;
 140         if new != old {
 141             self.put_probability(new)?;
 142         }
 143         Ok(())
 144     }
 145     pub fn write_el<T: PartialEq>(&mut self, el: T, tree: &[TokenSeq<T>], probs: &[u8]) -> EncoderResult<()> {
 146         for entry in tree.iter() {
 147             if entry.val == el {
 148                 for seq in entry.seq.iter() {
 149                     self.put_bool(seq.bit, probs[seq.idx as usize])?;
 150                 }
 151                 return Ok(());
 152             }
 153         }
 154         Err(EncoderError::Bug)
 155     }
 156 }
 157
 158 pub struct Estimator {}
 159
 160 #[allow(dead_code)]
 161 impl Estimator {
 162     pub fn new() -> Self { Self{} }
 163     pub fn write_el<T: PartialEq>(&self, el: T, tree: &[TokenSeq<T>], probs: &mut [ProbCounter]) {
 164         for entry in tree.iter() {
 165             if entry.val == el {
 166                 for seq in entry.seq.iter() {
 167                     probs[seq.idx as usize].add(seq.bit);
 168                 }
 169                 return;
 170             }
 171         }
 172     }
 173     pub fn est_nits(bit: bool, prob: u8) -> u32 {
 174         if !bit {
 175             u32::from(PROB_BITS[prob as usize])
 176         } else {
 177             u32::from(PROB_BITS[256 - (prob as usize)])
 178         }
 179     }
 180     pub fn nits_to_bits(nits: u32) -> u32 { (nits + 7) >> 3 }
 181 }