add parameter to print out margin for overlapped unpacking

Cargo.lock

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

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "upkr"
-version = "0.2.0-pre1"
+version = "0.2.0-pre2"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

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);
 
@@ -20,6 +20,8 @@ pub struct Config {
 
     pub bitstream_is_big_endian: bool,
     pub simplified_prob_update: bool,
+
+    pub no_repeated_offsets: bool,
 }
 
 impl Default for Config {
@@ -35,6 +37,8 @@ impl Default for Config {
 
             bitstream_is_big_endian: false,
             simplified_prob_update: false,
+
+            no_repeated_offsets: false,
         }
     }
 }

src/lz.rs

@@ -25,7 +25,7 @@ impl Op {
             }
             &Op::Match { offset, len } => {
                 encode_bit(coder, state, literal_base, config.is_match_bit);
-                if !state.prev_was_match {
+                if !state.prev_was_match && !config.no_repeated_offsets {
                     encode_bit(
                         coder,
                         state,
@@ -33,9 +33,9 @@ impl Op {
                         (offset != state.last_offset) == config.new_offset_bit,
                     );
                 } else {
-                    assert!(offset != state.last_offset);
+                    assert!(offset != state.last_offset || config.no_repeated_offsets);
                 }
-                if offset != state.last_offset {
+                if offset != state.last_offset || config.no_repeated_offsets {
                     encode_length(
                         coder,
                         state,
@@ -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 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,11 +167,13 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
     }
 
     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
         {
-            if prev_was_match
+            if config.no_repeated_offsets
+                || prev_was_match
                 || decoder
                     .decode_with_context(&mut contexts.context_mut(256 * config.parity_contexts))
                     == config.new_offset_bit
@@ -177,9 +194,12 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
                 256 * config.parity_contexts + 65,
                 &config,
             );
+            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;
@@ -190,10 +210,13 @@ pub fn unpack(packed_data: &[u8], config: Config) -> Vec<u8> {
                 context_index = (context_index << 1) | bit as usize;
                 byte |= (bit as u8) << i;
             }
+            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
 }

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<PathBuf> = None;
     let mut outfile: Option<PathBuf> = None;
@@ -28,6 +29,7 @@ fn main() -> Result<()> {
                 config.use_bitstream = true;
                 config.bitstream_is_big_endian = true;
             }
+            Long("no-repeated-offsets") => config.no_repeated_offsets = true,
 
             Long("z80") => {
                 config.use_bitstream = true;
@@ -42,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()?),
@@ -75,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 {
@@ -111,12 +114,17 @@ fn main() -> Result<()> {
         if reverse {
             data.reverse();
         }
-        let mut unpacked_data = upkr::unpack(&data, config);
+        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));
+        }
+    }
 
     Ok(())
 }
@@ -125,9 +133,11 @@ fn print_help(exit_code: i32) -> ! {
     eprintln!("Usage:");
     eprintln!("  upkr [-l level(0-9)] [config options] <infile> [<outfile>]");
     eprintln!("  upkr -u [config options] <infile> [<outfile>]");
+    eprintln!("  upkr --margin [config options] <infile>");
     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");
@@ -145,5 +155,6 @@ fn print_help(exit_code: i32) -> ! {
     eprintln!(" --invert-bit-encoding");
     eprintln!(" --simplified-prob-update");
     eprintln!(" --big-endian-bitstream   (implies --bitstream)");
+    eprintln!(" --no-repeated-offsets");
     process::exit(exit_code);
 }

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.byte = self.data[self.pos];
-                    self.data = &self.data[1..];
+                    self.pos += 1;
                     self.bits_left = 8;
                 }
                 if self.bitstream_is_big_endian {