src/io/bitreader.rs

   1 #[derive(Debug)]
   2 pub enum BitReaderMode {
   3     BE,
   4     LE16,
   5     LE32,
   6 }
   7
   8 #[derive(Debug)]
   9 pub enum BitReaderError {
  10     BitstreamEnd,
  11     TooManyBitsRequested,
  12 }
  13
  14 use self::BitReaderError::*;
  15
  16 type BitReaderResult<T> = Result<T, BitReaderError>;
  17
  18 #[derive(Debug)]
  19 pub struct BitReader<'a> {
  20     cache: u64,
  21     bits:  u8,
  22     pos:   usize,
  23     end:   usize,
  24     src:   &'a [u8],
  25     mode:  BitReaderMode,
  26 }
  27
  28 impl<'a> BitReader<'a> {
  29
  30     pub fn new(src: &'a [u8], size: usize, mode: BitReaderMode) -> Self {
  31         if src.len() < size { panic!("size is less than needed"); }
  32         BitReader{ cache: 0, pos: 0, bits: 0, end: size, src: src, mode: mode }
  33     }
  34
  35     pub fn tell(&self) -> usize {
  36         self.pos * 8 - (self.bits as usize)
  37     }
  38
  39     pub fn left(&self) -> isize {
  40         ((self.end as isize) - (self.pos as isize)) * 8 + (self.bits as isize)
  41     }
  42
  43     fn fill32be(&mut self, src: &[u8]) {
  44         let nw = (((src[0] as u32) << 24) |
  45                   ((src[1] as u32) << 16) |
  46                   ((src[2] as u32) <<  8) |
  47                   ((src[3] as u32) <<  0)) as u64;
  48         self.cache |= nw << (32 - self.bits);
  49     }
  50
  51     fn fill32le16(&mut self, src: &[u8], realbits: u8) {
  52         let mut nw = (((src[1] as u32) << 24) |
  53                       ((src[0] as u32) << 16) |
  54                       ((src[3] as u32) <<  8) |
  55                       ((src[2] as u32) <<  0)) as u64;
  56         if realbits <= 16 { nw >>= 16; }
  57         self.cache |= nw << self.bits;
  58     }
  59
  60     fn fill32le32(&mut self, src: &[u8]) {
  61         let nw = (((src[3] as u32) << 24) |
  62                   ((src[2] as u32) << 16) |
  63                   ((src[1] as u32) <<  8) |
  64                   ((src[0] as u32) <<  0)) as u64;
  65         self.cache |= nw << self.bits;
  66     }
  67
  68     fn refill(&mut self) -> BitReaderResult<()> {
  69         if self.pos >= self.end { return Err(BitstreamEnd) }
  70         while self.bits <= 32 {
  71             if self.pos + 4 <= self.end {
  72                 let buf = &self.src[self.pos..];
  73                 match self.mode {
  74                     BitReaderMode::BE   => self.fill32be  (buf),
  75                     BitReaderMode::LE16 => self.fill32le16(buf, 32),
  76                     BitReaderMode::LE32 => self.fill32le32(buf),
  77                 }
  78                 self.pos  +=  4;
  79                 self.bits += 32;
  80             } else {
  81                 let mut buf: [u8; 4] = [0, 0, 0, 0];
  82                 let mut newbits: u8 = 0;
  83                 for i in 0..3 {
  84                     if self.pos < self.end {
  85                         buf[i] = self.src[self.pos];
  86                         self.pos = self.pos + 1;
  87                         newbits += 8;
  88                     }
  89                 }
  90                 if newbits == 0 { break; }
  91                 match self.mode {
  92                     BitReaderMode::BE   => self.fill32be  (&buf),
  93                     BitReaderMode::LE16 => self.fill32le16(&buf, newbits),
  94                     BitReaderMode::LE32 => self.fill32le32(&buf),
  95                 }
  96                 self.bits += newbits;
  97             }
  98         }
  99         Ok(())
 100     }
 101
 102     fn read_cache(&mut self, nbits: u8) -> u32 {
 103         let res = match self.mode {
 104             BitReaderMode::BE => (self.cache as u64) >> (64 - nbits),
 105             _                 => ((1u64 << nbits) - 1) & self.cache,
 106         };
 107         res as u32
 108     }
 109
 110     fn skip_cache(&mut self, nbits: u8) {
 111         match self.mode {
 112             BitReaderMode::BE => self.cache <<= nbits,
 113             _                 => self.cache >>= nbits,
 114         };
 115         self.bits -= nbits;
 116     }
 117
 118     fn reset_cache(&mut self) {
 119         self.bits = 0;
 120         self.cache = 0;
 121     }
 122
 123     pub fn read(&mut self, nbits: u8) -> BitReaderResult<u32> {
 124         if nbits > 32 { return Err(TooManyBitsRequested) }
 125         if self.bits < nbits {
 126             if let Err(err) = self.refill() { return Err(err) }
 127             if self.bits < nbits { return Err(BitstreamEnd) }
 128         }
 129         let res = self.read_cache(nbits);
 130         self.skip_cache(nbits);
 131         Ok(res as u32)
 132     }
 133
 134     pub fn peek(&mut self, nbits: u8) -> u32 {
 135         if nbits > 32 { return 0 }
 136         if self.bits < nbits { let _ = self.refill(); }
 137         self.read_cache(nbits)
 138     }
 139
 140     pub fn skip(&mut self, nbits: u32) -> BitReaderResult<()> {
 141         if self.bits as u32 >= nbits {
 142             self.skip_cache(nbits as u8);
 143             return Ok(());
 144         }
 145         let mut skip_bits = nbits - (self.bits as u32);
 146         self.reset_cache();
 147         self.pos += ((skip_bits / 32) * 4) as usize;
 148         skip_bits = skip_bits & 0x1F;
 149         self.refill()?;
 150         if skip_bits > 0 {
 151             self.skip_cache(skip_bits as u8);
 152         }
 153         Ok(())
 154     }
 155
 156     pub fn seek(&mut self, nbits: u32) -> BitReaderResult<()> {
 157         if ((nbits + 7) >> 3) as usize > self.end { return Err(TooManyBitsRequested); }
 158         self.reset_cache();
 159         self.pos = ((nbits / 32) * 4) as usize;
 160         self.skip(nbits & 0x1F)
 161     }
 162 }
 163
 164 #[cfg(test)]
 165 mod test {
 166     use super::*;
 167
 168     #[test]
 169     fn br_works() {
 170         const DATA: [u8; 18] = [0b00011011; 18];
 171         let src = &DATA;
 172         let mut br = BitReader::new(src, src.len(), BitReaderMode::LE16);
 173
 174         for _ in 0..8 {
 175             assert_eq!(br.read(16).unwrap(), 0x1B1B);
 176         }
 177     }
 178 }