scroll_buffer/
lib.rs

1use std::fmt::{Debug, Formatter};
2use scroll::{Endian, Pwrite, ctx::TryIntoCtx};
3use std::ops::*;
4
5/// A buffer designed for `Pwrite` that dynamically expands to hold all types written into it.
6/// Uses a standard `Vec` under the hood.
7///
8/// # Use Cases
9///
10/// * You don't know exactly how much you're writing ahead of time.
11/// * Your types change size depending on context parameters or specific values.
12///
13/// # Allocations
14///
15/// By default, this buffer expands every time one byte is written. The `with_increment`
16/// constructor can be used to change how much the buffer expands to reduce overall number of allocations.
17///
18/// # Pwrite and TryIntoCtx
19///
20/// Types must implement `TryIntoCtx` with a second generic parameter of `DynamicBuffer` to be used.
21///
22/// ```rust
23/// use scroll::{ctx::TryIntoCtx, Error, LE, Pwrite};
24/// use scroll_buffer::DynamicBuffer;
25///
26/// struct EvenOdd(u32);
27///
28/// impl TryIntoCtx<(), DynamicBuffer> for EvenOdd {
29///     type Error = Error;
30///
31///     fn try_into_ctx(self, buf: &mut DynamicBuffer, _: ()) -> Result<usize, Self::Error> {
32///         let offset = &mut 0;
33///
34///         if self.0 % 2 == 0 {
35///             buf.gwrite_with(self.0, offset, LE)?;
36///         } else {
37///             buf.gwrite_with(self.0 as u64, offset, LE)?;
38///         }
39///
40///         Ok(*offset)
41///     }
42/// }
43///
44/// let mut buf = DynamicBuffer::new();
45///
46/// let offset = buf.pwrite(EvenOdd(2), 0).unwrap();
47/// assert_eq!(buf.get(), [2, 0, 0, 0]);
48/// buf.pwrite(EvenOdd(3), offset).unwrap();
49/// assert_eq!(buf.get()[offset..], [3, 0, 0, 0, 0, 0, 0, 0]);
50/// ```
51///
52/// ## Why can't I use types that already implement the default TryIntoCtx?
53///
54/// The default type parameters for `TryIntoCtx` expose a mutable raw slice to write into.
55/// This is great for pre-allocated buffers and types that have a static written size; however,
56/// we cannot create a custom slice type that performs dynamic expansions under the hood.
57///
58/// As a result, we need to expose a separate writing type that can track offsets and calculate when to allocate.
59///
60/// However, if your `TryIntoCtx` impls don't use any special slice APIs and just use `Pwrite` and/or
61/// basic indexing, it's extremely easy to migrate! Just add the `DynamicBuffer` generic type.
62pub struct DynamicBuffer {
63    pub(crate) buffer: Vec<u8>,
64    alloc_increment: usize,
65    pub(crate) start_offset: usize,
66    write_end: usize,
67}
68
69impl DynamicBuffer {
70    /// Constructs an empty `DynamicBuffer` with the default allocation increment of one byte.
71    pub fn new() -> DynamicBuffer {
72        Self::with_increment(1)
73    }
74
75    /// Constructs an empty `DynamicBuffer` with a custom allocation increment.
76    pub fn with_increment(alloc_increment: usize) -> DynamicBuffer {
77        DynamicBuffer {
78            buffer: vec![],
79            alloc_increment,
80            start_offset: 0,
81            write_end: 0,
82        }
83    }
84
85    /// Gets a slice of all written bytes. Does not include extra bytes allocated by a large increment.
86    /// ```rust
87    /// use scroll::{LE, Pwrite};
88    /// use scroll_buffer::DynamicBuffer;
89    ///
90    /// let mut buf = DynamicBuffer::new();
91    ///
92    /// assert_eq!(buf.get(), []);
93    /// buf.pwrite_with(2u32, 0, LE).unwrap();
94    /// assert_eq!(buf.get(), [2, 0, 0, 0]);
95    /// ```
96    pub fn get(&self) -> &[u8] {
97        &self.buffer[..self.write_end]
98    }
99
100    /// Consumes the buffer and returns everything written to it, discarding any excess allocated zeroes.
101    /// ```rust
102    /// use scroll::{LE, Pwrite};
103    /// use scroll_buffer::DynamicBuffer;
104    ///
105    /// let mut buf = DynamicBuffer::with_increment(8);
106    /// buf.pwrite_with(2u32, 0, LE).unwrap();
107    ///
108    /// assert_eq!(buf.into_vec(), vec![2, 0, 0, 0]);
109    /// ```
110    pub fn into_vec(mut self) -> Vec<u8> {
111        self.buffer.truncate(self.write_end);
112        self.buffer
113    }
114
115    /// Resets the buffer's contents. Maintains already allocated capacity.
116    /// ```rust
117    /// use scroll::{LE, Pwrite};
118    /// use scroll_buffer::DynamicBuffer;
119    ///
120    /// let mut buf = DynamicBuffer::new();
121    ///
122    /// buf.pwrite_with(2u16, 0, LE).unwrap();
123    /// assert_eq!(buf.get(), [2, 0]);
124    ///
125    /// buf.clear();
126    /// assert_eq!(buf.get(), []);
127    ///
128    /// // does not reallocate!
129    /// buf.pwrite_with(0xffu8, 0, LE).unwrap();
130    /// assert_eq!(buf.get(), [0xff]);
131    /// ```
132    pub fn clear(&mut self) {
133        self.buffer.clear();
134        self.start_offset = 0;
135        self.write_end = 0;
136    }
137}
138
139impl Debug for DynamicBuffer {
140    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
141        Debug::fmt(self.get(), f)
142    }
143}
144
145impl Index<usize> for DynamicBuffer {
146    type Output = u8;
147
148    fn index(&self, index: usize) -> &Self::Output {
149        self.buffer.index(self.start_offset + index)
150    }
151}
152impl IndexMut<usize> for DynamicBuffer {
153    fn index_mut(&mut self, mut index: usize) -> &mut Self::Output {
154        index += self.start_offset;
155
156        if let Some(diff) = index.checked_sub(self.buffer.len()) {
157            self.buffer
158                .extend(std::iter::repeat(0).take(std::cmp::max(self.alloc_increment, diff + 1)));
159        }
160
161        self.write_end = self.write_end.max(index + 1);
162
163        self.buffer.index_mut(index)
164    }
165}
166
167impl<Ctx: Copy, E> Pwrite<Ctx, E> for DynamicBuffer {
168    fn pwrite_with<N: TryIntoCtx<Ctx, Self, Error = E>>(
169        &mut self,
170        n: N,
171        offset: usize,
172        ctx: Ctx,
173    ) -> Result<usize, E> {
174        self.start_offset += offset;
175
176        let end = n.try_into_ctx(self, ctx)?;
177
178        self.start_offset -= offset;
179
180        Ok(end)
181    }
182}
183
184impl TryIntoCtx<(), DynamicBuffer> for &'_ [u8] {
185    type Error = scroll::Error; // doesn't matter, this is infallible
186
187    fn try_into_ctx(self, into: &mut DynamicBuffer, _: ()) -> Result<usize, Self::Error> {
188        let len = self.len();
189
190        // set final byte to 0
191        // if the buffer is not large enough, this will automatically allocate through IndexMut
192        into[len - 1] = 0;
193
194        // the following range is now guaranteed to be valid
195        // TODO: should consumers have this kind of IndexMut<Range> access?
196        into.buffer[into.start_offset..into.start_offset + len].copy_from_slice(self);
197
198        Ok(len)
199    }
200}
201
202macro_rules! num_impl {
203    ($t:ty) => {
204        impl TryIntoCtx<scroll::Endian, DynamicBuffer> for $t {
205            type Error = scroll::Error; // also infallible
206
207            fn try_into_ctx(
208                self,
209                buf: &mut DynamicBuffer,
210                ctx: Endian,
211            ) -> Result<usize, Self::Error> {
212                let bytes = if ctx.is_little() {
213                    self.to_le_bytes()
214                } else {
215                    self.to_be_bytes()
216                };
217                TryIntoCtx::try_into_ctx(&bytes[..], buf, ())
218            }
219        }
220
221        impl TryIntoCtx<scroll::Endian, DynamicBuffer> for &$t {
222            type Error = scroll::Error;
223
224            fn try_into_ctx(
225                self,
226                buf: &mut DynamicBuffer,
227                ctx: Endian,
228            ) -> Result<usize, Self::Error> {
229                (*self).try_into_ctx(buf, ctx)
230            }
231        }
232    };
233}
234
235num_impl!(i8);
236num_impl!(i16);
237num_impl!(i32);
238num_impl!(i64);
239num_impl!(i128);
240
241num_impl!(u8);
242num_impl!(u16);
243num_impl!(u32);
244num_impl!(u64);
245num_impl!(u128);
246
247num_impl!(f32);
248num_impl!(f64);
249
250#[cfg(test)]
251mod tests {
252    use super::DynamicBuffer;
253    use scroll::{Endian, Pwrite, ctx::TryIntoCtx};
254
255    struct Test {
256        a: u16,
257        b: u32,
258        c: u64,
259    }
260
261    #[test]
262    fn int_write() {
263        let mut buf = DynamicBuffer::new();
264
265        buf.pwrite_with(0x1234u16, 0, Endian::Little).unwrap();
266        buf.pwrite_with(0x5678i16, 2, Endian::Big).unwrap();
267
268        assert_eq!(buf.get(), [0x34, 0x12, 0x56, 0x78]);
269    }
270
271    #[test]
272    fn offset_write() {
273        let mut buf = DynamicBuffer::new();
274
275        buf.pwrite_with(0x1234u16, 2, Endian::Big).unwrap();
276
277        assert_eq!(buf.get(), [0, 0, 0x12, 0x34]);
278    }
279
280    #[test]
281    fn slice_write() {
282        let mut buf = DynamicBuffer::new();
283
284        buf.pwrite([1u8; 4].as_slice(), 0).unwrap();
285        assert_eq!(buf.get(), [1, 1, 1, 1]);
286
287        buf.pwrite([2u8; 2].as_slice(), 2).unwrap();
288        assert_eq!(buf.get(), [1, 1, 2, 2]);
289    }
290
291    #[test]
292    fn basic_write() {
293        impl TryIntoCtx<Endian, DynamicBuffer> for Test {
294            type Error = scroll::Error;
295
296            fn try_into_ctx(
297                self,
298                buf: &mut DynamicBuffer,
299                ctx: Endian,
300            ) -> Result<usize, Self::Error> {
301                let offset = &mut 0;
302
303                buf.gwrite_with(self.a, offset, ctx)?;
304                buf.gwrite_with(self.b, offset, ctx)?;
305                buf.gwrite_with(self.c, offset, ctx)?;
306
307                Ok(*offset)
308            }
309        }
310
311        let mut buf = DynamicBuffer::new();
312
313        buf.pwrite_with(Test { a: 1, b: 2, c: 3 }, 0, Endian::Little)
314            .unwrap();
315
316        assert_eq!(buf.get(), [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0]);
317    }
318
319    #[test]
320    fn dyn_size_write() {
321        impl TryIntoCtx<(bool, Endian), DynamicBuffer> for Test {
322            type Error = scroll::Error;
323
324            fn try_into_ctx(
325                self,
326                buf: &mut DynamicBuffer,
327                (is16, ctx): (bool, Endian),
328            ) -> Result<usize, Self::Error> {
329                let offset = &mut 0;
330
331                if is16 {
332                    buf.gwrite_with(self.a, offset, ctx)?;
333                } else {
334                    buf.gwrite_with(self.a as u32, offset, ctx)?;
335                }
336
337                buf.gwrite_with(self.b, offset, ctx)?;
338                buf.gwrite_with(self.c, offset, ctx)?;
339
340                Ok(*offset)
341            }
342        }
343
344        let mut buf1 = DynamicBuffer::new();
345        buf1.pwrite_with(Test { a: 1, b: 2, c: 3 }, 0, (true, Endian::Little))
346            .unwrap();
347        assert_eq!(buf1.get(), [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0]);
348
349        let mut buf1 = DynamicBuffer::new();
350        buf1.pwrite_with(Test { a: 1, b: 2, c: 3 }, 0, (false, Endian::Little))
351            .unwrap();
352        assert_eq!(buf1.get(), [1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0]);
353    }
354
355    #[test]
356    fn alloc_size() {
357        let mut buf = DynamicBuffer::with_increment(8);
358
359        buf.pwrite_with(0xbabecafe_u32, 0, Endian::Big).unwrap();
360
361        assert_eq!(buf.buffer.len(), 8);
362        assert_eq!(buf.get(), [0xba, 0xbe, 0xca, 0xfe]);
363    }
364}
365
366#[cfg(doctest)]
367doc_comment::doctest!("../README.md");
scroll_buffer/lib.rs

scroll_buffer/
lib.rs
