add parameter to print out margin for overlapped unpacking

This commit is contained in:
2022-09-25 23:44:03 +02:00
parent af5fe898bf
commit b12c8f8d93
5 changed files with 53 additions and 16 deletions

2
Cargo.lock generated
View File

@@ -117,7 +117,7 @@ dependencies = [
[[package]] [[package]]
name = "upkr" name = "upkr"
version = "0.2.0-pre1" version = "0.2.0-pre2"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"cdivsufsort", "cdivsufsort",

View File

@@ -5,7 +5,7 @@ mod match_finder;
mod parsing_packer; mod parsing_packer;
mod rans; mod rans;
pub use lz::unpack; pub use lz::{calculate_margin, unpack};
pub type ProgressCallback<'a> = &'a mut dyn FnMut(usize); pub type ProgressCallback<'a> = &'a mut dyn FnMut(usize);

View File

@@ -124,12 +124,27 @@ impl CoderState {
} }
} }
pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> { pub fn unpack(packed_data: &[u8], config: &Config) -> Vec<u8> {
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<u8>>,
packed_data: &[u8],
config: &Config,
) -> isize {
let mut decoder = RansDecoder::new(packed_data, &config); let mut decoder = RansDecoder::new(packed_data, &config);
let mut contexts = ContextState::new((1 + 255) * config.parity_contexts + 1 + 64 + 64, &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 offset = 0;
let mut position = 0usize;
let mut prev_was_match = false; let mut prev_was_match = false;
let mut margin = 0isize;
fn decode_length( fn decode_length(
decoder: &mut RansDecoder, decoder: &mut RansDecoder,
@@ -152,7 +167,8 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
} }
loop { 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)) if decoder.decode_with_context(&mut contexts.context_mut(literal_base))
== config.is_match_bit == config.is_match_bit
{ {
@@ -178,9 +194,12 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
256 * config.parity_contexts + 65, 256 * config.parity_contexts + 65,
&config, &config,
); );
for _ in 0..length { if let Some(ref mut result) = result {
result.push(result[result.len() - offset]); for _ in 0..length {
result.push(result[result.len() - offset]);
}
} }
position += length;
prev_was_match = true; prev_was_match = true;
} else { } else {
let mut context_index = 1; let mut context_index = 1;
@@ -191,10 +210,13 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
context_index = (context_index << 1) | bit as usize; context_index = (context_index << 1) | bit as usize;
byte |= (bit as u8) << i; byte |= (bit as u8) << i;
} }
result.push(byte); if let Some(ref mut result) = result {
result.push(byte);
}
position += 1;
prev_was_match = false; prev_was_match = false;
} }
} }
result margin + decoder.pos() as isize - position as isize
} }

View File

@@ -8,6 +8,7 @@ fn main() -> Result<()> {
let mut config = upkr::Config::default(); let mut config = upkr::Config::default();
let mut reverse = false; let mut reverse = false;
let mut unpack = false; let mut unpack = false;
let mut calculate_margin = false;
let mut level = 2; let mut level = 2;
let mut infile: Option<PathBuf> = None; let mut infile: Option<PathBuf> = None;
let mut outfile: Option<PathBuf> = None; let mut outfile: Option<PathBuf> = None;
@@ -43,6 +44,7 @@ fn main() -> Result<()> {
} }
Short('u') | Long("unpack") => unpack = true, Short('u') | Long("unpack") => unpack = true,
Long("margin") => calculate_margin = true,
Short('l') | Long("level") => level = parser.value()?.parse()?, Short('l') | Long("level") => level = parser.value()?.parse()?,
Short('h') | Long("help") => print_help(0), Short('h') | Long("help") => print_help(0),
Value(val) if infile.is_none() => infile = Some(val.try_into()?), Value(val) if infile.is_none() => infile = Some(val.try_into()?),
@@ -76,7 +78,7 @@ fn main() -> Result<()> {
process::exit(1); process::exit(1);
} }
if !unpack { if !unpack && !calculate_margin {
let mut data = vec![]; let mut data = vec![];
File::open(infile)?.read_to_end(&mut data)?; File::open(infile)?.read_to_end(&mut data)?;
if reverse { if reverse {
@@ -112,11 +114,16 @@ fn main() -> Result<()> {
if reverse { if reverse {
data.reverse(); data.reverse();
} }
let mut unpacked_data = upkr::unpack(&data, config); if unpack {
if reverse { let mut unpacked_data = upkr::unpack(&data, &config);
unpacked_data.reverse(); 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(()) Ok(())
@@ -126,9 +133,11 @@ fn print_help(exit_code: i32) -> ! {
eprintln!("Usage:"); eprintln!("Usage:");
eprintln!(" upkr [-l level(0-9)] [config options] <infile> [<outfile>]"); eprintln!(" upkr [-l level(0-9)] [config options] <infile> [<outfile>]");
eprintln!(" upkr -u [config options] <infile> [<outfile>]"); eprintln!(" upkr -u [config options] <infile> [<outfile>]");
eprintln!(" upkr --margin [config options] <infile>");
eprintln!(); eprintln!();
eprintln!(" -l, --level N compression level 0-9"); eprintln!(" -l, --level N compression level 0-9");
eprintln!(" -u, --unpack unpack infile"); eprintln!(" -u, --unpack unpack infile");
eprintln!(" --margin calculate margin for overlapped unpacking of a packed file");
eprintln!(); eprintln!();
eprintln!("Config presets for specific unpackers:"); eprintln!("Config presets for specific unpackers:");
eprintln!(" --z80 --big-endian-bitstream --invert-bit-encoding --simplified-prob-update"); eprintln!(" --z80 --big-endian-bitstream --invert-bit-encoding --simplified-prob-update");

View File

@@ -149,6 +149,7 @@ impl EntropyCoder for CostCounter {
pub struct RansDecoder<'a> { pub struct RansDecoder<'a> {
data: &'a [u8], data: &'a [u8],
pos: usize,
state: u32, state: u32,
use_bitstream: bool, use_bitstream: bool,
byte: u8, byte: u8,
@@ -163,6 +164,7 @@ impl<'a> RansDecoder<'a> {
pub fn new(data: &'a [u8], config: &Config) -> RansDecoder<'a> { pub fn new(data: &'a [u8], config: &Config) -> RansDecoder<'a> {
RansDecoder { RansDecoder {
data, data,
pos: 0,
state: 0, state: 0,
use_bitstream: config.use_bitstream, use_bitstream: config.use_bitstream,
byte: 0, 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 { pub fn decode_with_context(&mut self, context: &mut Context) -> bool {
let bit = self.decode_bit(context.prob()); let bit = self.decode_bit(context.prob());
context.update(bit); context.update(bit);
@@ -183,8 +189,8 @@ impl<'a> RansDecoder<'a> {
if self.use_bitstream { if self.use_bitstream {
while self.state < 32768 { while self.state < 32768 {
if self.bits_left == 0 { if self.bits_left == 0 {
self.byte = self.data[0]; self.byte = self.data[self.pos];
self.data = &self.data[1..]; self.pos += 1;
self.bits_left = 8; self.bits_left = 8;
} }
if self.bitstream_is_big_endian { if self.bitstream_is_big_endian {