load_buffer/
lib.rs

1//!
2//! ```no_run
3//! # use load_buffer::*;
4//! # #[derive(Default)]
5//! # struct Loader;
6//! # impl Load for Loader {
7//! #   type Error = core::convert::Infallible;
8//! #   fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error> { unimplemented!() }
9//! # }
10//! // Static, backed by the array size you give it:
11//! let mut loader: BufferedLoader<[u8; 64], _> =
12//!     BufferedLoader::new_static(Loader::default());
13//!
14//! // Dynamic, backed by alloc::vec::Vec<u8>:
15//! let mut loader = BufferedLoader::new_dynamic(Loader::default(), 64);
16//!
17//! // Dynamic can be resized.  This invalidates the buffer.
18//! loader.resize(268435456);
19//!
20//! // Heap, backed by alloc::boxed::Box<[u8]>:
21//! let mut loader = BufferedLoader::new_heap(Loader::default(), 268435456);
22//!
23//! // Heap can not be resized, but will not fill the stack for huge sizes.
24//! // You can also supply your own buffer type, as long as it implements
25//! // AsRef<[u8]> + AsMut<[u8]> + Debug.
26//! ```
27#![cfg_attr(not(feature = "std"), no_std)]
28
29#[cfg(feature = "alloc")]
30extern crate alloc;
31
32#[cfg(feature = "alloc")]
33use alloc::vec;
34
35#[cfg(feature = "alloc")]
36use alloc::vec::Vec;
37
38#[cfg(feature = "alloc")]
39use alloc::boxed::Box;
40
41use core::fmt::Debug;
42
43/// Like [`std::io::Read`], but only supporting [`read_exact`], and with an
44/// arbitrary error type.
45///
46/// [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
47/// [`read_exact`]: https://doc.rust-lang.org/std/io/trait.Read.html#method.read_exact
48pub trait Load {
49    type Error;
50
51    fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error>;
52}
53
54#[cfg(feature = "std")]
55impl<T> Load for T
56where
57    T: std::io::Read,
58{
59    type Error = std::io::Error;
60    fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error> {
61        self.read_exact(destination)
62    }
63}
64
65/// A buffering wrapper around a [`Load`].
66///
67/// This fills some of the simpler niches of a [`std::io::BufReader`] for no_std
68/// use-cases, and allowing a static buffer.
69///
70/// [`std::io::BufReader`]: https://doc.rust-lang.org/std/io/struct.BufReader.html
71/// [`Load`]: trait.Load.html
72#[derive(Clone, Debug)]
73pub struct BufferedLoader<Buffer, Loader>
74where
75    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
76    Loader: Load,
77{
78    loader: Loader,
79    buffer: Buffer,
80
81    /// The amount of buffer that has been read already.
82    offset: usize,
83}
84
85impl<const N: usize, Loader> BufferedLoader<[u8; N], Loader>
86where
87    Loader: Load,
88{
89    /// Creates a new stack-based BufferedLoader.  Does not allocate.
90    pub fn new_static(loader: Loader) -> Self {
91        Self::new([0u8; N], loader)
92    }
93}
94
95#[cfg(feature = "alloc")]
96impl<Loader> BufferedLoader<Vec<u8>, Loader>
97where
98    Loader: Load,
99{
100    /// Creates a new heap-based BufferedLoader.  Allocates only once.
101    pub fn new_dynamic(loader: Loader, buffer_size: usize) -> Self {
102        Self::new(vec![0u8; buffer_size], loader)
103    }
104
105    pub fn resize(&mut self, new_size: usize) {
106        self.buffer.resize(new_size, 0);
107        self.offset = new_size;
108    }
109}
110
111#[cfg(feature = "alloc")]
112impl<Loader> BufferedLoader<Box<[u8]>, Loader>
113where
114    Loader: Load,
115{
116    /// Creates a new heap-based BufferedLoader.  Allocates only once.
117    pub fn new_heap(loader: Loader, buffer_size: usize) -> Self {
118        Self::new(vec![0u8; buffer_size].into_boxed_slice(), loader)
119    }
120}
121
122impl<Buffer, Loader> BufferedLoader<Buffer, Loader>
123where
124    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
125    Loader: Load,
126{
127    pub fn new(buffer: Buffer, loader: Loader) -> Self {
128        let offset = buffer.as_ref().len();
129        Self {
130            loader,
131            buffer,
132            offset,
133        }
134    }
135
136    /// Drop what's remaining in the buffer and give a mutable reference to the
137    /// inner loader, so it can be seeked or otherwise manipulated.
138    pub fn loader(&mut self) -> &mut Loader {
139        self.offset = self.buffer.as_ref().len();
140        &mut self.loader
141    }
142}
143
144impl<Buffer, Loader> Load for BufferedLoader<Buffer, Loader>
145where
146    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
147    Loader: Load,
148{
149    type Error = Loader::Error;
150
151    /// Efficiently fill the destination buffer, calling the underlying
152    /// [`Load::load`] as few times as possible.
153    ///
154    /// [`Load::Load`]: trait.Load.html#method.load
155    fn load(&mut self, mut destination: &mut [u8]) -> Result<(), Self::Error> {
156        if destination.is_empty() {
157            return Ok(());
158        }
159
160        let buffer = self.buffer.as_mut();
161        let buffer_size = buffer.len();
162
163        let buffer_remaining = buffer_size - self.offset;
164
165        if destination.len() <= buffer_remaining {
166            // There are enough bytes left in the buffer to consume without
167            // reading.
168            let end = self.offset + destination.len();
169            destination.copy_from_slice(&buffer[self.offset..end]);
170            self.offset = end;
171        } else {
172            // Empty what's remaining in the buffer.
173            if buffer_remaining > 0 {
174                destination[..buffer_remaining].copy_from_slice(&buffer[self.offset..buffer_size]);
175                let copied = buffer_size - self.offset;
176                destination = &mut destination[copied..];
177            }
178
179            let buffers = destination.len() / buffer_size;
180            let remainder = destination.len() % buffer_size;
181
182            // Load full-sized chunks directly to the destination, bypassing
183            // the buffer.
184            if buffers > 0 {
185                let buffers_bytes = buffers * buffer_size;
186                self.loader.load(&mut destination[..buffers_bytes])?;
187
188                destination = &mut destination[buffers_bytes..];
189            }
190
191            // Fill the buffer for the remainder, if there is any.
192            if remainder > 0 {
193                self.loader.load(buffer)?;
194                destination.copy_from_slice(&buffer[..remainder]);
195                self.offset = remainder;
196            } else {
197                // We have emptied the remaining buffer, so mark this empty.
198                self.offset = buffer_size;
199            }
200        }
201        Ok(())
202    }
203}
204
205#[cfg(test)]
206pub mod test {
207    use core::convert::Infallible;
208
209    use super::*;
210
211    #[derive(Default)]
212    struct Increment(u8);
213
214    impl Load for Increment {
215        type Error = Infallible;
216
217        fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error> {
218            for each in destination {
219                *each = self.0;
220                self.0 = self.0.wrapping_add(1);
221            }
222            Ok(())
223        }
224    }
225
226    #[test]
227    fn test_static_buffered_loader() {
228        let mut loader: BufferedLoader<[u8; 8], _> =
229            BufferedLoader::new_static(Increment::default());
230        let mut buf = [0u8; 4];
231        loader.load(&mut buf).unwrap();
232        assert_eq!(buf, [0, 1, 2, 3]);
233        loader.load(&mut buf).unwrap();
234        assert_eq!(buf, [4, 5, 6, 7]);
235        loader.load(&mut buf).unwrap();
236        assert_eq!(buf, [8, 9, 10, 11]);
237
238        let mut buf = [0u8; 12];
239        loader.load(&mut buf).unwrap();
240        assert_eq!(buf, [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);
241
242        let mut buf = [0u8; 3];
243
244        loader.load(&mut buf).unwrap();
245        assert_eq!(buf, [24, 25, 26]);
246        loader.load(&mut buf).unwrap();
247        assert_eq!(buf, [27, 28, 29]);
248        loader.load(&mut buf).unwrap();
249        assert_eq!(buf, [30, 31, 32]);
250        loader.load(&mut buf).unwrap();
251        assert_eq!(buf, [33, 34, 35]);
252        loader.load(&mut buf).unwrap();
253        assert_eq!(buf, [36, 37, 38]);
254        loader.load(&mut buf).unwrap();
255        assert_eq!(buf, [39, 40, 41]);
256        loader.load(&mut buf).unwrap();
257        assert_eq!(buf, [42, 43, 44]);
258        loader.load(&mut buf).unwrap();
259        assert_eq!(buf, [45, 46, 47]);
260        loader.load(&mut buf).unwrap();
261        assert_eq!(buf, [48, 49, 50]);
262    }
263
264    #[test]
265    fn test_dynamic_buffered_loader() {
266        let mut loader = BufferedLoader::new_dynamic(Increment::default(), 8);
267        let mut buf = [0u8; 4];
268        loader.load(&mut buf).unwrap();
269        assert_eq!(buf, [0, 1, 2, 3]);
270        loader.load(&mut buf).unwrap();
271        assert_eq!(buf, [4, 5, 6, 7]);
272        loader.load(&mut buf).unwrap();
273        assert_eq!(buf, [8, 9, 10, 11]);
274
275        let mut buf = [0u8; 12];
276        loader.load(&mut buf).unwrap();
277        assert_eq!(buf, [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);
278
279        let mut buf = [0u8; 3];
280
281        loader.load(&mut buf).unwrap();
282        assert_eq!(buf, [24, 25, 26]);
283        loader.load(&mut buf).unwrap();
284        assert_eq!(buf, [27, 28, 29]);
285        loader.load(&mut buf).unwrap();
286        assert_eq!(buf, [30, 31, 32]);
287        loader.load(&mut buf).unwrap();
288        assert_eq!(buf, [33, 34, 35]);
289        loader.load(&mut buf).unwrap();
290        assert_eq!(buf, [36, 37, 38]);
291        loader.load(&mut buf).unwrap();
292        assert_eq!(buf, [39, 40, 41]);
293        loader.load(&mut buf).unwrap();
294        assert_eq!(buf, [42, 43, 44]);
295        loader.load(&mut buf).unwrap();
296        assert_eq!(buf, [45, 46, 47]);
297        loader.load(&mut buf).unwrap();
298        assert_eq!(buf, [48, 49, 50]);
299    }
300
301    #[test]
302    fn test_boxed_slice_buffered_loader() {
303        let mut loader = BufferedLoader::new_heap(Increment::default(), 8);
304        let mut buf = [0u8; 4];
305        loader.load(&mut buf).unwrap();
306        assert_eq!(buf, [0, 1, 2, 3]);
307        loader.load(&mut buf).unwrap();
308        assert_eq!(buf, [4, 5, 6, 7]);
309        loader.load(&mut buf).unwrap();
310        assert_eq!(buf, [8, 9, 10, 11]);
311
312        let mut buf = [0u8; 12];
313        loader.load(&mut buf).unwrap();
314        assert_eq!(buf, [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);
315
316        let mut buf = [0u8; 3];
317
318        loader.load(&mut buf).unwrap();
319        assert_eq!(buf, [24, 25, 26]);
320        loader.load(&mut buf).unwrap();
321        assert_eq!(buf, [27, 28, 29]);
322        loader.load(&mut buf).unwrap();
323        assert_eq!(buf, [30, 31, 32]);
324        loader.load(&mut buf).unwrap();
325        assert_eq!(buf, [33, 34, 35]);
326        loader.load(&mut buf).unwrap();
327        assert_eq!(buf, [36, 37, 38]);
328        loader.load(&mut buf).unwrap();
329        assert_eq!(buf, [39, 40, 41]);
330        loader.load(&mut buf).unwrap();
331        assert_eq!(buf, [42, 43, 44]);
332        loader.load(&mut buf).unwrap();
333        assert_eq!(buf, [45, 46, 47]);
334        loader.load(&mut buf).unwrap();
335        assert_eq!(buf, [48, 49, 50]);
336    }
337}
load_buffer/lib.rs

load_buffer/
lib.rs
