From f1f1c64a767f23a64d3f26a575afbd2efe1f6a1d Mon Sep 17 00:00:00 2001
From: Dennis Ranke <dennis.ranke@gmail.com>
Date: Sat, 10 Sep 2022 12:01:42 +0200
Subject: [PATCH] implement simplified prob update, update unpack.c

---
 c_unpacker/unpack.c  | 25 ++++++++++++-------------
 src/context_state.rs | 13 ++++++++-----
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/c_unpacker/unpack.c b/c_unpacker/unpack.c
index c893d56..7de3deb 100644
--- a/c_unpacker/unpack.c
+++ b/c_unpacker/unpack.c
@@ -19,8 +19,8 @@ int upkr_decode_bit(int context_index) {
             upkr_current_byte = *upkr_data_ptr++;
             upkr_bits_left = 8;
         }
-        upkr_state = (upkr_state << 1) + (upkr_current_byte & 1);
-        upkr_current_byte >>= 1;
+        upkr_state = (upkr_state << 1) + (upkr_current_byte >> 7);
+        upkr_current_byte <<= 1;
         --upkr_bits_left;
     }
 #else
@@ -30,19 +30,18 @@ int upkr_decode_bit(int context_index) {
 #endif
    
     int prob = upkr_probs[context_index];
-    int bit = (upkr_state & 255) < prob ? 1 : 0;
+    int bit = (upkr_state & 255) >= prob ? 1 : 0;
     
-    int tmp = prob;
-    if(!bit) {
-        tmp = 256 - tmp;
+    int prob_offset = 16;
+    int state_offset = 0;
+    int state_scale = prob;
+    if(bit) {
+        state_offset = -prob;
+        state_scale = 256 - prob;
+        prob_offset = 0;
     }
-    upkr_state = tmp * (upkr_state >> 8) + (upkr_state & 255);
-    tmp += (256 - tmp + 8) >> 4;
-    if(!bit) {
-        upkr_state -= prob;
-        tmp = 256 - tmp;
-    }
-    upkr_probs[context_index] = tmp;
+    upkr_state = state_offset + state_scale * (upkr_state >> 8) + (upkr_state & 255);
+    upkr_probs[context_index] = prob_offset + prob - ((prob + 8) >> 4); 
 
     return bit;
 }
diff --git a/src/context_state.rs b/src/context_state.rs
index 92d5fb2..b7a75f2 100644
--- a/src/context_state.rs
+++ b/src/context_state.rs
@@ -1,8 +1,8 @@
 use crate::rans::{ONE_PROB, PROB_BITS};
 
 const INIT_PROB: u16 = 1 << (PROB_BITS - 1);
-const UPDATE_RATE: u32 = 4;
-const UPDATE_ADD: u32 = 8;
+const UPDATE_RATE: i32 = 4;
+const UPDATE_ADD: i32 = 8;
 
 #[derive(Clone)]
 pub struct ContextState {
@@ -33,10 +33,13 @@ 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_ADD) >> UPDATE_RATE) as u8
+        let offset = if !bit {
+            ONE_PROB as i32 >> UPDATE_RATE
         } else {
-            old - ((old as u32 + UPDATE_ADD) >> UPDATE_RATE) as u8
+            0
         };
+
+        self.state.contexts[self.index] =
+            (offset + old as i32 - ((old as i32 + UPDATE_ADD) >> UPDATE_RATE)) as u8;
     }
 }