Refactor FrameBuf and MD5 computation.

This commit also includes minor improvements in BitSink.

report/report.md

@@ -22,15 +22,15 @@ Sources used: wikimedia.i_love_you_california, wikimedia.winter_kiss, wikimedia.
 
 ### Average compression speed (inverse RTF)
   - Reference
-    - opt8lax: 260.5208675241677
-    - opt8: 260.4015355192523
-    - opt5: 496.04988453221785
+    - opt8lax: 263.34007288854724
+    - opt8: 263.4615363107738
+    - opt5: 499.8657156848808
 
   - Ours
-    - default: 167.2145640416671
-    - st: 83.99637957900106
-    - dmse: 111.79190850520624
-    - bsbs: 11.622540675983512
-    - mae: 26.102371937358594
+    - default: 181.4254484243478
+    - st: 86.49084643182348
+    - dmse: 113.97081675514823
+    - bsbs: 11.718727584136694
+    - mae: 25.477642469493656
 
 

src/bitsink.rs

@@ -14,14 +14,24 @@
 
 //! Abstract interface for bit-based output.
 
-use std::ops::Shl;
-
+use num_traits::PrimInt;
 use num_traits::ToBytes;
 
 /// Alias trait for the bit-addressible integers.
-pub trait PackedBits: ToBytes + Into<u64> + Shl<usize, Output = Self> + Copy {}
+pub trait PackedBits:
+    ToBytes + From<u8> + Into<u64> + PrimInt + std::ops::ShlAssign<usize>
+where
+    Self::Bytes: AsRef<[u8]>,
+{
+    const PACKED_BITS: usize;
+}
 
-impl<T: ToBytes + Into<u64> + Shl<usize, Output = T> + Copy> PackedBits for T {}
+impl<T: ToBytes + From<u8> + Into<u64> + PrimInt + std::ops::ShlAssign<usize>> PackedBits for T
+where
+    <T as ToBytes>::Bytes: AsRef<[u8]>,
+{
+    const PACKED_BITS: usize = std::mem::size_of::<T>() * 8;
+}
 
 /// Storage-agnostic interface trait for bit-based output.
 ///
@@ -161,34 +171,6 @@ impl ByteVec {
         ret
     }
 
-    /// Appends first `n` bits (from MSB) to the `ByteVec`.
-    #[inline]
-    fn push_u64_msbs(&mut self, val: u64, n: usize) {
-        let mut val: u64 = val;
-        let mut n = n;
-        let nbitlength = self.bitlength + n;
-        let r = self.tail_len();
-
-        if r != 0 {
-            let b: u8 = (val >> (64 - r)) as u8;
-            let tail = self.bytes.len() - 1;
-            self.bytes[tail] |= b;
-            val <<= r;
-            n = if n > r { n - r } else { 0 };
-        }
-        while n >= 8 {
-            let b: u8 = (val >> (64 - 8) & 0xFFu64) as u8;
-            self.bytes.push(b);
-            val <<= 8;
-            n -= 8;
-        }
-        if n > 0 {
-            let b: u8 = ((val >> (64 - n)) << (8 - n)) as u8;
-            self.bytes.push(b);
-        }
-        self.bitlength = nbitlength;
-    }
-
     pub fn as_byte_slice(&self) -> &[u8] {
         &self.bytes
     }
@@ -229,18 +211,41 @@ impl BitSink for ByteVec {
         if n == 0 {
             return;
         }
-        let initial_shift = 64 - (std::mem::size_of::<T>() * 8);
-        let val: u64 = val.into();
-        self.push_u64_msbs(val << initial_shift, n);
+        let mut val: T = val;
+        let mut n = n;
+        let nbitlength = self.bitlength + n;
+        let r = self.tail_len();
+
+        if r != 0 {
+            let b = (val >> (T::PACKED_BITS - r)).to_u8().unwrap();
+            let tail = self.bytes.len() - 1;
+            self.bytes[tail] |= b;
+            val <<= r;
+            n = n.saturating_sub(r);
+        }
+        let bytes_to_write = n / 8;
+        if bytes_to_write > 0 {
+            let bytes = val.to_be_bytes();
+            self.bytes
+                .extend_from_slice(&bytes.as_ref()[0..bytes_to_write]);
+            n %= 8;
+        }
+        if n > 0 {
+            val <<= bytes_to_write * 8;
+            let mask = !((1 << (8 - n)) - 1);
+            let shifted = val >> (T::PACKED_BITS - 8) & mask.into();
+            let tail_byte = shifted.to_u8().unwrap();
+            self.bytes.push(tail_byte);
+        }
+        self.bitlength = nbitlength;
     }
 
     #[inline]
     fn write_lsbs<T: PackedBits>(&mut self, val: T, n: usize) {
         if n == 0 {
             return;
         }
-        let val: u64 = val.into();
-        self.push_u64_msbs(val << (64 - n), n);
+        self.write_msbs(val << (T::PACKED_BITS - n), n);
     }
 }
 

src/coding.rs

@@ -43,6 +43,7 @@ use super::constant::MAX_LPC_ORDER;
 use super::error::SourceError;
 use super::lpc;
 use super::rice;
+use super::source::Context;
 use super::source::FrameBuf;
 use super::source::Seekable;
 use super::source::Source;
@@ -475,20 +476,15 @@ pub fn parallel_encode_with_fixed_block_size<T: Source>(
 ) -> Result<Stream, SourceError> {
     let mut stream = Stream::new(src.sample_rate(), src.channels(), src.bits_per_sample());
     let channels = src.channels();
-    let mut md5_context = md5::Context::new();
-    let mut offset: usize = 0;
-    let mut frame_number: usize = 0;
+    let mut context = Context::new(src.bits_per_sample(), channels);
     let mut framebufs = vec![];
     loop {
         let mut framebuf = FrameBuf::with_size(channels, block_size);
-        let read_samples = src.read_samples(&mut framebuf)?;
+        let read_samples = src.read_samples(&mut framebuf, &mut context)?;
         if read_samples == 0 {
             break;
         }
-        framebuf.update_md5(src.bits_per_sample(), block_size, &mut md5_context);
-        framebufs.push((frame_number, framebuf));
-        frame_number += 1;
-        offset += read_samples;
+        framebufs.push((context.current_frame_number(), framebuf));
     }
     let stream_info = stream.stream_info();
     let frames = framebufs.into_par_iter().map(|(frame_number, framebuf)| {
@@ -500,10 +496,10 @@ pub fn parallel_encode_with_fixed_block_size<T: Source>(
     }
     stream
         .stream_info_mut()
-        .set_total_samples(src.len_hint().unwrap_or(offset));
+        .set_total_samples(src.len_hint().unwrap_or_else(|| context.total_samples()));
     stream
         .stream_info_mut()
-        .set_md5_digest(&md5_context.compute().into());
+        .set_md5_digest(&context.md5_digest());
     Ok(stream)
 }
 
@@ -521,30 +517,27 @@ pub fn encode_with_fixed_block_size<T: Source>(
         return parallel_encode_with_fixed_block_size(config, src, block_size);
     }
     let mut stream = Stream::new(src.sample_rate(), src.channels(), src.bits_per_sample());
-    let channels = src.channels();
-    let mut framebuf = FrameBuf::with_size(channels, block_size);
-
-    let mut md5_context = md5::Context::new();
-    let mut offset: usize = 0;
-    let mut frame_number: usize = 0;
+    let mut framebuf = FrameBuf::with_size(src.channels(), block_size);
+    let mut context = Context::new(src.bits_per_sample(), src.channels());
     loop {
-        let read_samples = src.read_samples(&mut framebuf)?;
+        let read_samples = src.read_samples(&mut framebuf, &mut context)?;
         if read_samples == 0 {
             break;
         }
-        // For the last frame, FLAC uses all the samples in the zero-padded block.
-        framebuf.update_md5(src.bits_per_sample(), block_size, &mut md5_context);
-        let frame = encode_fixed_size_frame(config, &framebuf, frame_number, stream.stream_info());
+        let frame = encode_fixed_size_frame(
+            config,
+            &framebuf,
+            context.current_frame_number(),
+            stream.stream_info(),
+        );
         stream.add_frame(frame);
-        offset += read_samples;
-        frame_number += 1;
     }
     stream
         .stream_info_mut()
-        .set_total_samples(src.len_hint().unwrap_or(offset));
+        .set_md5_digest(&context.md5_digest());
     stream
         .stream_info_mut()
-        .set_md5_digest(&md5_context.compute().into());
+        .set_total_samples(src.len_hint().unwrap_or_else(|| context.total_samples()));
     Ok(stream)
 }
 
@@ -558,7 +551,7 @@ struct Hyp {
     bits: usize,
     back_pointer: Option<Arc<Hyp>>,
     is_last: bool,
-    md5_context: md5::Context,
+    context: Context,
 }
 
 #[cfg(feature = "experimental")]
@@ -578,20 +571,16 @@ impl Hyp {
         // OPTIMIZE ME: Heap allocation done for each hyp.
         let mut framebuf = FrameBuf::with_size(stream_info.channels(), block_size);
         let offset = back_pointer.as_ref().map_or(0, |h| h.next_offset());
-        src.read_samples_from(offset, &mut framebuf)?;
+        let mut ctx = back_pointer.as_ref().map_or_else(
+            || Context::new(stream_info.bits_per_sample(), stream_info.channels()),
+            |hyp| hyp.context.clone(),
+        );
+        src.read_samples_from(offset, &mut framebuf, &mut ctx)?;
         let frame_samples = std::cmp::min(src.len() - offset, block_size);
         let frame = encode_frame(config, &framebuf, offset, stream_info);
         let bits = frame.count_bits();
         let prev_samples = back_pointer.as_ref().map_or(0, |h| h.samples);
         let prev_bits = back_pointer.as_ref().map_or(0, |h| h.bits);
-        let mut md5_context = back_pointer
-            .as_ref()
-            .map_or_else(md5::Context::new, |h| h.md5_context.clone());
-        framebuf.update_md5(
-            stream_info.bits_per_sample(),
-            frame_samples,
-            &mut md5_context,
-        );
         Ok(Arc::new(Self {
             frame,
             offset,
@@ -600,7 +589,7 @@ impl Hyp {
             bits: prev_bits + bits,
             back_pointer,
             is_last: frame_samples != block_size,
-            md5_context,
+            context: ctx,
         }))
     }