bitbuffer/
writestream.rs

1use num_traits::{Float, PrimInt};
2use std::cmp::min;
3use std::mem::size_of;
4use std::ops::{BitOrAssign, BitXor};
5
6use crate::endianness::Endianness;
7use crate::num_traits::{IsSigned, SplitFitUsize, UncheckedPrimitiveFloat, UncheckedPrimitiveInt};
8use crate::writebuffer::WriteBuffer;
9use crate::{BitError, BitReadStream, BitWrite, BitWriteSized, Result};
10use std::fmt::Debug;
11
12const USIZE_SIZE: usize = size_of::<usize>();
13const USIZE_BITS: usize = USIZE_SIZE * 8;
14
15/// Stream that provides an a way to write non bit aligned adata
16///
17/// # Examples
18///
19/// ```
20/// use bitbuffer::{BitWriteStream, LittleEndian};
21/// # use bitbuffer::Result;
22///
23/// # fn main() -> Result<()> {
24/// let mut data = Vec::new();
25/// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
26///
27/// stream.write_bool(false)?;
28/// stream.write_int(123u16, 15)?;
29/// # Ok(())
30/// # }
31/// ```
32///
33/// [`BitBuffer`]: struct.BitBuffer.html
34pub struct BitWriteStream<'a, E>
35where
36    E: Endianness,
37{
38    buffer: WriteBuffer<'a, E>,
39}
40
41impl<'a, E> BitWriteStream<'a, E>
42where
43    E: Endianness,
44{
45    /// Create a new write stream
46    ///
47    /// # Examples
48    ///
49    /// ```
50    /// use bitbuffer::{BitWriteStream, LittleEndian};
51    ///
52    /// let mut data = Vec::new();
53    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
54    /// ```
55    pub fn new(data: &'a mut Vec<u8>, endianness: E) -> Self {
56        BitWriteStream {
57            buffer: WriteBuffer::new(data, endianness),
58        }
59    }
60
61    /// Create a new write stream
62    ///
63    /// Note that the resulting stream will panic when trying to write more data then fits
64    /// in the provided slice.
65    pub fn from_slice(data: &'a mut [u8], endianness: E) -> Self {
66        BitWriteStream {
67            buffer: WriteBuffer::for_slice(data, endianness),
68        }
69    }
70}
71
72impl<E> BitWriteStream<'_, E>
73where
74    E: Endianness,
75{
76    /// The number of written bits in the buffer
77    pub fn bit_len(&self) -> usize {
78        self.buffer.bit_len()
79    }
80
81    /// The number of written bytes in the buffer
82    pub fn byte_len(&self) -> usize {
83        self.buffer.bit_len().div_ceil(8)
84    }
85
86    fn push_non_fit_bits<I>(&mut self, bits: I, count: usize)
87    where
88        I: ExactSizeIterator,
89        I: DoubleEndedIterator<Item = (usize, u8)>,
90    {
91        self.buffer.push_non_fit_bits(bits, count)
92    }
93
94    /// Push up to an usize worth of bits
95    fn push_bits(&mut self, bits: usize, count: usize) {
96        if count > 0 {
97            self.buffer.push_bits(bits, count)
98        }
99    }
100
101    /// Align the stream on the next byte by writing zero bits and returns the amount of bits written
102    ///
103    /// # Examples
104    ///
105    /// ```
106    /// # use bitbuffer::{BitReadBuffer, LittleEndian, Result};
107    /// #
108    /// # fn main() -> Result<()> {
109    /// # use bitbuffer::{BitWriteStream, LittleEndian};
110    ///
111    /// let mut data = Vec::new();
112    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
113    /// stream.write_bool(true)?;
114    /// stream.align();
115    /// assert_eq!(stream.bit_len(), 8);
116    /// assert_eq!(data, [0b0000_0001]);
117    /// #
118    /// #     Ok(())
119    /// # }
120    /// ```
121    ///
122    /// [`ReadError::NotEnoughData`]: enum.ReadError.html#variant.NotEnoughData
123    pub fn align(&mut self) -> usize {
124        match self.bit_len() % 8 {
125            0 => 0,
126            n => {
127                self.push_bits(0, 8 - n);
128                8 - n
129            }
130        }
131    }
132
133    /// Write a boolean into the buffer
134    ///
135    /// # Examples
136    ///
137    /// ```
138    /// # use bitbuffer::{BitReadBuffer, LittleEndian, Result};
139    /// #
140    /// # fn main() -> Result<()> {
141    /// # use bitbuffer::{BitWriteStream, LittleEndian};
142    ///
143    /// let mut data = Vec::new();
144    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
145    /// stream.write_bool(true)?;
146    /// #
147    /// #     Ok(())
148    /// # }
149    /// ```
150    #[inline]
151    pub fn write_bool(&mut self, value: bool) -> Result<()> {
152        self.buffer.push_bool(value);
153        Ok(())
154    }
155
156    /// Write an integer into the buffer
157    ///
158    /// # Examples
159    ///
160    /// ```
161    /// # use bitbuffer::{BitReadBuffer, LittleEndian, Result};
162    /// #
163    /// # fn main() -> Result<()> {
164    /// # use bitbuffer::{BitWriteStream, LittleEndian};
165    ///
166    /// let mut data = Vec::new();
167    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
168    /// stream.write_int(123u16, 15)?;
169    /// #
170    /// #     Ok(())
171    /// # }
172    /// ```
173    #[inline]
174    pub fn write_int<T>(&mut self, value: T, count: usize) -> Result<()>
175    where
176        T: PrimInt
177            + BitOrAssign
178            + IsSigned
179            + UncheckedPrimitiveInt
180            + BitXor
181            + Debug
182            + SplitFitUsize,
183    {
184        let type_bit_size = size_of::<T>() * 8;
185
186        if type_bit_size < count {
187            return Err(BitError::TooManyBits {
188                requested: count,
189                max: type_bit_size,
190            });
191        }
192
193        if type_bit_size < USIZE_BITS || count <= (USIZE_BITS - (self.bit_len() % 8)) {
194            self.push_bits(value.into_usize_unchecked(), count);
195        } else {
196            self.push_non_fit_bits(value.split_fit_usize::<E>(count as u8), count)
197        }
198
199        Ok(())
200    }
201
202    /// Write a float into the buffer
203    ///
204    /// # Examples
205    ///
206    /// ```
207    /// # use bitbuffer::{BitReadBuffer, LittleEndian, Result};
208    /// #
209    /// # fn main() -> Result<()> {
210    /// # use bitbuffer::{BitWriteStream, LittleEndian};
211    ///
212    /// let mut data = Vec::new();
213    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
214    /// stream.write_float(123.15f32)?;
215    /// #
216    /// #     Ok(())
217    /// # }
218    /// ```
219    #[inline]
220    pub fn write_float<T>(&mut self, value: T) -> Result<()>
221    where
222        T: Float + UncheckedPrimitiveFloat,
223    {
224        if self.buffer.bit_len() & 7 == 0 {
225            let bytes = value.to_bytes::<E>();
226            self.buffer.extends_from_slice(bytes.as_ref());
227        } else {
228            self.write_int(value.to_int(), size_of::<T>() * 8)?;
229        }
230
231        Ok(())
232    }
233
234    /// Write a number of bytes into the buffer
235    ///
236    /// # Examples
237    ///
238    /// ```
239    /// # use bitbuffer::{BitReadBuffer, LittleEndian, Result};
240    /// #
241    /// # fn main() -> Result<()> {
242    /// # use bitbuffer::{BitWriteStream, LittleEndian};
243    ///
244    /// let mut data = Vec::new();
245    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
246    /// stream.write_bytes(&[0, 1, 2 ,3])?;
247    /// #
248    /// #     Ok(())
249    /// # }
250    /// ```
251    #[inline]
252    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
253        if self.buffer.bit_len() & 7 == 0 {
254            self.buffer.extends_from_slice(bytes);
255        } else {
256            bytes
257                .iter()
258                .copied()
259                .for_each(|chunk| self.push_bits(chunk as usize, 8));
260        }
261        Ok(())
262    }
263
264    /// Write bits from a read stream into the buffer
265    #[inline]
266    pub fn write_bits(&mut self, bits: &BitReadStream<E>) -> Result<()> {
267        let mut bits = bits.clone();
268        let bit_offset = self.bit_len() % 8;
269        if bit_offset > 0 {
270            let bit_count = min(8 - bit_offset, bits.bits_left());
271            let start = bits.read_int::<u8>(bit_count)?;
272            self.push_bits(start as usize, bit_count);
273        }
274
275        while bits.bits_left() > 32 {
276            let chunk = bits.read::<u32>()?;
277            self.push_bits(chunk as usize, 32);
278        }
279
280        if bits.bits_left() > 0 {
281            let end_bits = bits.bits_left();
282            let end = bits.read_int::<u32>(end_bits)?;
283            self.push_bits(end as usize, end_bits);
284        }
285        Ok(())
286    }
287
288    /// Write a string into the buffer
289    ///
290    /// # Examples
291    ///
292    /// ```
293    /// # use bitbuffer::{BitReadBuffer, LittleEndian, Result};
294    /// #
295    /// # fn main() -> Result<()> {
296    /// # use bitbuffer::{BitWriteStream, LittleEndian};
297    ///
298    /// let mut data = Vec::new();
299    /// let mut stream = BitWriteStream::new(&mut data, LittleEndian);
300    /// stream.write_string("zero terminated string", None)?;
301    /// stream.write_string("fixed size string, zero padded", Some(64))?;
302    /// #
303    /// #     Ok(())
304    /// # }
305    /// ```
306    pub fn write_string(&mut self, string: &str, length: Option<usize>) -> Result<()> {
307        match length {
308            Some(length) => {
309                if length < string.len() {
310                    return Err(BitError::StringToLong {
311                        string_length: string.len(),
312                        requested_length: length,
313                    });
314                }
315                self.write_bytes(string.as_bytes())?;
316                for _ in 0..(length - string.len()) {
317                    self.push_bits(0, 8)
318                }
319            }
320            None => {
321                self.write_bytes(string.as_bytes())?;
322                self.push_bits(0, 8)
323            }
324        }
325        Ok(())
326    }
327
328    /// Write the type to stream
329    #[inline]
330    pub fn write<T: BitWrite<E>>(&mut self, value: &T) -> Result<()> {
331        value.write(self)
332    }
333
334    /// Write the type to stream
335    #[inline]
336    pub fn write_sized<T: BitWriteSized<E>>(&mut self, value: &T, length: usize) -> Result<()> {
337        value.write_sized(self, length)
338    }
339
340    /// Write the length of a section before the section
341    pub fn reserve_length<Err: From<BitError>, F: Fn(&mut BitWriteStream<E>) -> Result<(), Err>>(
342        &mut self,
343        length_bit_size: usize,
344        body_fn: F,
345    ) -> Result<(), Err> {
346        self.reserve_int(length_bit_size, |stream| {
347            let start = stream.bit_len();
348            body_fn(stream)?;
349            let end = stream.bit_len();
350            Ok((end - start) as u64)
351        })
352    }
353
354    /// Write the length in bytes of a section before the section, the section will be 0 padded to an even byte length
355    pub fn reserve_byte_length<
356        Err: From<BitError>,
357        F: Fn(&mut BitWriteStream<E>) -> Result<(), Err>,
358    >(
359        &mut self,
360        length_bit_size: usize,
361        body_fn: F,
362    ) -> Result<(), Err> {
363        self.reserve_int(length_bit_size, |stream| {
364            let start = stream.bit_len();
365            body_fn(stream)?;
366            let end = stream.bit_len();
367            let bit_len = end - start;
368
369            let pad_len = (8 - (bit_len & 7)) & 7;
370            stream.push_bits(0, pad_len);
371
372            let byte_len = (bit_len + pad_len) / 8;
373            Ok(byte_len as u64)
374        })
375    }
376
377    /// Reserve the length to write an integer
378    pub fn reserve_int<Err: From<BitError>, F: Fn(&mut BitWriteStream<E>) -> Result<u64, Err>>(
379        &mut self,
380        count: usize,
381        body_fn: F,
382    ) -> Result<(), Err> {
383        let start = self.bit_len();
384        self.write_int(0u64, count)?;
385
386        let head_int = body_fn(self)?;
387        self.buffer.set_at(start, head_int, count);
388
389        Ok(())
390    }
391}
bitbuffer/writestream.rs

bitbuffer/
writestream.rs
