Bool writer optimization (#114)

src/structs/vpx_bool_writer.rs

@@ -52,6 +52,7 @@ impl<W: Write> VPXBoolWriter<W> {
         };
 
         let mut dummy_branch = Branch::new();
+        // initial false bit is put to not get carry out of stream bits
         retval.put_bit(false, &mut dummy_branch, ModelComponent::Dummy)?;
 
         Ok(retval)
@@ -99,15 +100,12 @@ impl<W: Write> VPXBoolWriter<W> {
         if value {
             *tmp_value += split;
             *tmp_range -= split;
-
-            shift = (*tmp_range as u8).leading_zeros() as i32;
         } else {
             *tmp_range = split;
-
-            // optimizer understands that split > 0, so it can optimize this
-            shift = (split as u8).leading_zeros() as i32;
         }
 
+        shift = (*tmp_range as u8).leading_zeros() as i32;
+
         #[cfg(feature = "compression_stats")]
         {
             self.model_statistics
@@ -118,25 +116,18 @@ impl<W: Write> VPXBoolWriter<W> {
         *tmp_count += shift;
 
         if *tmp_count >= 0 {
-            let offset = shift - *tmp_count;
-
-            if ((*tmp_value << (offset - 1)) & 0x80000000) != 0 {
-                let mut x = self.buffer.len() - 1;
-
-                while self.buffer[x] == 0xFF {
-                    self.buffer[x] = 0;
+            let offset = shift - *tmp_count - 1;
 
-                    assert!(x > 0);
-                    x -= 1;
-                }
-
-                self.buffer[x] += 1;
+            *tmp_value <<= offset;
+            if (*tmp_value & 0x80000000) != 0 {
+                self.carry();
             }
 
-            self.buffer.push((*tmp_value >> (24 - offset)) as u8);
-            *tmp_value <<= offset;
-            shift = *tmp_count;
+            *tmp_value <<= 1;
+            self.buffer.push((*tmp_value >> 24) as u8);
             *tmp_value &= 0xffffff;
+
+            shift = *tmp_count;
             *tmp_count -= 8;
         }
 
@@ -145,6 +136,24 @@ impl<W: Write> VPXBoolWriter<W> {
         Ok(())
     }
 
+    /// Safe as: at the stream beginning initially put `false` ensure that carry cannot get out
+    /// of the first stream byte - then `carry` cannot be invoked on empty `buffer`,
+    /// and after the stream beginning `flush_non_final_data` keeps carry-terminating
+    /// byte sequence (one non-255-byte before any number of 255-bytes) inside the `buffer`.
+    #[inline(always)]
+    fn carry(&mut self) {
+        let mut x = self.buffer.len() - 1;
+
+        while self.buffer[x] == 0xFF {
+            self.buffer[x] = 0;
+
+            assert!(x > 0);
+            x -= 1;
+        }
+
+        self.buffer[x] += 1;
+    }
+
     #[inline(always)]
     pub fn put_grid<const A: usize>(
         &mut self,
@@ -284,23 +293,30 @@ impl<W: Write> VPXBoolWriter<W> {
         Ok(())
     }
 
+    // Typically all bytes of `low_value` will have stream bits,
+    // so just write them all - that is what initial Lepton implementation does.
     pub fn finish(&mut self) -> Result<()> {
-        // push real stream bits out of `value`
-        for _i in 0..32 {
-            let mut dummy_branch = Branch::new();
-            self.put_bit(false, &mut dummy_branch, ModelComponent::Dummy)?;
+        let tmp_value = self.low_value << (-self.count - 1);
+
+        if (tmp_value & 0x80000000) != 0 {
+            self.carry();
         }
+        self.buffer.push((tmp_value >> 23) as u8);
+        self.buffer.push((tmp_value >> 15) as u8);
+        self.buffer.push((tmp_value >> 7) as u8);
+        self.buffer.push((tmp_value << 1) as u8);
 
         self.writer.write_all(&self.buffer[..])?;
         Ok(())
     }
 
     /// When buffer is full and is going to be sent to output, preserve buffer data that
-    /// is not final and should carried over to the next buffer.
+    /// is not final and should be carried over to the next buffer. At least one byte
+    /// will remain in `buffer` if it is non-empty.
     pub fn flush_non_final_data(&mut self) -> Result<()> {
         // carry over buffer data that might be not final
         let mut i = self.buffer.len();
-        if i > 0 {
+        if i > 1 {
             i -= 1;
             while self.buffer[i] == 0xFF {
                 assert!(i > 0);