[nihav.git] / src / io / bitreader.rs

#[derive(Debug)]
pub enum BitReaderMode {
    BE,
    LE,
    LE16MSB,
    LE32MSB,
}

#[derive(Debug)]
pub enum BitReaderError {
    BitstreamEnd,
    TooManyBitsRequested,
    InvalidValue,
}

use self::BitReaderError::*;

pub type BitReaderResult<T> = Result<T, BitReaderError>;

#[derive(Debug)]
pub struct BitReader<'a> {
    cache: u64,
    bits:  u8,
    pos:   usize,
    end:   usize,
    src:   &'a [u8],
    mode:  BitReaderMode,
}

impl<'a> BitReader<'a> {

    pub fn new(src: &'a [u8], size: usize, mode: BitReaderMode) -> Self {
        if src.len() < size { panic!("size is less than needed"); }
        BitReader{ cache: 0, pos: 0, bits: 0, end: size, src: src, mode: mode }
    }

    pub fn tell(&self) -> usize {
        self.pos * 8 - (self.bits as usize)
    }

    pub fn left(&self) -> isize {
        ((self.end as isize) - (self.pos as isize)) * 8 + (self.bits as isize)
    }

    fn fill32be(&mut self, src: &[u8]) {
        let nw = (((src[0] as u32) << 24) |
                  ((src[1] as u32) << 16) |
                  ((src[2] as u32) <<  8) |
                  ((src[3] as u32) <<  0)) as u64;
        self.cache |= nw << (32 - self.bits);
    }

    fn fill32le16(&mut self, src: &[u8]) {
        let nw = (((src[1] as u32) << 24) |
                  ((src[0] as u32) << 16) |
                  ((src[3] as u32) <<  8) |
                  ((src[2] as u32) <<  0)) as u64;
        self.cache |= nw << (32 - self.bits);
    }

    fn fill32le32(&mut self, src: &[u8], lsb: bool) {
        let nw = (((src[3] as u32) << 24) |
                  ((src[2] as u32) << 16) |
                  ((src[1] as u32) <<  8) |
                  ((src[0] as u32) <<  0)) as u64;
        if lsb {
            self.cache |= nw << self.bits;
        } else {
            self.cache |= nw << (32 - self.bits);
        }
    }

    fn refill(&mut self) -> BitReaderResult<()> {
        if self.pos >= self.end { return Err(BitstreamEnd) }
        while self.bits <= 32 {
            if self.pos + 4 <= self.end {
                let buf = &self.src[self.pos..];
                match self.mode {
                    BitReaderMode::BE      => self.fill32be  (buf),
                    BitReaderMode::LE16MSB => self.fill32le16(buf),
                    BitReaderMode::LE      => self.fill32le32(buf, true),
                    BitReaderMode::LE32MSB => self.fill32le32(buf, false),
                }
                self.pos  +=  4;
                self.bits += 32;
            } else {
                let mut buf: [u8; 4] = [0, 0, 0, 0];
                let mut newbits: u8 = 0;
                for i in 0..3 {
                    if self.pos < self.end {
                        buf[i] = self.src[self.pos];
                        self.pos = self.pos + 1;
                        newbits += 8;
                    }
                }
                if newbits == 0 { break; }
                match self.mode {
                    BitReaderMode::BE      => self.fill32be  (&buf),
                    BitReaderMode::LE16MSB => self.fill32le16(&buf),
                    BitReaderMode::LE      => self.fill32le32(&buf, true),
                    BitReaderMode::LE32MSB => self.fill32le32(&buf, false),
                }
                self.bits += newbits;
            }
        }
        Ok(())
    }

    fn read_cache(&mut self, nbits: u8) -> u32 {
        let res = match self.mode {
            BitReaderMode::LE => ((1u64 << nbits) - 1) & self.cache,
            _                 => (self.cache as u64) >> (64 - nbits),
        };
        res as u32
    }

    fn read_cache_s(&mut self, nbits: u8) -> i32 {
        let res = match self.mode {
            BitReaderMode::LE => ((self.cache as i64) << (64 - nbits)) >> (64 - nbits),
            _                 => (self.cache as i64) >> (64 - nbits),
        };
        res as i32
    }

    fn skip_cache(&mut self, nbits: u8) {
        match self.mode {
            BitReaderMode::LE => self.cache >>= nbits,
            _                 => self.cache <<= nbits,
        };
        self.bits -= nbits;
    }

    fn reset_cache(&mut self) {
        self.bits = 0;
        self.cache = 0;
    }

    pub fn read(&mut self, nbits: u8) -> BitReaderResult<u32> {
        if nbits == 0 { return Ok(0) }
        if nbits > 32 { return Err(TooManyBitsRequested) }
        if self.bits < nbits {
            if let Err(err) = self.refill() { return Err(err) }
            if self.bits < nbits { return Err(BitstreamEnd) }
        }
        let res = self.read_cache(nbits);
        self.skip_cache(nbits);
        Ok(res)
    }

    pub fn read_s(&mut self, nbits: u8) -> BitReaderResult<i32> {
        if nbits == 0 || nbits > 32 { return Err(TooManyBitsRequested) }
        if self.bits < nbits {
            if let Err(err) = self.refill() { return Err(err) }
            if self.bits < nbits { return Err(BitstreamEnd) }
        }
        let res = self.read_cache_s(nbits);
        self.skip_cache(nbits);
        Ok(res)
    }

    pub fn read_bool(&mut self) -> BitReaderResult<bool> {
        if self.bits < 1 {
            if let Err(err) = self.refill() { return Err(err) }
            if self.bits < 1 { return Err(BitstreamEnd) }
        }
        let res = self.read_cache(1);
        self.skip_cache(1);
        Ok(res == 1)
    }

    pub fn peek(&mut self, nbits: u8) -> u32 {
        if nbits > 32 { return 0 }
        if self.bits < nbits { let _ = self.refill(); }
        self.read_cache(nbits)
    }

    pub fn skip(&mut self, nbits: u32) -> BitReaderResult<()> {
        if self.bits as u32 >= nbits {
            self.skip_cache(nbits as u8);
            return Ok(());
        }
        let mut skip_bits = nbits - (self.bits as u32);
        self.reset_cache();
        self.pos += ((skip_bits / 32) * 4) as usize;
        skip_bits = skip_bits & 0x1F;
        self.refill()?;
        if skip_bits > 0 {
            self.skip_cache(skip_bits as u8);
        }
        Ok(())
    }

    pub fn seek(&mut self, nbits: u32) -> BitReaderResult<()> {
        if ((nbits + 7) >> 3) as usize > self.end { return Err(TooManyBitsRequested); }
        self.reset_cache();
        self.pos = ((nbits / 32) * 4) as usize;
        self.skip(nbits & 0x1F)
    }

    pub fn align(&mut self) {
        let pos = self.bits & 7;
        if pos != 0 {
            self.skip_cache(pos);
        }
    }
}

pub fn reverse_bits(inval: u32, len: u8) -> u32 {
    if len == 0 { return 0; }
    const REV_TAB: [u8; 16] = [
        0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
        0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
    ];

    let mut ret = 0;
    let mut val = inval;
    for _ in 0..8 {
        ret = (ret << 4) | (REV_TAB[(val & 0xF) as usize] as u32);
        val = val >> 4;
    }
    ret >> (32 - len)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn br_works() {
        const DATA: [u8; 18] = [0b00011011; 18];
        let src = &DATA;
        let mut br = BitReader::new(src, src.len(), BitReaderMode::LE16MSB);

        for _ in 0..8 {
            assert_eq!(br.read(16).unwrap(), 0x1B1B);
        }
        const DATA2: [u8; 1] = [ 0b00011011 ];
        let src = &DATA2;
        let mut br = BitReader::new(src, src.len(), BitReaderMode::LE);
        assert_eq!(br.read_s(5).unwrap(), -5);
    }
}
Commit	Line	Data
90aa4e6b KS	1	#[derive(Debug)]
	2	pub enum BitReaderMode {
	3	BE,
a961f92b KS	4	LE,
	5	LE16MSB,
	6	LE32MSB,
90aa4e6b KS	7	}
	8
	9	#[derive(Debug)]
	10	pub enum BitReaderError {
	11	BitstreamEnd,
	12	TooManyBitsRequested,
a961f92b	13	InvalidValue,
90aa4e6b KS	14	}
	15
	16	use self::BitReaderError::*;
	17
a961f92b	18	pub type BitReaderResult<T> = Result<T, BitReaderError>;
90aa4e6b KS	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
8ab3904c	53	fn fill32le16(&mut self, src: &[u8]) {
dfc5d644 KS	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;
8ab3904c	58	self.cache \|= nw << (32 - self.bits);
90aa4e6b KS	59	}
90aa4e6b KS	60
8ab3904c	61	fn fill32le32(&mut self, src: &[u8], lsb: bool) {
90aa4e6b KS	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;
8ab3904c KS	66	if lsb {
	67	self.cache \|= nw << self.bits;
	68	} else {
	69	self.cache \|= nw << (32 - self.bits);
	70	}
90aa4e6b KS	71	}
	72
	73	fn refill(&mut self) -> BitReaderResult<()> {
	74	if self.pos >= self.end { return Err(BitstreamEnd) }
	75	while self.bits <= 32 {
	76	if self.pos + 4 <= self.end {
	77	let buf = &self.src[self.pos..];
	78	match self.mode {
a961f92b	79	BitReaderMode::BE => self.fill32be (buf),
8ab3904c KS	80	BitReaderMode::LE16MSB => self.fill32le16(buf),
	81	BitReaderMode::LE => self.fill32le32(buf, true),
	82	BitReaderMode::LE32MSB => self.fill32le32(buf, false),
90aa4e6b KS	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 i in 0..3 {
	90	if self.pos < self.end {
	91	buf[i] = self.src[self.pos];
	92	self.pos = self.pos + 1;
	93	newbits += 8;
	94	}
	95	}
	96	if newbits == 0 { break; }
	97	match self.mode {
a961f92b	98	BitReaderMode::BE => self.fill32be (&buf),
8ab3904c KS	99	BitReaderMode::LE16MSB => self.fill32le16(&buf),
	100	BitReaderMode::LE => self.fill32le32(&buf, true),
	101	BitReaderMode::LE32MSB => self.fill32le32(&buf, false),
90aa4e6b KS	102	}
	103	self.bits += newbits;
	104	}
	105	}
	106	Ok(())
	107	}
	108
	109	fn read_cache(&mut self, nbits: u8) -> u32 {
	110	let res = match self.mode {
a961f92b KS	111	BitReaderMode::LE => ((1u64 << nbits) - 1) & self.cache,
a961f92b KS	112	_ => (self.cache as u64) >> (64 - nbits),
90aa4e6b KS	113	};
	114	res as u32
	115	}
	116
a961f92b KS	117	fn read_cache_s(&mut self, nbits: u8) -> i32 {
	118	let res = match self.mode {
	119	BitReaderMode::LE => ((self.cache as i64) << (64 - nbits)) >> (64 - nbits),
	120	_ => (self.cache as i64) >> (64 - nbits),
	121	};
	122	res as i32
	123	}
	124
90aa4e6b KS	125	fn skip_cache(&mut self, nbits: u8) {
90aa4e6b KS	126	match self.mode {
a961f92b KS	127	BitReaderMode::LE => self.cache >>= nbits,
a961f92b KS	128	_ => self.cache <<= nbits,
90aa4e6b KS	129	};
	130	self.bits -= nbits;
	131	}
	132
	133	fn reset_cache(&mut self) {
	134	self.bits = 0;
	135	self.cache = 0;
	136	}
	137
	138	pub fn read(&mut self, nbits: u8) -> BitReaderResult<u32> {
a961f92b	139	if nbits == 0 { return Ok(0) }
90aa4e6b KS	140	if nbits > 32 { return Err(TooManyBitsRequested) }
	141	if self.bits < nbits {
	142	if let Err(err) = self.refill() { return Err(err) }
	143	if self.bits < nbits { return Err(BitstreamEnd) }
	144	}
	145	let res = self.read_cache(nbits);
	146	self.skip_cache(nbits);
a961f92b KS	147	Ok(res)
	148	}
	149
	150	pub fn read_s(&mut self, nbits: u8) -> BitReaderResult<i32> {
	151	if nbits == 0 \|\| nbits > 32 { return Err(TooManyBitsRequested) }
	152	if self.bits < nbits {
	153	if let Err(err) = self.refill() { return Err(err) }
	154	if self.bits < nbits { return Err(BitstreamEnd) }
	155	}
	156	let res = self.read_cache_s(nbits);
	157	self.skip_cache(nbits);
	158	Ok(res)
90aa4e6b KS	159	}
90aa4e6b KS	160
0a1d94d9 KS	161	pub fn read_bool(&mut self) -> BitReaderResult<bool> {
	162	if self.bits < 1 {
	163	if let Err(err) = self.refill() { return Err(err) }
	164	if self.bits < 1 { return Err(BitstreamEnd) }
	165	}
	166	let res = self.read_cache(1);
	167	self.skip_cache(1);
	168	Ok(res == 1)
	169	}
	170
90aa4e6b KS	171	pub fn peek(&mut self, nbits: u8) -> u32 {
	172	if nbits > 32 { return 0 }
	173	if self.bits < nbits { let _ = self.refill(); }
	174	self.read_cache(nbits)
	175	}
	176
	177	pub fn skip(&mut self, nbits: u32) -> BitReaderResult<()> {
	178	if self.bits as u32 >= nbits {
	179	self.skip_cache(nbits as u8);
	180	return Ok(());
	181	}
	182	let mut skip_bits = nbits - (self.bits as u32);
	183	self.reset_cache();
	184	self.pos += ((skip_bits / 32) * 4) as usize;
	185	skip_bits = skip_bits & 0x1F;
	186	self.refill()?;
	187	if skip_bits > 0 {
	188	self.skip_cache(skip_bits as u8);
	189	}
	190	Ok(())
	191	}
	192
	193	pub fn seek(&mut self, nbits: u32) -> BitReaderResult<()> {
	194	if ((nbits + 7) >> 3) as usize > self.end { return Err(TooManyBitsRequested); }
	195	self.reset_cache();
	196	self.pos = ((nbits / 32) * 4) as usize;
	197	self.skip(nbits & 0x1F)
	198	}
06fd8c88 KS	199
	200	pub fn align(&mut self) {
	201	let pos = self.bits & 7;
	202	if pos != 0 {
	203	self.skip_cache(pos);
	204	}
	205	}
	206	}
	207
	208	pub fn reverse_bits(inval: u32, len: u8) -> u32 {
	209	if len == 0 { return 0; }
	210	const REV_TAB: [u8; 16] = [
	211	0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
	212	0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
	213	];
	214
	215	let mut ret = 0;
	216	let mut val = inval;
	217	for _ in 0..8 {
	218	ret = (ret << 4) \| (REV_TAB[(val & 0xF) as usize] as u32);
	219	val = val >> 4;
	220	}
	221	ret >> (32 - len)
90aa4e6b KS	222	}
	223
	224	#[cfg(test)]
	225	mod test {
	226	use super::*;
	227
	228	#[test]
	229	fn br_works() {
	230	const DATA: [u8; 18] = [0b00011011; 18];
	231	let src = &DATA;
a961f92b	232	let mut br = BitReader::new(src, src.len(), BitReaderMode::LE16MSB);
90aa4e6b KS	233
	234	for _ in 0..8 {
	235	assert_eq!(br.read(16).unwrap(), 0x1B1B);
	236	}
a961f92b KS	237	const DATA2: [u8; 1] = [ 0b00011011 ];
	238	let src = &DATA2;
	239	let mut br = BitReader::new(src, src.len(), BitReaderMode::LE);
	240	assert_eq!(br.read_s(5).unwrap(), -5);
90aa4e6b KS	241	}
90aa4e6b KS	242	}