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