diff --git a/src/greedy_packer.rs b/src/greedy_packer.rs index 5e91a95..a0f160a 100644 --- a/src/greedy_packer.rs +++ b/src/greedy_packer.rs @@ -6,11 +6,12 @@ use crate::ProgressCallback; pub fn pack( data: &[u8], use_bitstream: bool, + parity_contexts: usize, mut progress_callback: Option, ) -> Vec { let mut match_finder = MatchFinder::new(data); let mut rans_coder = RansCoder::new(use_bitstream); - let mut state = lz::CoderState::new(); + let mut state = lz::CoderState::new(parity_contexts); let mut pos = 0; while pos < data.len() { diff --git a/src/lib.rs b/src/lib.rs index d53854b..65f2afe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,12 +13,19 @@ pub fn pack( data: &[u8], level: u8, use_bitstream: bool, + parity_contexts: usize, progress_callback: Option, ) -> Vec { if level == 0 { - greedy_packer::pack(data, use_bitstream, progress_callback) + greedy_packer::pack(data, use_bitstream, parity_contexts, progress_callback) } else { - parsing_packer::pack(data, level, use_bitstream, progress_callback) + parsing_packer::pack( + data, + level, + use_bitstream, + parity_contexts, + progress_callback, + ) } } diff --git a/src/lz.rs b/src/lz.rs index 1d01584..c22af22 100644 --- a/src/lz.rs +++ b/src/lz.rs @@ -9,41 +9,49 @@ pub enum Op { impl Op { pub fn encode(&self, coder: &mut dyn EntropyCoder, state: &mut CoderState) { + let literal_base = state.pos % state.parity_contexts * 256; match self { &Op::Literal(lit) => { - encode_bit(coder, state, 0, false); + encode_bit(coder, state, literal_base, false); let mut context_index = 1; for i in (0..8).rev() { let bit = (lit >> i) & 1 != 0; - encode_bit(coder, state, context_index, bit); + encode_bit(coder, state, literal_base + context_index, bit); context_index = (context_index << 1) | bit as usize; } state.prev_was_match = false; + state.pos += 1; } &Op::Match { offset, len } => { - encode_bit(coder, state, 0, true); + encode_bit(coder, state, literal_base, true); if !state.prev_was_match { - encode_bit(coder, state, 256, offset != state.last_offset); + encode_bit( + coder, + state, + 256 * state.parity_contexts, + offset != state.last_offset, + ); } else { assert!(offset != state.last_offset); } if offset != state.last_offset { - encode_length(coder, state, 257, offset + 1); + encode_length(coder, state, 256 * state.parity_contexts + 1, offset + 1); state.last_offset = offset; } - encode_length(coder, state, 257 + 64, len); + encode_length(coder, state, 256 * state.parity_contexts + 65, len); state.prev_was_match = true; + state.pos += len as usize; } } } } pub fn encode_eof(coder: &mut dyn EntropyCoder, state: &mut CoderState) { - encode_bit(coder, state, 0, true); + encode_bit(coder, state, state.pos % state.parity_contexts * 256, true); if !state.prev_was_match { - encode_bit(coder, state, 256, true); + encode_bit(coder, state, 256 * state.parity_contexts, true); } - encode_length(coder, state, 257, 1); + encode_length(coder, state, 256 * state.parity_contexts + 1, 1); } fn encode_bit( @@ -76,16 +84,20 @@ fn encode_length( #[derive(Clone)] pub struct CoderState { contexts: ContextState, + parity_contexts: usize, last_offset: u32, prev_was_match: bool, + pos: usize, } impl CoderState { - pub fn new() -> CoderState { + pub fn new(parity_contexts: usize) -> CoderState { CoderState { - contexts: ContextState::new(1 + 255 + 1 + 64 + 64), + contexts: ContextState::new((1 + 255) * parity_contexts + 1 + 64 + 64), last_offset: 0, + parity_contexts, prev_was_match: false, + pos: 0, } } @@ -94,9 +106,9 @@ impl CoderState { } } -pub fn unpack(packed_data: &[u8], use_bitstream: bool) -> Vec { +pub fn unpack(packed_data: &[u8], use_bitstream: bool, parity_contexts: usize) -> Vec { let mut decoder = RansDecoder::new(packed_data, use_bitstream); - let mut contexts = ContextState::new(1 + 255 + 1 + 64 + 64); + let mut contexts = ContextState::new((1 + 255) * parity_contexts + 1 + 64 + 64); let mut result = vec![]; let mut offset = 0; let mut prev_was_match = false; @@ -119,14 +131,17 @@ pub fn unpack(packed_data: &[u8], use_bitstream: bool) -> Vec { } loop { - if decoder.decode_with_context(&mut contexts.context_mut(0)) { - if prev_was_match || decoder.decode_with_context(&mut contexts.context_mut(256)) { - offset = decode_length(&mut decoder, &mut contexts, 257) - 1; + let literal_base = result.len() % parity_contexts * 256; + if decoder.decode_with_context(&mut contexts.context_mut(literal_base)) { + if prev_was_match + || decoder.decode_with_context(&mut contexts.context_mut(256 * parity_contexts)) + { + offset = decode_length(&mut decoder, &mut contexts, 256 * parity_contexts + 1) - 1; if offset == 0 { break; } } - let length = decode_length(&mut decoder, &mut contexts, 257 + 64); + let length = decode_length(&mut decoder, &mut contexts, 256 * parity_contexts + 65); for _ in 0..length { result.push(result[result.len() - offset]); } @@ -135,7 +150,8 @@ pub fn unpack(packed_data: &[u8], use_bitstream: bool) -> Vec { let mut context_index = 1; let mut byte = 0; for i in (0..8).rev() { - let bit = decoder.decode_with_context(&mut contexts.context_mut(context_index)); + let bit = decoder + .decode_with_context(&mut contexts.context_mut(literal_base + context_index)); context_index = (context_index << 1) | bit as usize; byte |= (bit as u8) << i; } diff --git a/src/main.rs b/src/main.rs index 711cb32..2f49d58 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ use anyhow::{bail, Result}; use std::io::prelude::*; +use std::process; use std::{fs::File, path::PathBuf}; fn main() -> Result<()> { @@ -10,6 +11,14 @@ fn main() -> Result<()> { Some("pack") => { let level = args.opt_value_from_str(["-l", "--level"])?.unwrap_or(2u8); let use_bitstream = args.contains(["-b", "--bitstream"]); + let parity_contexts = args + .opt_value_from_str(["-p", "--parity"])? + .unwrap_or(1usize); + + if parity_contexts != 1 && parity_contexts != 2 && parity_contexts != 4 { + eprintln!("--parity has to be 1, 2 or 4"); + process::exit(1); + } let infile = args.free_from_os_str::(|s| Ok(s.into()))?; let outfile = args.free_from_os_str::(|s| Ok(s.into()))?; @@ -23,6 +32,7 @@ fn main() -> Result<()> { &data, level, use_bitstream, + parity_contexts, Some(&mut |pos| { pb.set(pos as u64); }), @@ -39,13 +49,21 @@ fn main() -> Result<()> { } Some("unpack") => { let use_bitstream = args.contains(["-b", "--bitstream"]); + let parity_contexts = args + .opt_value_from_str(["-p", "--parity"])? + .unwrap_or(1usize); + + if parity_contexts != 1 && parity_contexts != 2 && parity_contexts != 4 { + eprintln!("--parity has to be 1, 2 or 4"); + process::exit(1); + } 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, use_bitstream); + let packed_data = upkr::unpack(&data, use_bitstream, parity_contexts); File::create(outfile)?.write_all(&packed_data)?; } Some(other) => { @@ -58,10 +76,11 @@ fn main() -> Result<()> { fn print_help() { eprintln!("Usage:"); - eprintln!(" upkr pack [-b] [-l level(0-9)] "); - eprintln!(" upkr unpack [-b] "); + eprintln!(" upkr pack [-b] [-l level(0-9)] [-p N] "); + eprintln!(" upkr unpack [-b] [-p N] "); eprintln!(); eprintln!(" -b, --bitstream bitstream mode"); eprintln!(" -l, --level N compression level 0-9"); - std::process::exit(1); + eprintln!(" -p, --parity N use N (2/4) parity contexts"); + process::exit(1); } diff --git a/src/parsing_packer.rs b/src/parsing_packer.rs index d2a6885..f110fe0 100644 --- a/src/parsing_packer.rs +++ b/src/parsing_packer.rs @@ -6,14 +6,25 @@ use crate::match_finder::MatchFinder; use crate::rans::{CostCounter, RansCoder}; use crate::{lz, ProgressCallback}; -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); +pub fn pack( + data: &[u8], + level: u8, + use_bitstream: bool, + parity_contexts: usize, + progress_cb: Option, +) -> Vec { + let mut parse = parse( + data, + Config::from_level(level), + parity_contexts, + progress_cb, + ); let mut ops = vec![]; while let Some(link) = parse { ops.push(link.op); parse = link.prev.clone(); } - let mut state = lz::CoderState::new(); + let mut state = lz::CoderState::new(parity_contexts); let mut coder = RansCoder::new(use_bitstream); for op in ops.into_iter().rev() { op.encode(&mut coder, &mut state); @@ -38,6 +49,7 @@ type Arrivals = HashMap>; fn parse( data: &[u8], config: Config, + parity_contexts: usize, mut progress_cb: Option, ) -> Option> { let mut match_finder = MatchFinder::new(data) @@ -129,7 +141,7 @@ fn parse( 0, Arrival { parse: None, - state: lz::CoderState::new(), + state: lz::CoderState::new(parity_contexts), cost: 0.0, }, max_arrivals,