hxd/
reader.rs

1use std::{cmp::min, convert::Infallible, fmt::Debug, io::Read};
2
3use crate::Endianness;
4
5#[doc(hidden)]
6pub trait GroupedReader<const N: usize> {
7    fn read_next(&mut self, end: Endianness) -> Option<[u8; N]>;
8    fn size(&self) -> usize {
9        N
10    }
11}
12
13pub struct ByteSliceReader<'a> {
14    slice: &'a [u8],
15    index: usize,
16}
17
18impl<'a> ByteSliceReader<'a> {
19    pub fn new(slice: &'a [u8]) -> ByteSliceReader<'a> {
20        Self {
21            slice,
22            index: 0usize,
23        }
24    }
25}
26
27impl<'a> ReadBytes for ByteSliceReader<'a> {
28    type Error = Infallible;
29
30    fn next_n<'buf>(&mut self, buf: &'buf mut [u8]) -> Result<&'buf [u8], Self::Error> {
31        if self.index >= self.slice.len() {
32            return Ok(&[]);
33        }
34        let end = min(self.index + buf.len(), self.slice.len()) - self.index;
35        buf[..end].copy_from_slice(&self.slice[self.index..self.index + end]);
36        self.index += end;
37        Ok(&buf[..end])
38    }
39
40    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
41        self.index += n;
42        Ok(self.index)
43    }
44
45    fn total_byte_hint(&self) -> Option<usize> {
46        Some(self.slice.len())
47    }
48}
49
50/// This trait provides a method to convert
51/// integer types into sized byte arrays.
52/// Under the hood, implementations for primitive integer
53/// types call `to_be_bytes()` or `to_le_bytes()`
54/// depending on the [endianness](crate::options::Endianness).
55pub trait EndianBytes<const N: usize> {
56    fn to_bytes(&self, end: Endianness) -> [u8; N];
57}
58
59impl EndianBytes<1> for u8 {
60    fn to_bytes(&self, _: Endianness) -> [u8; 1] {
61        [*self]
62    }
63}
64
65impl EndianBytes<1> for i8 {
66    fn to_bytes(&self, _: Endianness) -> [u8; 1] {
67        [*self as u8]
68    }
69}
70
71impl EndianBytes<2> for u16 {
72    fn to_bytes(&self, endianness: Endianness) -> [u8; 2] {
73        match endianness {
74            Endianness::BigEndian => self.to_be_bytes(),
75            Endianness::LittleEndian => self.to_le_bytes(),
76        }
77    }
78}
79
80impl EndianBytes<2> for i16 {
81    fn to_bytes(&self, endianness: Endianness) -> [u8; 2] {
82        match endianness {
83            Endianness::BigEndian => self.to_be_bytes(),
84            Endianness::LittleEndian => self.to_le_bytes(),
85        }
86    }
87}
88
89impl EndianBytes<4> for u32 {
90    fn to_bytes(&self, endianness: Endianness) -> [u8; 4] {
91        match endianness {
92            Endianness::BigEndian => self.to_be_bytes(),
93            Endianness::LittleEndian => self.to_le_bytes(),
94        }
95    }
96}
97
98impl EndianBytes<4> for i32 {
99    fn to_bytes(&self, endianness: Endianness) -> [u8; 4] {
100        match endianness {
101            Endianness::BigEndian => self.to_be_bytes(),
102            Endianness::LittleEndian => self.to_le_bytes(),
103        }
104    }
105}
106
107impl EndianBytes<8> for u64 {
108    fn to_bytes(&self, endianness: Endianness) -> [u8; 8] {
109        match endianness {
110            Endianness::BigEndian => self.to_be_bytes(),
111            Endianness::LittleEndian => self.to_le_bytes(),
112        }
113    }
114}
115
116impl EndianBytes<8> for i64 {
117    fn to_bytes(&self, endianness: Endianness) -> [u8; 8] {
118        match endianness {
119            Endianness::BigEndian => self.to_be_bytes(),
120            Endianness::LittleEndian => self.to_le_bytes(),
121        }
122    }
123}
124
125impl EndianBytes<16> for u128 {
126    fn to_bytes(&self, endianness: Endianness) -> [u8; 16] {
127        match endianness {
128            Endianness::BigEndian => self.to_be_bytes(),
129            Endianness::LittleEndian => self.to_le_bytes(),
130        }
131    }
132}
133
134impl EndianBytes<16> for i128 {
135    fn to_bytes(&self, endianness: Endianness) -> [u8; 16] {
136        match endianness {
137            Endianness::BigEndian => self.to_be_bytes(),
138            Endianness::LittleEndian => self.to_le_bytes(),
139        }
140    }
141}
142pub struct GroupedSliceReader<'a, U: EndianBytes<N>, const N: usize> {
143    slice: &'a [U],
144    index: usize,
145}
146
147pub struct GroupedSliceByteReader<'a, U: EndianBytes<N>, const N: usize> {
148    slice: &'a [U],
149    elt_index: usize,
150    u_index: usize,
151    current_elt: Option<[u8; N]>,
152    endianness: Endianness,
153}
154
155impl<'a, U: EndianBytes<N>, const N: usize> ReadBytes for GroupedSliceByteReader<'a, U, N> {
156    type Error = Infallible;
157
158    fn next_n<'buf>(&mut self, buf: &'buf mut [u8]) -> Result<&'buf [u8], Self::Error> {
159        Ok(self.next_bytes(buf))
160    }
161
162    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
163        self.advance_indices_by(n);
164        Ok(n)
165    }
166
167    fn total_byte_hint(&self) -> Option<usize> {
168        Some(self.slice.len() * N)
169    }
170}
171
172impl<'a, U: EndianBytes<N>, const N: usize> GroupedSliceByteReader<'a, U, N> {
173    pub fn new(slice: &'a [U], endianness: Endianness) -> Self {
174        let current_elt = if slice.len() > 0 {
175            Some(slice[0].to_bytes(endianness))
176        } else {
177            None
178        };
179        Self {
180            slice,
181            elt_index: 0,
182            u_index: 0,
183            current_elt,
184            endianness,
185        }
186    }
187    pub fn next_bytes<'buf>(&mut self, o: &'buf mut [u8]) -> &'buf [u8] {
188        for i in 0..o.len() {
189            if let Some(cb) = self.next_byte() {
190                o[i] = cb;
191            } else {
192                return &o[..i];
193            }
194        }
195        &o[..]
196    }
197
198    fn next_byte(&mut self) -> Option<u8> {
199        let o = self.current_elt.map(|ce| ce[self.u_index]);
200        self.advance_indices();
201        o
202    }
203
204    fn advance_indices(&mut self) {
205        self.u_index += 1;
206        if self.u_index >= N {
207            self.u_index = 0;
208            self.elt_index += 1;
209            self.current_elt = if self.elt_index < self.slice.len() {
210                Some(self.slice[self.elt_index].to_bytes(self.endianness))
211            } else {
212                None
213            }
214        }
215    }
216
217    fn advance_indices_by(&mut self, adv: usize) {
218        if N == 1 {
219            self.elt_index = adv;
220            self.u_index = 0;
221            self.current_elt = if self.elt_index < self.slice.len() {
222                Some(self.slice[self.elt_index].to_bytes(self.endianness))
223            } else {
224                None
225            };
226            return;
227        }
228        let mut adv = adv;
229        if self.u_index > 0 {
230            adv -= N - self.u_index;
231            self.u_index = 0;
232            self.elt_index += 1;
233        }
234
235        while adv > N {
236            adv -= N;
237            self.u_index = 0;
238            self.elt_index += 1;
239        }
240
241        self.current_elt = if self.elt_index < self.slice.len() {
242            Some(self.slice[self.elt_index].to_bytes(self.endianness))
243        } else {
244            None
245        };
246
247        if adv > 0 {
248            self.u_index = adv;
249        }
250    }
251}
252
253impl<'a, U: EndianBytes<N>, const N: usize> GroupedSliceReader<'a, U, N> {
254    pub fn new(slice: &'a [U]) -> Self {
255        Self { slice, index: 0 }
256    }
257}
258
259impl<'a, U: EndianBytes<N>, const N: usize> GroupedSliceReader<'a, U, N> {
260    pub fn next(&mut self, end: Endianness) -> Option<[u8; N]> {
261        if self.index < self.slice.len() {
262            let s = Some(self.slice[self.index].to_bytes(end));
263            self.index += 1;
264            s
265        } else {
266            None
267        }
268    }
269}
270
271impl<'a, const N: usize, U: EndianBytes<N>> GroupedReader<N> for GroupedSliceReader<'a, U, N> {
272    fn read_next(&mut self, end: Endianness) -> Option<[u8; N]> {
273        self.next(end)
274    }
275}
276
277pub struct IteratorByteReader<I: Iterator<Item = u8>> {
278    iterator: I,
279}
280
281impl<I: Iterator<Item = u8>> IteratorByteReader<I> {
282    pub fn new(iterator: I) -> Self {
283        Self { iterator }
284    }
285}
286
287impl<I: Iterator<Item = u8>> ReadBytes for IteratorByteReader<I> {
288    type Error = Infallible;
289
290    fn next_n<'buf>(&mut self, buf: &'buf mut [u8]) -> Result<&'buf [u8], Self::Error> {
291        let mut i = 0usize;
292        while i < buf.len() {
293            match self.iterator.next() {
294                Some(b) => {
295                    buf[i] = b;
296                }
297                None => {
298                    return Ok(&buf[..i]);
299                }
300            }
301            i += 1;
302        }
303        Ok(&buf[..i])
304    }
305
306    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
307        for i in 0..n {
308            if let None = self.iterator.next() {
309                return Ok(i);
310            }
311        }
312        Ok(n)
313    }
314}
315
316pub struct GroupedIteratorReader<U: EndianBytes<N>, I: Iterator<Item = U>, const N: usize> {
317    iterator: I,
318    current: Option<[u8; N]>,
319    index: usize,
320    endianness: Endianness,
321}
322
323impl<U: EndianBytes<N>, I: Iterator<Item = U>, const N: usize> GroupedIteratorReader<U, I, N> {
324    pub fn new(mut iterator: I, endianness: Endianness) -> Self {
325        let current = iterator.next().map(|u| u.to_bytes(endianness));
326        Self {
327            iterator,
328            current,
329            index: 0,
330            endianness,
331        }
332    }
333
334    pub fn next_byte(&mut self) -> Option<u8> {
335        let b = self.current.map(|c| c[self.index]);
336        self.index += 1;
337        if self.index >= N {
338            self.index = 0;
339            self.current = self.iterator.next().map(|u| u.to_bytes(self.endianness));
340        }
341        b
342    }
343}
344
345impl<U: EndianBytes<N>, I: Iterator<Item = U>, const N: usize> ReadBytes
346    for GroupedIteratorReader<U, I, N>
347{
348    type Error = Infallible;
349
350    fn next_n<'buf>(&mut self, buf: &'buf mut [u8]) -> Result<&'buf [u8], Self::Error> {
351        let mut i = 0usize;
352        while i < buf.len() {
353            if let Some(b) = self.next_byte() {
354                buf[i] = b;
355                i += 1;
356            } else {
357                return Ok(&buf[..i]);
358            }
359        }
360        Ok(&buf[..i])
361    }
362
363    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
364        for _ in 0..n {
365            if self.next_byte().is_none() {
366                return Ok(n);
367            }
368        }
369        Ok(n)
370    }
371
372    fn total_byte_hint(&self) -> Option<usize> {
373        None
374    }
375}
376
377pub trait ReadBytes {
378    type Error: Debug;
379    fn next_n<'buf>(&mut self, buf: &'buf mut [u8]) -> Result<&'buf [u8], Self::Error>;
380
381    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
382        const SKIP_LEN: usize = 64usize;
383        let mut skipbuf = [0u8; SKIP_LEN];
384        let mut i = 0usize;
385        while i < n {
386            let skipbuf = &mut skipbuf[..min(n - i, SKIP_LEN)];
387            let b = self.next_n(skipbuf)?;
388            if b.len() == 0 {
389                return Ok(i);
390            }
391            i += b.len();
392        }
393        Ok(n)
394    }
395    fn total_byte_hint(&self) -> Option<usize> {
396        None
397    }
398}
399
400impl<'b, T: Iterator<Item = &'b u8>> ReadBytes for T {
401    type Error = Infallible;
402    fn next_n<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a [u8], Self::Error> {
403        let mut i = 0;
404        while i < buf.len() {
405            match self.next() {
406                Some(u) => {
407                    buf[i] = *u;
408                }
409                None => {
410                    break;
411                }
412            };
413            i += 1;
414        }
415        Ok(&buf[..i])
416    }
417
418    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
419        for i in 0..n {
420            match self.next() {
421                Some(_) => {}
422                None => {
423                    return Ok(i);
424                }
425            }
426        }
427        Ok(n)
428    }
429
430    fn total_byte_hint(&self) -> Option<usize> {
431        match self.size_hint() {
432            (_, Some(upper)) => Some(upper),
433            (lower, None) if lower > 0 => Some(lower),
434            _ => None,
435        }
436    }
437}
438
439pub struct IOReader<R: Read>(R);
440
441impl<R: Read> IOReader<R> {
442    pub fn new(reader: R) -> IOReader<R> {
443        Self(reader)
444    }
445}
446
447impl<R: Read> ReadBytes for IOReader<R> {
448    type Error = std::io::Error;
449
450    fn next_n<'buf>(&mut self, buf: &'buf mut [u8]) -> Result<&'buf [u8], Self::Error> {
451        let n = self.0.read(buf)?;
452        Ok(&buf[..n])
453    }
454
455    fn skip_n(&mut self, n: usize) -> Result<usize, Self::Error> {
456        let mut skipped = 0;
457        while skipped < n {
458            let mut buf = [0u8; 1024];
459            let bytes = self.next_n(buf.as_mut_slice())?;
460            if bytes.len() == 0 {
461                break;
462            }
463            skipped += bytes.len();
464        }
465        Ok(skipped)
466    }
467}