diff --git a/src/greedy_packer.rs b/src/greedy_packer.rs index 4e7d271..5e91a95 100644 --- a/src/greedy_packer.rs +++ b/src/greedy_packer.rs @@ -3,9 +3,13 @@ use crate::match_finder::MatchFinder; use crate::rans::RansCoder; use crate::ProgressCallback; -pub fn pack(data: &[u8], mut progress_callback: Option) -> Vec { +pub fn pack( + data: &[u8], + use_bitstream: bool, + mut progress_callback: Option, +) -> Vec { let mut match_finder = MatchFinder::new(data); - let mut rans_coder = RansCoder::new(); + let mut rans_coder = RansCoder::new(use_bitstream); let mut state = lz::CoderState::new(); let mut pos = 0; diff --git a/src/lib.rs b/src/lib.rs index d5153b8..1060e27 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,17 +2,22 @@ mod context_state; mod greedy_packer; mod lz; mod match_finder; -mod rans; mod parsing_packer; +mod rans; pub use lz::unpack; pub type ProgressCallback<'a> = &'a mut dyn FnMut(usize); -pub fn pack(data: &[u8], level: u8, progress_callback: Option) -> Vec { +pub fn pack( + data: &[u8], + level: u8, + use_bitstream: bool, + progress_callback: Option, +) -> Vec { if level == 0 { - greedy_packer::pack(data, progress_callback) + greedy_packer::pack(data, use_bitstream, progress_callback) } else { - parsing_packer::pack(data, level, progress_callback) + parsing_packer::pack(data, level, use_bitstream, progress_callback) } } diff --git a/src/lz.rs b/src/lz.rs index a29d516..1d01584 100644 --- a/src/lz.rs +++ b/src/lz.rs @@ -77,7 +77,7 @@ fn encode_length( pub struct CoderState { contexts: ContextState, last_offset: u32, - prev_was_match: bool + prev_was_match: bool, } impl CoderState { @@ -85,7 +85,7 @@ impl CoderState { CoderState { contexts: ContextState::new(1 + 255 + 1 + 64 + 64), last_offset: 0, - prev_was_match: false + prev_was_match: false, } } @@ -94,8 +94,8 @@ impl CoderState { } } -pub fn unpack(packed_data: &[u8]) -> Vec { - let mut decoder = RansDecoder::new(packed_data); +pub fn unpack(packed_data: &[u8], use_bitstream: bool) -> Vec { + let mut decoder = RansDecoder::new(packed_data, use_bitstream); let mut contexts = ContextState::new(1 + 255 + 1 + 64 + 64); let mut result = vec![]; let mut offset = 0; diff --git a/src/main.rs b/src/main.rs index 4a669bb..ae90eff 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,6 +9,7 @@ fn main() -> Result<()> { None => print_help(), Some("pack") => { let level = args.opt_value_from_str(["-l", "--level"])?.unwrap_or(2u8); + let use_bitstream = args.contains(["-b", "--bitstream"]); let infile = args.free_from_os_str::(|s| Ok(s.into()))?; let outfile = args.free_from_os_str::(|s| Ok(s.into()))?; @@ -21,6 +22,7 @@ fn main() -> Result<()> { let packed_data = upkr::pack( &data, level, + use_bitstream, Some(&mut |pos| { pb.set(pos as u64); }), @@ -36,12 +38,14 @@ fn main() -> Result<()> { File::create(outfile)?.write_all(&packed_data)?; } Some("unpack") => { + let use_bitstream = args.contains(["-b", "--bitstream"]); + let infile = args.free_from_os_str::(|s| Ok(s.into()))?; let outfile = args.free_from_os_str::(|s| Ok(s.into()))?; let mut data = vec![]; File::open(infile)?.read_to_end(&mut data)?; - let packed_data = upkr::unpack(&data); + let packed_data = upkr::unpack(&data, use_bitstream); File::create(outfile)?.write_all(&packed_data)?; } Some(other) => { diff --git a/src/parsing_packer.rs b/src/parsing_packer.rs index 1cdb29f..d2a6885 100644 --- a/src/parsing_packer.rs +++ b/src/parsing_packer.rs @@ -6,7 +6,7 @@ use crate::match_finder::MatchFinder; use crate::rans::{CostCounter, RansCoder}; use crate::{lz, ProgressCallback}; -pub fn pack(data: &[u8], level: u8, progress_cb: Option) -> Vec { +pub fn pack(data: &[u8], level: u8, use_bitstream: bool, progress_cb: Option) -> Vec { let mut parse = parse(data, Config::from_level(level), progress_cb); let mut ops = vec![]; while let Some(link) = parse { @@ -14,7 +14,7 @@ pub fn pack(data: &[u8], level: u8, progress_cb: Option) -> Ve parse = link.prev.clone(); } let mut state = lz::CoderState::new(); - let mut coder = RansCoder::new(); + let mut coder = RansCoder::new(use_bitstream); for op in ops.into_iter().rev() { op.encode(&mut coder, &mut state); } diff --git a/src/rans.rs b/src/rans.rs index 3757447..e81308a 100644 --- a/src/rans.rs +++ b/src/rans.rs @@ -1,6 +1,5 @@ use crate::context_state::Context; -const L_BITS: u32 = 12; pub const PROB_BITS: u32 = 8; pub const ONE_PROB: u32 = 1 << PROB_BITS; @@ -13,43 +12,75 @@ pub trait EntropyCoder { } } -pub struct RansCoder(Vec); +pub struct RansCoder { + bits: Vec, + use_bitstream: bool, +} impl EntropyCoder for RansCoder { fn encode_bit(&mut self, bit: bool, prob: u16) { assert!(prob < 32768); - self.0.push(prob | ((bit as u16) << 15)); + self.bits.push(prob | ((bit as u16) << 15)); } } impl RansCoder { - pub fn new() -> RansCoder { - RansCoder(Vec::new()) + pub fn new(use_bitstream: bool) -> RansCoder { + RansCoder { + bits: Vec::new(), + use_bitstream, + } } pub fn finish(self) -> Vec { let mut buffer = vec![]; - let mut state = 1 << L_BITS; + let l_bits: u32 = if self.use_bitstream { 15 } else { 12 }; + let mut state = 1 << l_bits; - const MAX_STATE_FACTOR: u32 = 1 << (L_BITS + 8 - PROB_BITS); - for step in self.0.into_iter().rev() { + let mut byte = 0u8; + let mut bit = 8; + let mut flush_state: Box = if self.use_bitstream { + Box::new(|state: &mut u32| { + bit -= 1; + byte |= ((*state & 1) as u8) << bit; + if bit == 0 { + buffer.push(byte); + byte = 0; + bit = 8; + } + *state >>= 1; + }) + } else { + Box::new(|state: &mut u32| { + buffer.push(*state as u8); + *state >>= 8; + }) + }; + + let num_flush_bits = if self.use_bitstream { 1 } else { 8 }; + let max_state_factor: u32 = 1 << (l_bits + num_flush_bits - PROB_BITS); + for step in self.bits.into_iter().rev() { let prob = step as u32 & 32767; let (start, prob) = if step & 32768 != 0 { (0, prob) } else { (prob, ONE_PROB - prob) }; - let max_state = MAX_STATE_FACTOR * prob; + let max_state = max_state_factor * prob; while state >= max_state { - buffer.push(state as u8); - state >>= 8; + flush_state(&mut state); } state = ((state / prob) << PROB_BITS) + (state % prob) + start; } while state > 0 { - buffer.push(state as u8); - state >>= 8; + flush_state(&mut state); + } + + drop(flush_state); + + if self.use_bitstream && byte != 0 { + buffer.push(byte); } buffer.reverse(); @@ -99,14 +130,22 @@ impl EntropyCoder for CostCounter { pub struct RansDecoder<'a> { data: &'a [u8], state: u32, + use_bitstream: bool, + byte: u8, + bits_left: u8, } const PROB_MASK: u32 = ONE_PROB - 1; -const L: u32 = 1 << L_BITS; impl<'a> RansDecoder<'a> { - pub fn new(data: &'a [u8]) -> RansDecoder<'a> { - RansDecoder { data, state: 0 } + pub fn new(data: &'a [u8], use_bitstream: bool) -> RansDecoder<'a> { + RansDecoder { + data, + state: 0, + use_bitstream, + byte: 0, + bits_left: 0, + } } pub fn decode_with_context(&mut self, context: &mut Context) -> bool { @@ -117,9 +156,22 @@ impl<'a> RansDecoder<'a> { pub fn decode_bit(&mut self, prob: u16) -> bool { let prob = prob as u32; - while self.state < L { - self.state = (self.state << 8) | self.data[0] as u32; - self.data = &self.data[1..]; + if self.use_bitstream { + while self.state < 32768 { + if self.bits_left == 0 { + self.byte = self.data[0]; + self.data = &self.data[1..]; + self.bits_left = 8; + } + self.state = (self.state << 1) | (self.byte & 1) as u32; + self.byte >>= 1; + self.bits_left -= 1; + } + } else { + while self.state < 4096 { + self.state = (self.state << 8) | self.data[0] as u32; + self.data = &self.data[1..]; + } } let bit = (self.state & PROB_MASK) < prob;