1 // Copyright 2023 The ChromiumOS Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 use crate::bitstream_utils::BitReader; 6 7 use crate::codec::av1::helpers; 8 9 use super::parser::AnnexBState; 10 11 pub(crate) struct Reader<'a>(pub BitReader<'a>); 12 13 impl<'a> Reader<'a> { new(data: &'a [u8]) -> Self14 pub fn new(data: &'a [u8]) -> Self { 15 Self(BitReader::new(data, false)) 16 } 17 18 /// Implements uvlc(): Variable length unsigned n-bit number appearing 19 /// directly in the bitstream. See 4.10.3 read_uvlc(&mut self) -> Result<u32, String>20 pub fn read_uvlc(&mut self) -> Result<u32, String> { 21 let mut leading_zeroes = 0; 22 loop { 23 let done = self.0.read_bit()?; 24 25 if done { 26 break; 27 } 28 29 leading_zeroes += 1; 30 } 31 32 if leading_zeroes >= 32 { 33 return Ok(u32::MAX); 34 } 35 36 let value = self.0.read_bits::<u32>(leading_zeroes)?; 37 Ok(value + (1 << leading_zeroes) - 1) 38 } 39 40 /// Implements leb128(): Unsigned integer represented by a variable number 41 /// of little-endian bytes. See 4.10.5 read_leb128(&mut self) -> Result<u32, String>42 pub fn read_leb128(&mut self) -> Result<u32, String> { 43 let mut value = 0u64; 44 45 for i in 0..8 { 46 let byte = u64::from(self.0.read_bits_aligned::<u32>(8)?); 47 value |= (byte & 0x7f) << (i * 7); 48 49 if byte & 0x80 == 0 { 50 break; 51 } 52 } 53 54 Ok(value as u32) 55 } 56 57 /// Implements su(n): Signed integer converted from an n bits unsigned 58 /// integer in the bitstream. (The unsigned integer corresponds to the 59 /// bottom n bits of the signed integer.). See 4.10.6 read_su(&mut self, num_bits: usize) -> Result<i32, String>60 pub fn read_su(&mut self, num_bits: usize) -> Result<i32, String> { 61 let mut value: i32 = self 62 .0 63 .read_bits::<u32>(num_bits)? 64 .try_into() 65 .map_err(|_| String::from("Read more than 31 signed bits!"))?; 66 let sign_mask = 1 << (num_bits - 1); 67 68 if (value & sign_mask) != 0 { 69 value -= 2 * sign_mask; 70 } 71 72 Ok(value) 73 } 74 75 /// Implements ns(n): Unsigned encoded integer with maximum number of values 76 /// n (i.e. output in range 0..n-1). See 4.10.7 read_ns(&mut self, num_bits: usize) -> Result<u32, String>77 pub fn read_ns(&mut self, num_bits: usize) -> Result<u32, String> { 78 let w = helpers::floor_log2(num_bits as u32) + 1; 79 let m = (1 << w) - num_bits as u32; 80 let v = self.0.read_bits::<u32>( 81 usize::try_from(w).map_err(|_| String::from("Invalid num_bits"))? - 1, 82 )?; 83 84 if v < m.into() { 85 return Ok(v); 86 } 87 88 let extra_bit = self.0.read_bit()?; 89 Ok((v << 1) - u32::from(m) + u32::from(extra_bit)) 90 } 91 92 /// Implements 5.9.13: Delta quantizer syntax. read_delta_q(&mut self) -> Result<i32, String>93 pub fn read_delta_q(&mut self) -> Result<i32, String> { 94 let delta_coded = self.0.read_bit()?; 95 96 if delta_coded { 97 self.read_su(7) 98 } else { 99 Ok(0) 100 } 101 } 102 more_data_in_bitstream(&mut self) -> bool103 pub fn more_data_in_bitstream(&mut self) -> bool { 104 self.0.num_bits_left() > 0 105 } 106 consumed(&self, start_pos: u32) -> u32107 pub(crate) fn consumed(&self, start_pos: u32) -> u32 { 108 (self.0.position() / 8) as u32 - start_pos 109 } 110 111 /// Get the length of the current OBU in AnnexB format. current_annexb_obu_length( &mut self, annexb_state: &mut AnnexBState, ) -> Result<Option<usize>, String>112 pub fn current_annexb_obu_length( 113 &mut self, 114 annexb_state: &mut AnnexBState, 115 ) -> Result<Option<usize>, String> { 116 if !self.more_data_in_bitstream() { 117 return Ok(None); 118 } 119 120 #[allow(clippy::comparison_chain)] 121 if annexb_state.temporal_unit_consumed == annexb_state.temporal_unit_size { 122 annexb_state.temporal_unit_size = 0; 123 } else if annexb_state.temporal_unit_consumed > annexb_state.temporal_unit_size { 124 return Err(format!( 125 "temporal_unit_size is {} but we consumed {} bytes", 126 annexb_state.temporal_unit_size, annexb_state.temporal_unit_consumed, 127 )); 128 } 129 130 if annexb_state.temporal_unit_size == 0 { 131 annexb_state.temporal_unit_size = self.read_leb128()?; 132 if annexb_state.temporal_unit_size == 0 { 133 return Ok(None); 134 } 135 } 136 137 let start_pos = self.consumed(0); 138 139 #[allow(clippy::comparison_chain)] 140 if annexb_state.frame_unit_consumed == annexb_state.frame_unit_size { 141 annexb_state.frame_unit_size = 0; 142 } else if annexb_state.frame_unit_consumed > annexb_state.frame_unit_size { 143 return Err(format!( 144 "frame_unit_size is {} but we consumed {} bytes", 145 annexb_state.frame_unit_size, annexb_state.frame_unit_consumed, 146 )); 147 } 148 149 if annexb_state.frame_unit_size == 0 { 150 annexb_state.frame_unit_size = self.read_leb128()?; 151 if annexb_state.frame_unit_size == 0 { 152 return Ok(None); 153 } 154 annexb_state.temporal_unit_consumed += self.consumed(start_pos); 155 } 156 157 let start_pos = self.consumed(0); 158 let obu_length = self.read_leb128()?; 159 let consumed = self.consumed(start_pos); 160 161 annexb_state.temporal_unit_consumed += consumed; 162 annexb_state.frame_unit_consumed += consumed; 163 164 Ok(Some(obu_length.try_into().unwrap())) 165 } 166 167 /// Implements 5.3.4. read_trailing_bits(&mut self, mut num_bits: u64) -> Result<(), String>168 pub fn read_trailing_bits(&mut self, mut num_bits: u64) -> Result<(), String> { 169 let trailing_one_bit = self.0.read_bit()?; 170 num_bits -= 1; 171 172 if !trailing_one_bit { 173 return Err("bad padding: trailing_one_bit is not set".into()); 174 } 175 176 while num_bits > 0 { 177 let trailing_zero_bit = self.0.read_bit()?; 178 if trailing_zero_bit { 179 return Err("bad padding: trailing_zero_bit is set".into()); 180 } 181 num_bits -= 1; 182 } 183 184 Ok(()) 185 } 186 decode_subexp(&mut self, num_syms: i32) -> Result<u32, String>187 fn decode_subexp(&mut self, num_syms: i32) -> Result<u32, String> { 188 let mut i = 0; 189 let mut mk = 0; 190 let k = 3; 191 192 loop { 193 let b2 = if i != 0 { k + i - 1 } else { k }; 194 let a = 1 << b2; 195 if num_syms <= mk + 3 * a { 196 let num_bits = num_syms - mk; 197 let subexp_final_bits = self.read_ns(num_bits as usize)?; 198 return Ok(subexp_final_bits); 199 } else { 200 let subexp_more_bits = self.0.read_bit()?; 201 if subexp_more_bits { 202 i += 1; 203 mk += a; 204 } else { 205 let num_bits = b2 as usize; 206 let subexp_bits = self.0.read_bits::<u32>(num_bits)?; 207 return Ok(subexp_bits + mk as u32); 208 } 209 } 210 } 211 } 212 213 /// Implements 5.9.27. decode_unsigned_subexp_with_ref(&mut self, mx: i32, r: i32) -> Result<u32, String>214 pub fn decode_unsigned_subexp_with_ref(&mut self, mx: i32, r: i32) -> Result<u32, String> { 215 let v = self.decode_subexp(mx)?; 216 if (r << 1) <= mx { 217 Ok(helpers::inverse_recenter(r, v.try_into().unwrap()).try_into().unwrap()) 218 } else { 219 let res = mx - 1 - helpers::inverse_recenter(mx - 1 - r, v.try_into().unwrap()); 220 Ok(res.try_into().unwrap()) 221 } 222 } 223 224 /// Implements 5.9.26. decode_signed_subexp_with_ref( &mut self, low: i32, high: i32, r: i32, ) -> Result<i32, String>225 pub fn decode_signed_subexp_with_ref( 226 &mut self, 227 low: i32, 228 high: i32, 229 r: i32, 230 ) -> Result<i32, String> { 231 let x = self.decode_unsigned_subexp_with_ref(high - low, r - low)?; 232 Ok(i32::try_from(x).unwrap() + low) 233 } 234 235 /// Implements 5.3.5 Byte alignment syntax byte_alignment(&mut self) -> Result<(), String>236 pub fn byte_alignment(&mut self) -> Result<(), String> { 237 while (self.0.position() & 7) != 0 { 238 self.0.read_bit()?; 239 } 240 241 Ok(()) 242 } 243 } 244 245 impl<'a> Clone for Reader<'a> { clone(&self) -> Self246 fn clone(&self) -> Self { 247 Self(self.0.clone()) 248 } 249 } 250