From b12c8f8d93ee1be8c9fa5d3a26c380fc219d2118 Mon Sep 17 00:00:00 2001 From: Dennis Ranke Date: Sun, 25 Sep 2022 23:44:03 +0200 Subject: [PATCH] add parameter to print out margin for overlapped unpacking --- Cargo.lock | 2 +- src/lib.rs | 2 +- src/lz.rs | 36 +++++++++++++++++++++++++++++------- src/main.rs | 19 ++++++++++++++----- src/rans.rs | 10 ++++++++-- 5 files changed, 53 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 74cdb0c..d98035b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -117,7 +117,7 @@ dependencies = [ [[package]] name = "upkr" -version = "0.2.0-pre1" +version = "0.2.0-pre2" dependencies = [ "anyhow", "cdivsufsort", diff --git a/src/lib.rs b/src/lib.rs index 25b089d..ff133b8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,7 @@ mod match_finder; mod parsing_packer; mod rans; -pub use lz::unpack; +pub use lz::{calculate_margin, unpack}; pub type ProgressCallback<'a> = &'a mut dyn FnMut(usize); diff --git a/src/lz.rs b/src/lz.rs index 8ca8cd3..759282e 100644 --- a/src/lz.rs +++ b/src/lz.rs @@ -124,12 +124,27 @@ impl CoderState { } } -pub fn unpack(packed_data: &[u8], config: Config) -> Vec { +pub fn unpack(packed_data: &[u8], config: &Config) -> Vec { + let mut result = vec![]; + let _ = unpack_internal(Some(&mut result), packed_data, config); + result +} + +pub fn calculate_margin(packed_data: &[u8], config: &Config) -> isize { + unpack_internal(None, packed_data, config) +} + +pub fn unpack_internal( + mut result: Option<&mut Vec>, + packed_data: &[u8], + config: &Config, +) -> isize { let mut decoder = RansDecoder::new(packed_data, &config); let mut contexts = ContextState::new((1 + 255) * config.parity_contexts + 1 + 64 + 64, &config); - let mut result = vec![]; let mut offset = 0; + let mut position = 0usize; let mut prev_was_match = false; + let mut margin = 0isize; fn decode_length( decoder: &mut RansDecoder, @@ -152,7 +167,8 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec { } loop { - let literal_base = result.len() % config.parity_contexts * 256; + margin = margin.max(position as isize - decoder.pos() as isize); + let literal_base = position % config.parity_contexts * 256; if decoder.decode_with_context(&mut contexts.context_mut(literal_base)) == config.is_match_bit { @@ -178,9 +194,12 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec { 256 * config.parity_contexts + 65, &config, ); - for _ in 0..length { - result.push(result[result.len() - offset]); + if let Some(ref mut result) = result { + for _ in 0..length { + result.push(result[result.len() - offset]); + } } + position += length; prev_was_match = true; } else { let mut context_index = 1; @@ -191,10 +210,13 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec { context_index = (context_index << 1) | bit as usize; byte |= (bit as u8) << i; } - result.push(byte); + if let Some(ref mut result) = result { + result.push(byte); + } + position += 1; prev_was_match = false; } } - result + margin + decoder.pos() as isize - position as isize } diff --git a/src/main.rs b/src/main.rs index ab73858..a3708a6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ fn main() -> Result<()> { let mut config = upkr::Config::default(); let mut reverse = false; let mut unpack = false; + let mut calculate_margin = false; let mut level = 2; let mut infile: Option = None; let mut outfile: Option = None; @@ -43,6 +44,7 @@ fn main() -> Result<()> { } Short('u') | Long("unpack") => unpack = true, + Long("margin") => calculate_margin = true, Short('l') | Long("level") => level = parser.value()?.parse()?, Short('h') | Long("help") => print_help(0), Value(val) if infile.is_none() => infile = Some(val.try_into()?), @@ -76,7 +78,7 @@ fn main() -> Result<()> { process::exit(1); } - if !unpack { + if !unpack && !calculate_margin { let mut data = vec![]; File::open(infile)?.read_to_end(&mut data)?; if reverse { @@ -112,11 +114,16 @@ fn main() -> Result<()> { if reverse { data.reverse(); } - let mut unpacked_data = upkr::unpack(&data, config); - if reverse { - unpacked_data.reverse(); + if unpack { + let mut unpacked_data = upkr::unpack(&data, &config); + if reverse { + unpacked_data.reverse(); + } + File::create(outfile)?.write_all(&unpacked_data)?; + } + if calculate_margin { + println!("{}", upkr::calculate_margin(&data, &config)); } - File::create(outfile)?.write_all(&unpacked_data)?; } Ok(()) @@ -126,9 +133,11 @@ fn print_help(exit_code: i32) -> ! { eprintln!("Usage:"); eprintln!(" upkr [-l level(0-9)] [config options] []"); eprintln!(" upkr -u [config options] []"); + eprintln!(" upkr --margin [config options] "); eprintln!(); eprintln!(" -l, --level N compression level 0-9"); eprintln!(" -u, --unpack unpack infile"); + eprintln!(" --margin calculate margin for overlapped unpacking of a packed file"); eprintln!(); eprintln!("Config presets for specific unpackers:"); eprintln!(" --z80 --big-endian-bitstream --invert-bit-encoding --simplified-prob-update"); diff --git a/src/rans.rs b/src/rans.rs index ff8d770..3ee61d0 100644 --- a/src/rans.rs +++ b/src/rans.rs @@ -149,6 +149,7 @@ impl EntropyCoder for CostCounter { pub struct RansDecoder<'a> { data: &'a [u8], + pos: usize, state: u32, use_bitstream: bool, byte: u8, @@ -163,6 +164,7 @@ impl<'a> RansDecoder<'a> { pub fn new(data: &'a [u8], config: &Config) -> RansDecoder<'a> { RansDecoder { data, + pos: 0, state: 0, use_bitstream: config.use_bitstream, byte: 0, @@ -172,6 +174,10 @@ impl<'a> RansDecoder<'a> { } } + pub fn pos(&self) -> usize { + self.pos + } + pub fn decode_with_context(&mut self, context: &mut Context) -> bool { let bit = self.decode_bit(context.prob()); context.update(bit); @@ -183,8 +189,8 @@ impl<'a> RansDecoder<'a> { if self.use_bitstream { while self.state < 32768 { if self.bits_left == 0 { - self.byte = self.data[0]; - self.data = &self.data[1..]; + self.byte = self.data[self.pos]; + self.pos += 1; self.bits_left = 8; } if self.bitstream_is_big_endian {