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