implement simplified prob update, update unpack.c

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,16 +30,19 @@ 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 prob_offset = 16;
+    int state_offset = 0;
+    int state_scale = prob;
     if(bit) {
-        upkr_state = prob * (upkr_state >> 8) + (upkr_state & 255);
-        upkr_probs[context_index] = prob + ((256 - prob + 8) >> 4);
-    } else {
-        upkr_state = (256 - prob) * (upkr_state >> 8) + (upkr_state & 255) - prob;
-        upkr_probs[context_index] = prob - ((prob + 8) >> 4);
+        state_offset = -prob;
+        state_scale = 256 - prob;
+        prob_offset = 0;
     }
-    
+    upkr_state = state_offset + state_scale * (upkr_state >> 8) + (upkr_state & 255);
+    upkr_probs[context_index] = prob_offset + prob - ((prob + 8) >> 4); 
+
     return bit;
 }
 

src/context_state.rs

@@ -1,8 +1,8 @@
-use crate::rans::{PROB_BITS, ONE_PROB};
+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;
     }
 }

src/main.rs

@@ -16,7 +16,7 @@ fn main() -> Result<()> {
 
             let mut data = vec![];
             File::open(infile)?.read_to_end(&mut data)?;
-            
+
             let mut pb = pbr::ProgressBar::new(data.len() as u64);
             pb.set_units(pbr::Units::Bytes);
             let packed_data = upkr::pack(
@@ -58,7 +58,10 @@ fn main() -> Result<()> {
 
 fn print_help() {
     eprintln!("Usage:");
-    eprintln!("  upkr pack [-l level(0-9)] <infile> <outfile>");
-    eprintln!("  upkr unpack <infile> <outfile>");
+    eprintln!("  upkr pack [-b] [-l level(0-9)] <infile> <outfile>");
+    eprintln!("  upkr unpack [-b] <infile> <outfile>");
+    eprintln!();
+    eprintln!(" -b, --bitstream     bitstream mode");
+    eprintln!(" -l, --level N       compression level 0-9");
     std::process::exit(1);
 }

src/rans.rs

@@ -38,15 +38,15 @@ impl RansCoder {
         let mut state = 1 << l_bits;
 
         let mut byte = 0u8;
-        let mut bit = 8;
+        let mut bit = 0;
         let mut flush_state: Box<dyn FnMut(&mut u32)> = if self.use_bitstream {
             Box::new(|state: &mut u32| {
-                bit -= 1;
                 byte |= ((*state & 1) as u8) << bit;
-                if bit == 0 {
+                bit += 1;
+                if bit == 8 {
                     buffer.push(byte);
                     byte = 0;
-                    bit = 8;
+                    bit = 0;
                 }
                 *state >>= 1;
             })
@@ -61,7 +61,7 @@ impl RansCoder {
         let max_state_factor: u32 = 1 << (l_bits + num_flush_bits - PROB_BITS);
         for step in self.bits.into_iter().rev() {
             let prob = step as u32 & 32767;
-            let (start, prob) = if step & 32768 != 0 {
+            let (start, prob) = if step & 32768 == 0 {
                 (0, prob)
             } else {
                 (prob, ONE_PROB - prob)
@@ -118,7 +118,7 @@ impl CostCounter {
 
 impl EntropyCoder for CostCounter {
     fn encode_bit(&mut self, bit: bool, prob: u16) {
-        let prob = if bit {
+        let prob = if !bit {
             prob as u32
         } else {
             ONE_PROB - prob as u32
@@ -163,8 +163,8 @@ impl<'a> RansDecoder<'a> {
                     self.data = &self.data[1..];
                     self.bits_left = 8;
                 }
-                self.state = (self.state << 1) | (self.byte & 1) as u32;
-                self.byte >>= 1;
+                self.state = (self.state << 1) | (self.byte >> 7) as u32;
+                self.byte <<= 1;
                 self.bits_left -= 1;
             }
         } else {
@@ -174,12 +174,12 @@ impl<'a> RansDecoder<'a> {
             }
         }
 
-        let bit = (self.state & PROB_MASK) < prob;
+        let bit = (self.state & PROB_MASK) >= prob;
 
         let (start, prob) = if bit {
-            (0, prob)
-        } else {
             (prob, ONE_PROB - prob)
+        } else {
+            (0, prob)
         };
         self.state = prob * (self.state >> PROB_BITS) + (self.state & PROB_MASK) - start;