| 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 | } |