Merge pull request #2 from Danaozhong/Task/SupportWritingFullBytes

Add a write function to write full bytes

src/lib.rs

@@ -87,6 +87,12 @@ impl BitWriter {
     }
 
     pub fn align(&mut self, alignment_bytes: u32) -> Result<()> {
+        if alignment_bytes == 0 {
+            return Err(Error::new(
+                ErrorKind::InvalidInput,
+                "cannot align to 0 bytes",
+            ));
+        }
         let alignment_bits = alignment_bytes as u64 * 8;
         let cur_alignment = self.bit_count % alignment_bits;
         let bits_to_skip = (alignment_bits - cur_alignment) % alignment_bits;
@@ -155,6 +161,32 @@ impl BitWriter {
         Ok(())
     }
 
+    /// Writes a number of full bytes into the stream.
+    /// This will give best performance if the data is aligned
+    /// to a byte boundary. Otherwise, each data bytes will be
+    /// spread across two bytes in the bitstream.
+    /// Byte boundary can be ensured by using align().
+    pub fn write(&mut self, data: &Vec<u8>) -> Result<()> {
+        // If the data is byte-aligned, we can directly
+        // copy the bytes without performance penalty.
+        if self.bits == 0 {
+            // We are writing full bytes, so there is no
+            // need to update the bit count.
+            self.bit_count += 8 * data.len() as u64;
+            self.data.extend(data);
+            return Ok(());
+        }
+        // Since the buffer is not aligned, we need to
+        // bit-shift each byte.
+        for byte in data {
+            match self.write_u8(*byte, 8) {
+                Ok(()) => (),
+                Err(e) => return Err(e),
+            }
+        }
+        Ok(())
+    }
+
     pub fn close(&mut self) -> Result<()> {
         // align to the next byte boundary
         self.align(1)?;

src/tests.rs

@@ -1,4 +1,5 @@
 use super::*;
+use std::vec;
 
 #[test]
 fn simple_writing() {
@@ -31,3 +32,32 @@ fn test_bitshift_overflow() {
     assert_eq!(writer.bit_count, 3 * 64);
     writer.close().expect("failed to close byte vector");
 }
+
+#[test]
+fn test_byte_writing() {
+    let mut writer = BitWriter::new();
+
+    // First, test writing bytes aligned
+    writer
+        .write(&vec![0xFF, 0x22, 0x00, 0x12])
+        .expect("failed to write bytes aligned");
+    assert_eq!(writer.bit_count, 4 * 8);
+
+    // Make the buffer unaligned
+    writer.write_bool(true).expect("failed to write boolean");
+    assert_eq!(writer.bit_count, 4 * 8 + 1);
+    writer
+        .write(&vec![0xFF, 0x22, 0x00, 0x12])
+        .expect("failed to write bytes unaligned");
+    assert_eq!(writer.bit_count, 8 * 8 + 1);
+
+    // Align to byte boundary again
+    writer.align(1).expect("failed to align bit stream");
+    assert_eq!(writer.bit_count, 9 * 8);
+    writer
+        .write(&vec![0xFF])
+        .expect("failed to write bytes aligned");
+    assert_eq!(writer.bit_count, 10 * 8);
+
+    writer.close().expect("failed to close byte vector");
+}