From 08c86af06f106631ce0401c881bd66fb2c0271e6 Mon Sep 17 00:00:00 2001
From: Dennis Ranke <dennis.ranke@gmail.com>
Date: Mon, 27 Dec 2021 16:58:39 +0100
Subject: [PATCH] some slight tweaks to the compressed format

---
 src/context_state.rs |  5 +++--
 src/lz.rs            | 24 ++++++++++++++++++------
 src/rans.rs          |  4 ++--
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/context_state.rs b/src/context_state.rs
index 9906807..adec205 100644
--- a/src/context_state.rs
+++ b/src/context_state.rs
@@ -2,6 +2,7 @@ use crate::rans::{PROB_BITS, ONE_PROB};
 
 const INIT_PROB: u16 = 1 << (PROB_BITS - 1);
 const UPDATE_RATE: u32 = 4;
+const UPDATE_ADD: u32 = 8;
 
 #[derive(Clone)]
 pub struct ContextState {
@@ -33,9 +34,9 @@ impl<'a> Context<'a> {
     pub fn update(&mut self, bit: bool) {
         let old = self.state.contexts[self.index];
         self.state.contexts[self.index] = if bit {
-            old + ((ONE_PROB - old as u32) >> UPDATE_RATE) as u16
+            old + ((ONE_PROB - old as u32 + UPDATE_ADD) >> UPDATE_RATE) as u16
         } else {
-            old - (old >> UPDATE_RATE)
+            old - ((old + UPDATE_ADD as u16) >> UPDATE_RATE)
         };
     }
 }
diff --git a/src/lz.rs b/src/lz.rs
index 6eb128c..defaa76 100644
--- a/src/lz.rs
+++ b/src/lz.rs
@@ -18,15 +18,21 @@ impl Op {
                     encode_bit(coder, state, context_index, bit);
                     context_index = (context_index << 1) | bit as usize;
                 }
+                state.prev_was_match = false;
             }
             &Op::Match { offset, len } => {
                 encode_bit(coder, state, 0, true);
-                encode_bit(coder, state, 256, offset != state.last_offset);
+                if !state.prev_was_match {
+                    encode_bit(coder, state, 256, offset != state.last_offset);
+                } else {
+                    assert!(offset != state.last_offset);
+                }
                 if offset != state.last_offset {
                     encode_length(coder, state, 257, offset + 1);
                     state.last_offset = offset;
                 }
                 encode_length(coder, state, 257 + 64, len);
+                state.prev_was_match = true;
             }
         }
     }
@@ -51,15 +57,16 @@ fn encode_length(
     coder: &mut dyn EntropyCoder,
     state: &mut CoderState,
     context_start: usize,
-    value: u32,
+    mut value: u32,
 ) {
     assert!(value >= 1);
-    let top_bit = u32::BITS - 1 - value.leading_zeros();
+
     let mut context_index = context_start;
-    for i in 0..top_bit {
+    while value >= 2 {
         encode_bit(coder, state, context_index, true);
-        encode_bit(coder, state, context_index + 1, (value >> i) & 1 != 0);
+        encode_bit(coder, state, context_index + 1, value & 1 != 0);
         context_index += 2;
+        value >>= 1;
     }
     encode_bit(coder, state, context_index, false);
 }
@@ -68,6 +75,7 @@ fn encode_length(
 pub struct CoderState {
     contexts: ContextState,
     last_offset: u32,
+    prev_was_match: bool
 }
 
 impl CoderState {
@@ -75,6 +83,7 @@ impl CoderState {
         CoderState {
             contexts: ContextState::new(1 + 255 + 1 + 64 + 64),
             last_offset: 0,
+            prev_was_match: false
         }
     }
 
@@ -88,6 +97,7 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
     let mut contexts = ContextState::new(1 + 255 + 1 + 64 + 64);
     let mut result = vec![];
     let mut offset = 0;
+    let mut prev_was_match = false;
 
     fn decode_length(
         decoder: &mut RansDecoder,
@@ -108,7 +118,7 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
 
     loop {
         if decoder.decode_with_context(&mut contexts.context_mut(0)) {
-            if decoder.decode_with_context(&mut contexts.context_mut(256)) {
+            if prev_was_match || decoder.decode_with_context(&mut contexts.context_mut(256)) {
                 offset = decode_length(&mut decoder, &mut contexts, 257) - 1;
                 if offset == 0 {
                     break;
@@ -118,6 +128,7 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
             for _ in 0..length {
                 result.push(result[result.len() - offset]);
             }
+            prev_was_match = true;
         } else {
             let mut context_index = 1;
             let mut byte = 0;
@@ -127,6 +138,7 @@ pub fn unpack(packed_data: &[u8]) -> Vec<u8> {
                 byte |= (bit as u8) << i;
             }
             result.push(byte);
+            prev_was_match = false;
         }
     }
 
diff --git a/src/rans.rs b/src/rans.rs
index bead28a..3757447 100644
--- a/src/rans.rs
+++ b/src/rans.rs
@@ -1,7 +1,7 @@
 use crate::context_state::Context;
 
-const L_BITS: u32 = 16;
-pub const PROB_BITS: u32 = 12;
+const L_BITS: u32 = 12;
+pub const PROB_BITS: u32 = 8;
 pub const ONE_PROB: u32 = 1 << PROB_BITS;
 
 pub trait EntropyCoder {