1
use std::fmt::Debug;
2
use std::io::Read;
3

            
4
use byteorder::ReadBytesExt;
5

            
6
use crate::Error;
7

            
8
/// A reader that can temporarily buffer bytes read.
9
pub trait Reader<'de>: ReadBytesExt {
10
    /// Reads exactly `length` bytes.
11
    ///
12
    /// If the reader supports borrowing bytes, [`BufferedBytes::Data`] should
13
    /// be returned. Otherwise, the bytes will be read into `scratch`. `scratch`
14
    /// should only be assumed to be valid if [`BufferedBytes::Scratch`] is
15
    /// returned.
16
    fn buffered_read_bytes(
17
        &mut self,
18
        length: usize,
19
        scratch: &mut Vec<u8>,
20
    ) -> Result<BufferedBytes<'de>, Error>;
21
}
22

            
23
/// Bytes that have been read into a buffer.
24
4
#[derive(Debug)]
25
pub enum BufferedBytes<'de> {
26
    /// The bytes that have been read can be borrowed from the source.
27
    Data(&'de [u8]),
28
    /// The bytes that have been read have been stored in the scratch buffer
29
    /// passed to the function reading bytes.
30
    Scratch,
31
}
32

            
33
impl BufferedBytes<'_> {
34
    /// Resolves the bytes to a byte slice.
35
    #[inline]
36
    #[must_use]
37
6
    pub fn as_slice<'a>(&'a self, scratch: &'a [u8]) -> &'a [u8] {
38
6
        match self {
39
4
            BufferedBytes::Data(data) => data,
40
2
            BufferedBytes::Scratch => scratch,
41
        }
42
6
    }
43
}
44

            
45
/// Reads data from a slice.
46
#[allow(clippy::module_name_repetitions)]
47
pub struct SliceReader<'a> {
48
    pub(crate) data: &'a [u8],
49
}
50

            
51
impl<'a> SliceReader<'a> {
52
    /// Returns the remaining bytes to read.
53
    #[must_use]
54
    #[inline]
55
2
    pub const fn len(&self) -> usize {
56
2
        self.data.len()
57
2
    }
58

            
59
    /// Returns `true` if there are no bytes remaining to read.
60
    #[must_use]
61
    #[inline]
62
2
    pub const fn is_empty(&self) -> bool {
63
2
        self.data.is_empty()
64
2
    }
65
}
66

            
67
impl<'a> Debug for SliceReader<'a> {
68
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69
        f.debug_struct("SliceReader")
70
            .field(
71
                "preview",
72
                &format!("{:0x?}", &self.data[..8.min(self.data.len())]),
73
            )
74
            .finish()
75
    }
76
}
77

            
78
impl<'a> From<&'a [u8]> for SliceReader<'a> {
79
    #[inline]
80
169
    fn from(data: &'a [u8]) -> Self {
81
169
        Self { data }
82
169
    }
83
}
84

            
85
impl<'a> From<SliceReader<'a>> for &'a [u8] {
86
    #[inline]
87
1
    fn from(reader: SliceReader<'a>) -> Self {
88
1
        reader.data
89
1
    }
90
}
91

            
92
impl<'de> Reader<'de> for SliceReader<'de> {
93
    #[inline]
94
35126
    fn buffered_read_bytes(
95
35126
        &mut self,
96
35126
        length: usize,
97
35126
        _scratch: &mut Vec<u8>,
98
35126
    ) -> Result<BufferedBytes<'de>, Error> {
99
35126
        if length > self.data.len() {
100
1
            self.data = &self.data[self.data.len()..];
101
1
            Err(Error::Eof)
102
        } else {
103
35125
            let (start, remaining) = self.data.split_at(length);
104
35125
            self.data = remaining;
105
35125
            Ok(BufferedBytes::Data(start))
106
        }
107
35126
    }
108
}
109

            
110
impl<'a> Read for SliceReader<'a> {
111
    #[inline]
112
222858
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
113
222858
        let remaining_length = self.data.len();
114
222858
        let (to_copy, remaining) = self.data.split_at(remaining_length.min(buf.len()));
115
222858
        buf[..to_copy.len()].copy_from_slice(to_copy);
116
222858
        self.data = remaining;
117
222858
        Ok(to_copy.len())
118
222858
    }
119

            
120
    #[inline]
121
222858
    fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
122
888948
        self.read(buf).map(|_| ())
123
222858
    }
124
}
125

            
126
/// A reader over [`ReadBytesExt`].
127
#[allow(clippy::module_name_repetitions)]
128
pub struct IoReader<R: ReadBytesExt> {
129
    pub(crate) reader: R,
130
}
131
impl<R: ReadBytesExt> IoReader<R> {
132
118
    pub(crate) const fn new(reader: R) -> Self {
133
118
        Self { reader }
134
118
    }
135
}
136

            
137
impl<'de, R: ReadBytesExt> Reader<'de> for IoReader<R> {
138
    #[inline]
139
62
    fn buffered_read_bytes(
140
62
        &mut self,
141
62
        length: usize,
142
62
        scratch: &mut Vec<u8>,
143
62
    ) -> Result<BufferedBytes<'de>, Error> {
144
62
        scratch.resize(length, 0);
145
62
        self.reader.read_exact(scratch)?;
146
62
        Ok(BufferedBytes::Scratch)
147
62
    }
148
}
149

            
150
impl<R: ReadBytesExt> Read for IoReader<R> {
151
    #[inline]
152
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
153
        self.reader.read(buf)
154
    }
155

            
156
    #[inline]
157
    fn read_vectored(&mut self, bufs: &mut [std::io::IoSliceMut<'_>]) -> std::io::Result<usize> {
158
        self.reader.read_vectored(bufs)
159
    }
160

            
161
    #[inline]
162
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
163
        self.reader.read_to_end(buf)
164
    }
165

            
166
    #[inline]
167
    fn read_to_string(&mut self, buf: &mut String) -> std::io::Result<usize> {
168
        self.reader.read_to_string(buf)
169
    }
170

            
171
    #[inline]
172
858
    fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
173
858
        self.reader.read_exact(buf)
174
858
    }
175
}
176

            
177
1
#[test]
178
1
fn slice_reader_pub_methods() {
179
1
    let mut reader = SliceReader::from(&b"a"[..]);
180
1
    assert_eq!(reader.len(), 1);
181
1
    assert!(!reader.is_empty());
182
1
    reader.read_exact(&mut [0]).unwrap();
183
1

            
184
1
    assert_eq!(reader.len(), 0);
185
1
    assert!(reader.is_empty());
186
1
    assert_eq!(<&[u8]>::from(reader), b"");
187
1
}