codecs/vp6enc/huff.rs

   1 use nihav_core::io::byteio::*;
   2 use nihav_core::codecs::EncoderResult;
   3 use super::super::vpcommon::*;
   4 use super::super::vp6data::*;
   5 use super::models::{VP6HuffModels, VP6Huff};
   6
   7 #[derive(Default)]
   8 pub struct HuffState {
   9     pub dc_zero_run:    [usize; 2],
  10     pub dc_zr_coded:    [bool; 2],
  11     pub ac_zero_run:    [usize; 2],
  12     pub ac_zr_coded:    [bool; 2],
  13 }
  14
  15 impl HuffState {
  16     pub fn new() -> Self { Self::default() }
  17 }
  18
  19 pub const MAX_EOB_RUN: usize = 63 + 10;
  20
  21 pub struct HuffEncoder<'a, 'b> {
  22     bw:         &'a mut ByteWriter<'b>,
  23     bitbuf:     u32,
  24     bits:       u8,
  25 }
  26
  27 impl<'a, 'b> HuffEncoder<'a, 'b> {
  28     pub fn new(bw: &'a mut ByteWriter<'b>) -> Self {
  29         Self {
  30             bitbuf:     0,
  31             bits:       0,
  32             bw
  33         }
  34     }
  35     pub fn flush(mut self) -> EncoderResult<()> {
  36         while self.bits > 0 {
  37             self.bw.write_byte((self.bitbuf >> 24) as u8)?;
  38             self.bitbuf <<= 8;
  39             self.bits = self.bits.saturating_sub(8);
  40         }
  41         Ok(())
  42     }
  43     fn put_bits(&mut self, val: u16, bits: u8) -> EncoderResult<()> {
  44         self.bitbuf |= u32::from(val) << (32 - self.bits - bits);
  45         self.bits   += bits;
  46         while self.bits >= 8 {
  47             self.bw.write_byte((self.bitbuf >> 24) as u8)?;
  48             self.bitbuf <<= 8;
  49             self.bits    -= 8;
  50         }
  51         Ok(())
  52     }
  53     fn encode_eob(&mut self, mdl: &VP6Huff) -> EncoderResult<()> {
  54         self.put_bits(mdl.codes[11], mdl.bits[11])
  55     }
  56     fn encode_val(&mut self, val: i16, mdl: &VP6Huff) -> EncoderResult<()> {
  57         let idx = match val.abs() {
  58                       0 => 0,
  59                       1 => 1,
  60                       2 => 2,
  61                       3 => 3,
  62                       4 => 4,
  63                  5..= 6 => 5,
  64                  7..=10 => 6,
  65                 11..=18 => 7,
  66                 19..=34 => 8,
  67                 35..=66 => 9,
  68                 _ => 10,
  69             };
  70         self.put_bits(mdl.codes[idx], mdl.bits[idx])?;
  71         if idx >= 5 {
  72             self.put_bits((val.abs() - VP56_COEF_BASE[idx - 5]) as u16, VP6_COEF_ADD_BITS[idx - 5])?;
  73         }
  74         if idx > 0 {
  75             self.put_bits((val < 0) as u16, 1)?;
  76         }
  77         Ok(())
  78     }
  79     fn encode_eob_run(&mut self, val: usize) -> EncoderResult<()> {
  80         match val {
  81             0 => { self.put_bits(0, 2)?; },
  82             1 => { self.put_bits(1, 2)?; },
  83             2..=5 => {
  84                 self.put_bits(2, 2)?;
  85                 self.put_bits((val - 2) as u16, 2)?;
  86             },
  87             6..=9 => {
  88                 self.put_bits(3, 2)?;
  89                 self.put_bits(0, 1)?;
  90                 self.put_bits((val - 6) as u16, 2)?;
  91             },
  92             _ => {
  93                 self.put_bits(3, 2)?;
  94                 self.put_bits(1, 1)?;
  95                 self.put_bits((val - 10) as u16, 6)?;
  96             },
  97         };
  98         Ok(())
  99     }
 100     fn encode_zero_run(&mut self, val: usize, mdl: &VP6Huff) -> EncoderResult<()> {
 101         self.put_bits(mdl.codes[val.min(8)], mdl.bits[val.min(8)])?;
 102         if val >= 8 {
 103             self.put_bits((val - 8) as u16, 6)?;
 104         }
 105         Ok(())
 106     }
 107 }
 108
 109 pub fn encode_block_huff(huff: &mut HuffEncoder, scan: &[usize; 64], coeffs: &[i16; 64], plane: usize, hstate: &mut HuffState, model: &VP6HuffModels) -> EncoderResult<()> {
 110     let mut last_idx = 64;
 111     for i in (0..64).rev() {
 112         if coeffs[scan[i]] != 0 {
 113             last_idx = i;
 114             break;
 115         }
 116     }
 117
 118     if !hstate.dc_zr_coded[plane] {
 119         let mdl = &model.dc_token_tree[plane];
 120         huff.encode_val(coeffs[0], mdl)?;
 121         if coeffs[0] == 0 {
 122             huff.encode_eob_run(hstate.dc_zero_run[plane])?;
 123             hstate.dc_zr_coded[plane] = hstate.dc_zero_run[plane] > 0;
 124         }
 125     } else {
 126         hstate.dc_zero_run[plane] -= 1;
 127         if hstate.dc_zero_run[plane] == 0 {
 128             hstate.dc_zr_coded[plane] = false;
 129         }
 130     }
 131     if hstate.ac_zr_coded[plane] {
 132         hstate.ac_zero_run[plane] -= 1;
 133         if hstate.ac_zero_run[plane] == 0 {
 134             hstate.ac_zr_coded[plane] = false;
 135         }
 136         return Ok(());
 137     }
 138
 139     let mut last_val = coeffs[0];
 140
 141     if last_idx == 0 || last_idx == 64 {
 142         let ac_band = VP6_IDX_TO_AC_BAND[1].min(3);
 143         let ac_mode = last_val.abs().min(2) as usize;
 144         let mdl = &model.ac_token_tree[plane][ac_mode][ac_band];
 145         huff.encode_eob(mdl)?;
 146         huff.encode_eob_run(hstate.ac_zero_run[plane])?;
 147         hstate.ac_zr_coded[plane] = hstate.ac_zero_run[plane] > 0;
 148         return Ok(());
 149     }
 150
 151     let mut idx = 1;
 152     while idx < 64 {
 153         let ac_band = VP6_IDX_TO_AC_BAND[idx].min(3);
 154         let ac_mode = last_val.abs().min(2) as usize;
 155         let mdl = &model.ac_token_tree[plane][ac_mode][ac_band];
 156         if idx > last_idx {
 157             huff.encode_eob(mdl)?;
 158             break;
 159         }
 160         let val = coeffs[scan[idx]];
 161         huff.encode_val(val, mdl)?;
 162
 163         idx += 1;
 164         last_val = val;
 165
 166         if val == 0 {
 167             let first_idx = idx;
 168             while idx < 64 && coeffs[scan[idx]] == 0 {
 169                 idx += 1;
 170             }
 171             let zrun = idx - first_idx;
 172             huff.encode_zero_run(zrun, &model.zero_run_tree[if first_idx >= 7 { 1 } else { 0 }])?;
 173         }
 174     }
 175
 176     Ok(())
 177 }