BitWriteStream::from_slice

src/writestream.rs

@@ -2,6 +2,7 @@ use num_traits::{Float, PrimInt};
 use std::cmp::min;
 use std::mem::size_of;
 use std::ops::{BitOrAssign, BitXor};
+use std::rc::Rc;
 
 use crate::endianness::Endianness;
 use crate::num_traits::{IsSigned, SplitFitUsize, UncheckedPrimitiveFloat, UncheckedPrimitiveInt};
@@ -36,6 +37,7 @@ where
     E: Endianness,
 {
     buffer: WriteBuffer<'a, E>,
+    slice: Option<Rc<Vec<u8>>>,
 }
 
 impl<'a, E> BitWriteStream<'a, E>
@@ -55,6 +57,33 @@ where
     pub fn new(data: &'a mut Vec<u8>, endianness: E) -> Self {
         BitWriteStream {
             buffer: WriteBuffer::new(data, endianness),
+            slice: None,
+        }
+    }
+    /// Create a new write stream
+    ///
+    /// must assert data.len() is enough
+    pub unsafe fn from_slice(data: &'a mut [u8], endianness: E) -> Self {
+        let p = data.as_mut_ptr();
+        let cap = data.len();
+        let slice = Rc::new(Vec::from_raw_parts(p, 0, cap));
+        let data = &mut *(Rc::as_ptr(&slice) as *mut Vec<u8>);
+        BitWriteStream {
+            buffer: WriteBuffer::new(data, endianness),
+            slice: Some(slice),
+        }
+    }
+}
+
+impl<'a, E> Drop for BitWriteStream<'a, E>
+where
+    E: Endianness,
+{
+    fn drop(&mut self) {
+        if let Some(slice) = self.slice.take() {
+            let slice = Rc::try_unwrap(slice).unwrap();
+            // don't memory free
+            core::mem::forget(slice);
         }
     }
 }

tests/write_tests.rs

@@ -204,3 +204,24 @@ fn test_write_container() {
     assert_eq!(Rc::new(true), read.read().unwrap());
     assert_eq!(Arc::new(true), read.read().unwrap());
 }
+
+#[test]
+fn test_write_to_slice() {
+    let mut data = [0; 32];
+    {
+        let mut stream = unsafe { BitWriteStream::from_slice(&mut data[..], LittleEndian) };
+
+        stream.write_bool(true).unwrap();
+        stream.write_int(3253u16, 16).unwrap();
+        stream.write_int(13253u64, 64).unwrap();
+    }
+
+    let mut read = BitReadStream::from(BitReadBuffer::new(&data[..], LittleEndian));
+
+    assert_eq!(true, read.read_bool().unwrap());
+    assert_eq!(3253u16, read.read::<u16>().unwrap());
+    assert_eq!(13253u64, read.read::<u64>().unwrap());
+
+    // 0 padded
+    assert_eq!(false, read.read_bool().unwrap());
+}