From 90aa4e6be97ce0849901ce188b30773b2d6662ff Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Sat, 29 Apr 2017 13:22:49 +0200 Subject: [PATCH] initial work --- .gitignore | 2 + Cargo.toml | 6 + src/io/bitreader.rs | 178 ++++++++++++++++++++++++++ src/io/byteio.rs | 297 ++++++++++++++++++++++++++++++++++++++++++++ src/io/mod.rs | 2 + src/lib.rs | 2 + 6 files changed, 487 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 src/io/bitreader.rs create mode 100644 src/io/byteio.rs create mode 100644 src/io/mod.rs create mode 100644 src/lib.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a9d37c5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +target +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..910d4d1 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,6 @@ +[package] +name = "nihav" +version = "0.1.0" +authors = ["Kostya Shishkov "] + +[dependencies] diff --git a/src/io/bitreader.rs b/src/io/bitreader.rs new file mode 100644 index 0000000..9c702ff --- /dev/null +++ b/src/io/bitreader.rs @@ -0,0 +1,178 @@ +#[derive(Debug)] +pub enum BitReaderMode { + BE, + LE16, + LE32, +} + +#[derive(Debug)] +pub enum BitReaderError { + BitstreamEnd, + TooManyBitsRequested, +} + +use self::BitReaderError::*; + +type BitReaderResult = Result; + +#[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], realbits: u8) { + let mut nw = (((src[1] as u32) << 24) | + ((src[0] as u32) << 16) | + ((src[3] as u32) << 8) | + ((src[2] as u32) << 0)) as u64; + if realbits <= 16 { nw >>= 16; } + self.cache |= nw << self.bits; + } + + fn fill32le32(&mut self, src: &[u8]) { + let nw = (((src[3] as u32) << 24) | + ((src[2] as u32) << 16) | + ((src[1] as u32) << 8) | + ((src[0] as u32) << 0)) as u64; + self.cache |= nw << 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::LE16 => self.fill32le16(buf, 32), + BitReaderMode::LE32 => self.fill32le32(buf), + } + 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::LE16 => self.fill32le16(&buf, newbits), + BitReaderMode::LE32 => self.fill32le32(&buf), + } + self.bits += newbits; + } + } + Ok(()) + } + + fn read_cache(&mut self, nbits: u8) -> u32 { + let res = match self.mode { + BitReaderMode::BE => (self.cache as u64) >> (64 - nbits), + _ => ((1u64 << nbits) - 1) & self.cache, + }; + res as u32 + } + + fn skip_cache(&mut self, nbits: u8) { + match self.mode { + BitReaderMode::BE => 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 { + 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 as u32) + } + + 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) + } +} + +#[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::LE16); + + for _ in 0..8 { + assert_eq!(br.read(16).unwrap(), 0x1B1B); + } + } +} diff --git a/src/io/byteio.rs b/src/io/byteio.rs new file mode 100644 index 0000000..19d2ca5 --- /dev/null +++ b/src/io/byteio.rs @@ -0,0 +1,297 @@ +use std::io::SeekFrom; +use std::fs::File; +use std::io::prelude::*; + +#[derive(Debug)] +pub enum ByteIOError { + EOF, + WrongRange, + WrongIOMode, + NotImplemented, + ReadError, + WriteError, + SeekError, +} + +type ByteIOResult = Result; + +pub trait ByteIO { + fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult; + fn read_byte(&mut self) -> ByteIOResult; + fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult; + fn tell(&mut self) -> u64; + fn seek(&mut self, pos: SeekFrom) -> ByteIOResult; + fn is_eof(&mut self) -> bool; + fn is_seekable(&mut self) -> bool; +} + +#[allow(dead_code)] +pub struct ByteReader<'a> { + io: &'a mut ByteIO, +} + +pub struct MemoryReader<'a> { + buf: &'a [u8], + size: usize, + pos: usize, + do_write: bool, +} + +pub struct FileReader<'a> { + file: &'a File, + eof: bool, +} + +macro_rules! read_int { + ($s: ident, $inttype: ty, $size: expr, $which: ident) => ({ + let mut buf = [0; $size]; + try!($s.read_buf(&mut buf)); + unsafe { + Ok((*(buf.as_ptr() as *const $inttype)).$which()) + } + }) +} + +impl<'a> ByteReader<'a> { + pub fn new(io: &'a mut ByteIO) -> ByteReader { ByteReader { io: io } } + + pub fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult { + self.io.read_buf(buf) + } + + pub fn read_byte(&mut self) -> ByteIOResult { + self.io.read_byte() + } + + pub fn read_u16be(&mut self) -> ByteIOResult { + read_int!(self, u16, 2, to_be) + } + + pub fn read_u24be(&mut self) -> ByteIOResult { + let p16 = self.read_u16be(); + if let Err(e) = p16 { return Err(e); } + let p8 = self.read_byte(); + if let Err(e) = p8 { return Err(e); } + Ok(((p16.unwrap() as u32) << 8) | (p8.unwrap() as u32)) + } + + pub fn read_u32be(&mut self) -> ByteIOResult { + read_int!(self, u32, 4, to_be) + } + + pub fn read_u64be(&mut self) -> ByteIOResult { + read_int!(self, u64, 8, to_be) + } + + pub fn read_u16le(&mut self) -> ByteIOResult { + read_int!(self, u16, 2, to_le) + } + + pub fn read_u24le(&mut self) -> ByteIOResult { + let p8 = self.read_byte(); + if let Err(e) = p8 { return Err(e); } + let p16 = self.read_u16le(); + if let Err(e) = p16 { return Err(e); } + Ok(((p16.unwrap() as u32) << 8) | (p8.unwrap() as u32)) + } + + pub fn read_u32le(&mut self) -> ByteIOResult { + read_int!(self, u32, 4, to_le) + } + + pub fn read_u64le(&mut self) -> ByteIOResult { + read_int!(self, u64, 8, to_le) + } + + pub fn read_skip(&mut self, len: usize) -> ByteIOResult { + if self.io.is_seekable() { + self.io.seek(SeekFrom::Current(len as i64))?; + } else { + let mut ssize = len; + let mut buf : [u8; 16] = [0; 16]; + let mut bref = &mut buf; + while ssize > bref.len() { + self.io.read_buf(bref)?; + ssize -= bref.len(); + } + while ssize > 0 { + self.io.read_byte()?; + ssize = ssize - 1; + } + } + Ok(self.tell()) + } + + pub fn tell(&mut self) -> u64 { + self.io.tell() + } + + pub fn seek(&mut self, pos: SeekFrom) -> ByteIOResult { + self.io.seek(pos) + } + + pub fn is_eof(&mut self) -> bool { + self.io.is_eof() + } +} + +impl<'a> MemoryReader<'a> { + + pub fn new_read(buf: &'a [u8]) -> Self { + MemoryReader { buf: buf, size: buf.len(), pos: 0, do_write: false } + } + + fn real_seek(&mut self, pos: i64) -> ByteIOResult { + if pos < 0 || (pos as usize) > self.size { + return Err(ByteIOError::WrongRange) + } + self.pos = pos as usize; + Ok(pos as u64) + } + + fn check_read_perm(&self) -> ByteIOResult<()> { + if self.do_write { + Err(ByteIOError::WrongIOMode) + } else { + Ok(()) + } + } + + fn check_write_perm(&self) -> ByteIOResult<()> { + if !self.do_write { + Err(ByteIOError::WrongIOMode) + } else { + Ok(()) + } + } +} + +impl<'a> ByteIO for MemoryReader<'a> { + fn read_byte(&mut self) -> ByteIOResult { + self.check_read_perm()?; + if self.is_eof() { return Err(ByteIOError::EOF); } + let res = self.buf[self.pos]; + self.pos = self.pos + 1; + Ok(res) + } + + fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult { + self.check_read_perm()?; + let copy_size = if self.size - self.pos < buf.len() { self.size } else { buf.len() }; + if copy_size == 0 { return Err(ByteIOError::EOF); } + for i in 0..copy_size { + buf[i] = self.buf[self.pos + i]; + } + self.pos += copy_size; + Ok(copy_size) + } + + #[allow(unused_variables)] + fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult { + self.check_write_perm()?; + Err(ByteIOError::NotImplemented) + } + + fn tell(&mut self) -> u64 { + self.pos as u64 + } + + fn seek(&mut self, pos: SeekFrom) -> ByteIOResult { + let cur_pos = self.pos as i64; + let cur_size = self.size as i64; + match pos { + SeekFrom::Start(x) => self.real_seek(x as i64), + SeekFrom::Current(x) => self.real_seek(cur_pos + x), + SeekFrom::End(x) => self.real_seek(cur_size + x), + } + } + + fn is_eof(&mut self) -> bool { + self.pos >= self.size + } + + fn is_seekable(&mut self) -> bool { + true + } +} + +impl<'a> FileReader<'a> { + + pub fn new_read(file: &'a mut File) -> Self { + FileReader { file: file, eof : false } + } +} + +impl<'a> ByteIO for FileReader<'a> { + fn read_byte(&mut self) -> ByteIOResult { + let mut byte : [u8; 1] = [0]; + let err = self.file.read(&mut byte); + if let Err(_) = err { return Err(ByteIOError::ReadError); } + let sz = err.unwrap(); + if sz == 0 { self.eof = true; return Err(ByteIOError::EOF); } + Ok (byte[0]) + } + + fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult { + let res = self.file.read(buf); + if let Err(_) = res { return Err(ByteIOError::ReadError); } + let sz = res.unwrap(); + if sz < buf.len() { self.eof = true; } + Ok(sz) + } + + #[allow(unused_variables)] + fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult { + Err(ByteIOError::NotImplemented) + } + + fn tell(&mut self) -> u64 { + self.file.seek(SeekFrom::Current(0)).unwrap() + } + + fn seek(&mut self, pos: SeekFrom) -> ByteIOResult { + let res = self.file.seek(pos); + match res { + Ok(r) => Ok(r), + Err(_) => Err(ByteIOError::SeekError), + } + } + + fn is_eof(&mut self) -> bool { + self.eof + } + + fn is_seekable(&mut self) -> bool { + true + } +} + +#[cfg(test)] +mod test { + use super::*; + use std::fs::File; + + #[test] + fn test_read() { + //const DATA : &'static [u8] = include_bytes!("../../assets/file"); + let buf: [u8; 64] = [1; 64]; + let mut mr = MemoryReader::new_read(&buf); + let mut reader = ByteReader::new(&mut mr); + assert_eq!(reader.read_byte().unwrap(), 0x01u8); + assert_eq!(reader.read_u16le().unwrap(), 0x0101u16); + assert_eq!(reader.read_u24le().unwrap(), 0x010101u32); + assert_eq!(reader.read_u32le().unwrap(), 0x01010101u32); + assert_eq!(reader.read_u64le().unwrap(), 0x0101010101010101u64); + let mut file = File::open("assets/MaoMacha.asx").unwrap(); + let mut fr = FileReader::new_read(&mut file); + let mut br2 = ByteReader::new(&mut fr); + assert_eq!(br2.read_byte().unwrap(), 0x30); + assert_eq!(br2.read_u24be().unwrap(), 0x26B275); + assert_eq!(br2.read_u24le().unwrap(), 0xCF668E); + assert_eq!(br2.read_u32be().unwrap(), 0x11A6D900); + assert_eq!(br2.read_u32le().unwrap(), 0xCE6200AA); + } + #[test] + fn test_write() { + } +} diff --git a/src/io/mod.rs b/src/io/mod.rs new file mode 100644 index 0000000..cb908c9 --- /dev/null +++ b/src/io/mod.rs @@ -0,0 +1,2 @@ +pub mod bitreader; +pub mod byteio; diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..f039bdf --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,2 @@ +pub mod io; + -- 2.30.2