nihav-core/src/io/bitreader.rs

   1 #[derive(Debug)]
   2 pub enum BitReaderMode {
   3     BE,
   4     LE,
   5     LE16MSB,
   6     LE32MSB,
   7 }
   8
   9 #[derive(Debug)]
  10 pub enum BitReaderError {
  11     BitstreamEnd,
  12     TooManyBitsRequested,
  13     InvalidValue,
  14 }
  15
  16 use self::BitReaderError::*;
  17
  18 pub type BitReaderResult<T> = Result<T, BitReaderError>;
  19
  20 #[derive(Debug)]
  21 pub struct BitReader<'a> {
  22     cache: u64,
  23     bits:  u8,
  24     pos:   usize,
  25     src:   &'a [u8],
  26     mode:  BitReaderMode,
  27 }
  28
  29 #[allow(clippy::identity_op)]
  30 impl<'a> BitReader<'a> {
  31
  32     pub fn new(src: &'a [u8], mode: BitReaderMode) -> Self {
  33         BitReader{ cache: 0, pos: 0, bits: 0, src, mode }
  34     }
  35
  36     pub fn tell(&self) -> usize {
  37         self.pos * 8 - (self.bits as usize)
  38     }
  39
  40     pub fn left(&self) -> isize {
  41         ((self.src.len() as isize) - (self.pos as isize)) * 8 + (self.bits as isize)
  42     }
  43
  44     fn fill32be(&mut self, src: &[u8]) {
  45         let nw = (u32::from(src[0]) << 24) |
  46                  (u32::from(src[1]) << 16) |
  47                  (u32::from(src[2]) <<  8) |
  48                  (u32::from(src[3]) <<  0);
  49         self.cache |= u64::from(nw) << (32 - self.bits);
  50     }
  51
  52     fn fill32le16(&mut self, src: &[u8]) {
  53         let nw = (u32::from(src[1]) << 24) |
  54                  (u32::from(src[0]) << 16) |
  55                  (u32::from(src[3]) <<  8) |
  56                  (u32::from(src[2]) <<  0);
  57         self.cache |= u64::from(nw) << (32 - self.bits);
  58     }
  59
  60     fn fill32le32(&mut self, src: &[u8], lsb: bool) {
  61         let nw = (u32::from(src[3]) << 24) |
  62                  (u32::from(src[2]) << 16) |
  63                  (u32::from(src[1]) <<  8) |
  64                  (u32::from(src[0]) <<  0);
  65         if lsb {
  66             self.cache |= u64::from(nw) << self.bits;
  67         } else {
  68             self.cache |= u64::from(nw) << (32 - self.bits);
  69         }
  70     }
  71
  72     #[inline(always)]
  73     fn refill(&mut self) -> BitReaderResult<()> {
  74         if self.pos >= self.src.len() { return Err(BitstreamEnd) }
  75         while self.bits <= 32 {
  76             if self.pos + 4 <= self.src.len() {
  77                 let buf = &self.src[self.pos..];
  78                 match self.mode {
  79                     BitReaderMode::BE      => self.fill32be  (buf),
  80                     BitReaderMode::LE16MSB => self.fill32le16(buf),
  81                     BitReaderMode::LE      => self.fill32le32(buf, true),
  82                     BitReaderMode::LE32MSB => self.fill32le32(buf, false),
  83                 }
  84                 self.pos  +=  4;
  85                 self.bits += 32;
  86             } else {
  87                 let mut buf: [u8; 4] = [0, 0, 0, 0];
  88                 let mut newbits: u8 = 0;
  89                 for out in buf.iter_mut().take(3) {
  90                     if self.pos < self.src.len() {
  91                         *out = self.src[self.pos];
  92                         self.pos += 1;
  93                         newbits += 8;
  94                     }
  95                 }
  96                 if newbits == 0 { break; }
  97                 match self.mode {
  98                     BitReaderMode::BE      => self.fill32be  (&buf),
  99                     BitReaderMode::LE16MSB => self.fill32le16(&buf),
 100                     BitReaderMode::LE      => self.fill32le32(&buf, true),
 101                     BitReaderMode::LE32MSB => self.fill32le32(&buf, false),
 102                 }
 103                 self.bits += newbits;
 104             }
 105         }
 106         Ok(())
 107     }
 108
 109     #[inline(always)]
 110     fn read_cache(&mut self, nbits: u8) -> u32 {
 111         let res = match self.mode {
 112             BitReaderMode::LE => ((1u64 << nbits) - 1) & self.cache,
 113             _                 => self.cache >> (64 - nbits),
 114         };
 115         res as u32
 116     }
 117
 118     fn read_cache_s(&mut self, nbits: u8) -> i32 {
 119         let res = match self.mode {
 120             BitReaderMode::LE => ((self.cache as i64) << (64 - nbits)) >> (64 - nbits),
 121             _                 => (self.cache as i64) >> (64 - nbits),
 122         };
 123         res as i32
 124     }
 125
 126     #[inline(always)]
 127     fn skip_cache(&mut self, nbits: u8) {
 128         match self.mode {
 129             BitReaderMode::LE => self.cache >>= nbits,
 130             _                 => self.cache <<= nbits,
 131         };
 132         self.bits -= nbits;
 133     }
 134
 135     #[inline(always)]
 136     fn reset_cache(&mut self) {
 137         self.bits = 0;
 138         self.cache = 0;
 139     }
 140
 141     #[inline(always)]
 142     pub fn read(&mut self, nbits: u8) -> BitReaderResult<u32> {
 143         if nbits == 0 { return Ok(0) }
 144         if nbits > 32 { return Err(TooManyBitsRequested) }
 145         if self.bits < nbits {
 146             if let Err(err) = self.refill() { return Err(err) }
 147             if self.bits < nbits { return Err(BitstreamEnd) }
 148         }
 149         let res = self.read_cache(nbits);
 150         self.skip_cache(nbits);
 151         Ok(res)
 152     }
 153
 154     pub fn read_s(&mut self, nbits: u8) -> BitReaderResult<i32> {
 155         if nbits == 0 || nbits > 32 { return Err(TooManyBitsRequested) }
 156         if self.bits < nbits {
 157             if let Err(err) = self.refill() { return Err(err) }
 158             if self.bits < nbits { return Err(BitstreamEnd) }
 159         }
 160         let res = self.read_cache_s(nbits);
 161         self.skip_cache(nbits);
 162         Ok(res)
 163     }
 164
 165     #[inline(always)]
 166     pub fn read_bool(&mut self) -> BitReaderResult<bool> {
 167         if self.bits < 1 {
 168             if let Err(err) = self.refill() { return Err(err) }
 169             if self.bits < 1 { return Err(BitstreamEnd) }
 170         }
 171         let res = self.read_cache(1);
 172         self.skip_cache(1);
 173         Ok(res == 1)
 174     }
 175
 176     #[inline(always)]
 177     pub fn peek(&mut self, nbits: u8) -> u32 {
 178         if nbits > 32 { return 0 }
 179         if self.bits < nbits { let _ = self.refill(); }
 180         self.read_cache(nbits)
 181     }
 182
 183     #[inline(always)]
 184     pub fn skip(&mut self, nbits: u32) -> BitReaderResult<()> {
 185         if u32::from(self.bits) >= nbits {
 186             self.skip_cache(nbits as u8);
 187             return Ok(());
 188         }
 189         let mut skip_bits = nbits - u32::from(self.bits);
 190         self.reset_cache();
 191         self.pos += ((skip_bits / 32) * 4) as usize;
 192         skip_bits &= 0x1F;
 193         self.refill()?;
 194         if skip_bits > 0 {
 195             self.skip_cache(skip_bits as u8);
 196         }
 197         Ok(())
 198     }
 199
 200     pub fn seek(&mut self, nbits: u32) -> BitReaderResult<()> {
 201         if ((nbits + 7) >> 3) as usize > self.src.len() { return Err(TooManyBitsRequested); }
 202         self.reset_cache();
 203         self.pos = ((nbits / 32) * 4) as usize;
 204         self.skip(nbits & 0x1F)
 205     }
 206
 207     pub fn align(&mut self) {
 208         let pos = self.bits & 7;
 209         if pos != 0 {
 210             self.skip_cache(pos);
 211         }
 212     }
 213 }
 214
 215 pub fn reverse_bits(inval: u32, len: u8) -> u32 {
 216     if len == 0 { return 0; }
 217     const REV_TAB: [u8; 16] = [
 218         0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
 219         0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
 220     ];
 221
 222     let mut ret = 0;
 223     let mut val = inval;
 224     for _ in 0..8 {
 225         ret = (ret << 4) | u32::from(REV_TAB[(val & 0xF) as usize]);
 226         val >>= 4;
 227     }
 228     ret >> (32 - len)
 229 }
 230
 231 #[cfg(test)]
 232 mod test {
 233     use super::*;
 234
 235     #[test]
 236     fn br_works() {
 237         const DATA: [u8; 18] = [0b00011011; 18];
 238         let src = &DATA;
 239         let mut br = BitReader::new(src, BitReaderMode::LE16MSB);
 240
 241         for _ in 0..8 {
 242             assert_eq!(br.read(16).unwrap(), 0x1B1B);
 243         }
 244         const DATA2: [u8; 1] = [ 0b00011011 ];
 245         let src = &DATA2;
 246         let mut br = BitReader::new(src, BitReaderMode::LE);
 247         assert_eq!(br.read_s(5).unwrap(), -5);
 248     }
 249 }