1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::vec::Vec;

use core::fmt::Debug;

/// Like [`std::io::Read`], but only supporting [`read_exact`], and with an
/// arbitrary error type.
///
/// [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
/// [`read_exact`]: https://doc.rust-lang.org/std/io/trait.Read.html#method.read_exact
pub trait Load {
    type Error;

    fn load(&mut self, buf: &mut [u8]) -> Result<(), Self::Error>;
}

#[cfg(feature = "std")]
impl<T> Load for T
where
    T: std::io::Read,
{
    type Error = std::io::Error;
    fn load(&mut self, buf: &mut [u8]) -> Result<(), Self::Error> {
        self.read_exact(buf)
    }
}

/// A buffering wrapper around a [`Load`].
///
/// This fills some of the simpler niches of a [`std::io::BufReader`] for no_std
/// use-cases, and allowing a static buffer.
///
/// [`std::io::BufReader`]: https://doc.rust-lang.org/std/io/struct.BufReader.html
/// [`Load`]: trait.Load.html
#[derive(Clone, Debug)]
pub struct BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    loader: Loader,
    buffer: Buffer,

    /// The amount of buffer that has been read already.
    offset: usize,
}

impl<const N: usize, Loader> BufferedLoader<[u8; N], Loader>
where
    Loader: Load,
{
    pub fn new_static(loader: Loader) -> Self {
        Self {
            loader,
            buffer: [0u8; N],
            offset: N,
        }
    }
}

#[cfg(feature = "alloc")]
impl<Loader> BufferedLoader<Vec<u8>, Loader>
where
    Loader: Load,
{
    pub fn new_dynamic(loader: Loader, buffer_size: usize) -> Self {
        let mut buffer = Vec::with_capacity(buffer_size);
        buffer.resize(buffer_size, 0);
        Self {
            loader,
            buffer,
            offset: buffer_size,
        }
    }

    pub fn resize(&mut self, new_size: usize) {
        let invalid = new_size > self.buffer.len();

        self.buffer.resize(new_size, 0);
        if invalid {
            self.offset = new_size;
        }
    }
}

impl<Buffer, Loader> BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    /// Drop what's remaining in the buffer and give a mutable reference to the
    /// inner loader, so it can be seeked or otherwise manipulated.
    pub fn loader(&mut self) -> &mut Loader {
        self.offset = self.buffer.as_ref().len();
        &mut self.loader
    }
}

impl<Buffer, Loader> Load for BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    type Error = Loader::Error;

    /// Efficiently fill the destination buffer, calling the underlying
    /// [`Load::load`] as few times as possible.
    ///
    /// [`Load::Load`]: trait.Load.html#method.load
    fn load(&mut self, mut buf: &mut [u8]) -> Result<(), Self::Error> {
        if buf.is_empty() {
            return Ok(());
        }

        let buffer = self.buffer.as_mut();
        let buffer_size = buffer.len();

        let buffer_remaining = buffer_size - self.offset;

        if buf.len() <= buffer_remaining {
            // There are enough bytes left in the buffer to consume without
            // reading.
            let end = self.offset + buf.len();
            buf.copy_from_slice(&buffer[self.offset..end]);
            self.offset = end;
        } else {
            // First empty the buffer.
            if buffer_remaining > 0 {
                buf[..buffer_remaining].copy_from_slice(&buffer[self.offset..buffer_size]);
                let copied = buffer_size - self.offset;
                buf = &mut buf[copied..];
            }

            let buffers = buf.len() / buffer_size;
            let remainder = buf.len() % buffer_size;

            // Load full-sized chunks directly to the destination, bypassing
            // the buffer.
            if buffers > 0 {
                let buffers_bytes = buffers * buffer_size;
                self.loader.load(&mut buf[..buffers_bytes])?;
                buf = &mut buf[buffers_bytes..];
            }

            // Fill the buffer for the remainder, if there is any.
            if remainder > 0 {
                self.loader.load(buffer)?;
                buf.copy_from_slice(&buffer[..remainder]);
                self.offset = remainder;
            } else {
                // We have emptied the remaining buffer, so mark this empty.
                self.offset = buffer_size;
            }
        }
        Ok(())
    }
}