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                (&if ctx.is_little() {
213                    self.to_le_bytes()
214                } else {
215                    self.to_be_bytes()
216                }).try_into_ctx(buf, ())
217            }
218        }
219
220        impl TryIntoCtx<scroll::Endian, DynamicBuffer> for &$t {
221            type Error = scroll::Error;
222
223            fn try_into_ctx(
224                self,
225                buf: &mut DynamicBuffer,
226                ctx: Endian,
227            ) -> Result<usize, Self::Error> {
228                (*self).try_into_ctx(buf, ctx)
229            }
230        }
231    };
232}
233
234num_impl!(i8);
235num_impl!(i16);
236num_impl!(i32);
237num_impl!(i64);
238num_impl!(i128);
239
240num_impl!(u8);
241num_impl!(u16);
242num_impl!(u32);
243num_impl!(u64);
244num_impl!(u128);
245
246num_impl!(f32);
247num_impl!(f64);
248
249#[cfg(test)]
250mod tests {
251    use super::DynamicBuffer;
252    use scroll::{Endian, Pwrite, ctx::TryIntoCtx};
253
254    struct Test {
255        a: u16,
256        b: u32,
257        c: u64,
258    }
259
260    #[test]
261    fn int_write() {
262        let mut buf = DynamicBuffer::new();
263
264        buf.pwrite_with(0x1234u16, 0, Endian::Little).unwrap();
265        buf.pwrite_with(0x5678i16, 2, Endian::Big).unwrap();
266
267        assert_eq!(buf.get(), [0x34, 0x12, 0x56, 0x78]);
268    }
269
270    #[test]
271    fn offset_write() {
272        let mut buf = DynamicBuffer::new();
273
274        buf.pwrite_with(0x1234u16, 2, Endian::Big).unwrap();
275
276        assert_eq!(buf.get(), [0, 0, 0x12, 0x34]);
277    }
278
279    #[test]
280    fn slice_write() {
281        let mut buf = DynamicBuffer::new();
282
283        buf.pwrite([1u8; 4].as_slice(), 0).unwrap();
284        assert_eq!(buf.get(), [1, 1, 1, 1]);
285
286        buf.pwrite([2u8; 2].as_slice(), 2).unwrap();
287        assert_eq!(buf.get(), [1, 1, 2, 2]);
288    }
289
290    #[test]
291    fn basic_write() {
292        impl TryIntoCtx<Endian, DynamicBuffer> for Test {
293            type Error = scroll::Error;
294
295            fn try_into_ctx(
296                self,
297                buf: &mut DynamicBuffer,
298                ctx: Endian,
299            ) -> Result<usize, Self::Error> {
300                let offset = &mut 0;
301
302                buf.gwrite_with(self.a, offset, ctx)?;
303                buf.gwrite_with(self.b, offset, ctx)?;
304                buf.gwrite_with(self.c, offset, ctx)?;
305
306                Ok(*offset)
307            }
308        }
309
310        let mut buf = DynamicBuffer::new();
311
312        buf.pwrite_with(Test { a: 1, b: 2, c: 3 }, 0, Endian::Little)
313            .unwrap();
314
315        assert_eq!(buf.get(), [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0]);
316    }
317
318    #[test]
319    fn dyn_size_write() {
320        impl TryIntoCtx<(bool, Endian), DynamicBuffer> for Test {
321            type Error = scroll::Error;
322
323            fn try_into_ctx(
324                self,
325                buf: &mut DynamicBuffer,
326                (is16, ctx): (bool, Endian),
327            ) -> Result<usize, Self::Error> {
328                let offset = &mut 0;
329
330                if is16 {
331                    buf.gwrite_with(self.a, offset, ctx)?;
332                } else {
333                    buf.gwrite_with(self.a as u32, offset, ctx)?;
334                }
335
336                buf.gwrite_with(self.b, offset, ctx)?;
337                buf.gwrite_with(self.c, offset, ctx)?;
338
339                Ok(*offset)
340            }
341        }
342
343        let mut buf1 = DynamicBuffer::new();
344        buf1.pwrite_with(Test { a: 1, b: 2, c: 3 }, 0, (true, Endian::Little))
345            .unwrap();
346        assert_eq!(buf1.get(), [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0]);
347
348        let mut buf1 = DynamicBuffer::new();
349        buf1.pwrite_with(Test { a: 1, b: 2, c: 3 }, 0, (false, Endian::Little))
350            .unwrap();
351        assert_eq!(buf1.get(), [1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0]);
352    }
353
354    #[test]
355    fn alloc_size() {
356        let mut buf = DynamicBuffer::with_increment(8);
357
358        buf.pwrite_with(0xbabecafe_u32, 0, Endian::Big).unwrap();
359
360        assert_eq!(buf.buffer.len(), 8);
361        assert_eq!(buf.get(), [0xba, 0xbe, 0xca, 0xfe]);
362    }
363}
364
365#[cfg(doctest)]
366doc_comment::doctest!("../README.md");
scroll_buffer/lib.rs

scroll_buffer/
lib.rs
