lilliput_core/
io.rs

1//! I/O related stuff.
2
3use std::ops::Deref;
4
5use crate::error::{Error, Result};
6
7/// A reference to a decoded byte sequence's value.
8#[derive(Copy, Clone, Eq, PartialEq, Debug)]
9pub enum Reference<'b, 'c, T>
10where
11    T: ?Sized + 'static,
12{
13    /// A reference to a borrowed value.
14    Borrowed(&'b T),
15    /// A reference to a copied value.
16    Copied(&'c T),
17}
18
19impl<T> Deref for Reference<'_, '_, T>
20where
21    T: ?Sized + 'static,
22{
23    type Target = T;
24
25    fn deref(&self) -> &Self::Target {
26        match *self {
27            Self::Borrowed(b) => b,
28            Self::Copied(c) => c,
29        }
30    }
31}
32
33// MARK: - Read
34
35/// A trait for objects which are byte-oriented sources.
36///
37/// Implementors of the Read trait are called ‘readers’.
38pub trait Read<'r> {
39    /// Returns the next byte without advancing the position.
40    fn peek_one(&mut self) -> Result<u8>;
41
42    /// Skips the next byte, advancing the position.
43    fn skip_one(&mut self) -> Result<()> {
44        match self.read_one() {
45            Ok(_) => Ok(()),
46            Err(err) => Err(err),
47        }
48    }
49
50    /// Skips the next `len` bytes, advancing the position.
51    fn skip(&mut self, len: usize) -> Result<()> {
52        let mut to_read = len;
53
54        let mut bytes: [u8; 8] = [0b0; 8];
55        while to_read >= 8 {
56            self.read_into(&mut bytes)?;
57            to_read -= 8;
58        }
59
60        let mut bytes: [u8; 4] = [0b0; 4];
61        while to_read >= 4 {
62            self.read_into(&mut bytes)?;
63            to_read -= 4;
64        }
65
66        let mut bytes: [u8; 2] = [0b0; 2];
67        while to_read >= 2 {
68            self.read_into(&mut bytes)?;
69            to_read -= 2;
70        }
71
72        let mut bytes: [u8; 1] = [0b0; 1];
73        while to_read >= 1 {
74            self.read_into(&mut bytes)?;
75            to_read -= 1;
76        }
77
78        Ok(())
79    }
80
81    /// Reads the next byte, advancing the position.
82    fn read_one(&mut self) -> Result<u8> {
83        let mut bytes: [u8; 1] = [0b0];
84        self.read_into(&mut bytes)?;
85        Ok(bytes[0])
86    }
87
88    /// Reads the next `len` bytes into `scratch` (if necessary), advancing the position, returning a reference.
89    fn read<'s>(
90        &'s mut self,
91        len: usize,
92        scratch: &'s mut Vec<u8>,
93    ) -> Result<Reference<'r, 's, [u8]>>;
94
95    /// Reads the next `len` bytes into `buf`, advancing the position.
96    fn read_into(&mut self, buf: &mut [u8]) -> Result<()>;
97}
98
99// MARK: - StdIoReader
100
101/// A wrapper around instances of `std::io::Read`.
102pub struct StdIoReader<R> {
103    reader: R,
104    peeked: Option<u8>,
105}
106
107impl<R> StdIoReader<R> {
108    /// Creates an instance from a `reader`.
109    pub fn new(reader: R) -> Self {
110        Self {
111            reader,
112            peeked: None,
113        }
114    }
115
116    /// Returns the internal `reader`, consuming `self`.
117    pub fn into_reader(self) -> R {
118        self.reader
119    }
120}
121
122impl<'r, R> Read<'r> for StdIoReader<R>
123where
124    R: std::io::Read,
125{
126    fn peek_one(&mut self) -> Result<u8> {
127        if let Some(byte) = self.peeked {
128            return Ok(byte);
129        }
130
131        let byte = self.read_one()?;
132        self.peeked = Some(byte);
133
134        Ok(byte)
135    }
136
137    fn read_one(&mut self) -> Result<u8> {
138        if let Some(byte) = self.peeked.take() {
139            return Ok(byte);
140        }
141
142        let mut bytes: [u8; 1] = [0b0];
143        self.read_into(&mut bytes)?;
144
145        Ok(bytes[0])
146    }
147
148    fn read<'s>(
149        &'s mut self,
150        len: usize,
151        scratch: &'s mut Vec<u8>,
152    ) -> Result<Reference<'r, 's, [u8]>> {
153        // Copied from the default buffer length of `std::io::BufReader`:
154        const MAX_CHUNK_LENGTH: usize = 8192;
155
156        let mut total_read = 0;
157
158        if len == 0 {
159            return Ok(Reference::Copied(&[]));
160        }
161
162        if let Some(byte) = self.peeked.take() {
163            scratch.resize(1, byte);
164            total_read += 1;
165        }
166
167        while total_read < len {
168            let remaining = len - total_read;
169            let to_read = remaining.min(MAX_CHUNK_LENGTH);
170
171            let old_len = scratch.len();
172            scratch.resize(old_len + to_read, 0);
173
174            let read = self
175                .reader
176                .read(&mut scratch[old_len..])
177                .map_err(Error::io)?;
178
179            if read < to_read {
180                return Err(Error::end_of_file());
181            }
182
183            total_read += read;
184        }
185
186        Ok(Reference::Copied(scratch))
187    }
188
189    fn read_into(&mut self, buf: &mut [u8]) -> Result<()> {
190        if buf.is_empty() {
191            return Ok(());
192        }
193
194        let offset = if let Some(byte) = self.peeked.take() {
195            buf[0] = byte;
196            1
197        } else {
198            0
199        };
200
201        self.reader
202            .read_exact(&mut buf[offset..])
203            .map_err(Error::io)
204    }
205}
206
207// MARK: - SliceReader
208
209/// A wrapper around instances of `&[u8]`.
210pub struct SliceReader<'r> {
211    slice: &'r [u8],
212    pos: usize,
213}
214
215impl<'r> SliceReader<'r> {
216    /// Creates an instance from a `slice`.
217    pub fn new(slice: &'r [u8]) -> Self {
218        Self { slice, pos: 0 }
219    }
220
221    /// Returns the current position in the slice.
222    pub fn pos(&self) -> usize {
223        self.pos
224    }
225}
226
227impl<'r> Read<'r> for SliceReader<'r> {
228    fn peek_one(&mut self) -> Result<u8> {
229        if self.pos >= self.slice.len() {
230            return Err(Error::end_of_file());
231        }
232
233        Ok(self.slice[self.pos])
234    }
235
236    fn read<'s>(
237        &'s mut self,
238        len: usize,
239        _scratch: &'s mut Vec<u8>,
240    ) -> Result<Reference<'r, 's, [u8]>> {
241        if self.pos + len > self.slice.len() {
242            return Err(Error::end_of_file());
243        }
244
245        let range = self.pos..(self.pos + len);
246        self.pos += len;
247
248        Ok(Reference::Borrowed(&self.slice[range]))
249    }
250
251    fn read_into(&mut self, buf: &mut [u8]) -> Result<()> {
252        let len = buf.len();
253
254        if self.pos + len > self.slice.len() {
255            return Err(Error::end_of_file());
256        }
257
258        let range = self.pos..(self.pos + len);
259        self.pos += len;
260
261        buf.copy_from_slice(&self.slice[range]);
262
263        Ok(())
264    }
265}
266
267// MARK: - Write
268
269/// A trait for objects which are byte-oriented sinks.
270///
271/// Implementors of the Write trait are sometimes called ‘writers’.
272pub trait Write {
273    /// Writes a buffer into this writer, returning how many bytes were written.
274    fn write(&mut self, buf: &[u8]) -> Result<usize>;
275
276    /// Flushes this output stream, ensuring that all intermediately
277    /// buffered contents reach their destination.
278    fn flush(&mut self) -> Result<()>;
279}
280
281// MARK: - MutSliceWriter
282
283/// A wrapper around instances of `&mut [u8]`.
284pub struct MutSliceWriter<'w> {
285    slice: &'w mut [u8],
286    pos: usize,
287}
288
289impl<'w> MutSliceWriter<'w> {
290    /// Creates a writer from a mutable `slice`.
291    pub fn new(slice: &'w mut Vec<u8>) -> Self {
292        Self { slice, pos: 0 }
293    }
294}
295
296impl Write for MutSliceWriter<'_> {
297    fn write(&mut self, buf: &[u8]) -> Result<usize> {
298        let len = buf.len();
299
300        if self.pos + len > self.slice.len() {
301            return Err(Error::end_of_file());
302        }
303
304        let range = self.pos..(self.pos + len);
305        self.slice[range].copy_from_slice(buf);
306
307        self.pos += len;
308
309        Ok(len)
310    }
311
312    fn flush(&mut self) -> Result<()> {
313        Ok(())
314    }
315}
316
317// MARK: - VecWriter
318
319/// A wrapper around instances of `Vec<u8>`.
320pub struct VecWriter<'w> {
321    vec: &'w mut Vec<u8>,
322}
323
324impl<'w> VecWriter<'w> {
325    /// Creates a writer from a `vec`.
326    pub fn new(vec: &'w mut Vec<u8>) -> Self {
327        Self { vec }
328    }
329
330    /// Returns a slice into the inner `vec`.
331    pub fn vec(&self) -> &[u8] {
332        self.vec
333    }
334}
335
336impl Write for VecWriter<'_> {
337    fn write(&mut self, buf: &[u8]) -> Result<usize> {
338        self.vec.extend_from_slice(buf);
339        Ok(buf.len())
340    }
341
342    fn flush(&mut self) -> Result<()> {
343        Ok(())
344    }
345}
346
347// MARK: - StdIoBufWriter
348
349/// A wrapper around instances of `std::io::Write`.
350pub struct StdIoWriter<W> {
351    writer: W,
352}
353
354impl<W> StdIoWriter<W> {
355    /// Creates an instance from a `reader`.
356    pub fn new(writer: W) -> Self {
357        Self { writer }
358    }
359
360    /// Returns the internal `writer`, consuming `self`.
361    pub fn into_writer(self) -> W {
362        self.writer
363    }
364}
365
366impl<W> Write for StdIoWriter<W>
367where
368    W: std::io::Write,
369{
370    fn write(&mut self, buf: &[u8]) -> Result<usize> {
371        self.writer.write(buf).map_err(Error::io)
372    }
373
374    fn flush(&mut self) -> Result<()> {
375        self.writer.flush().map_err(Error::io)
376    }
377}
378
379#[cfg(test)]
380mod test {
381    use crate::error::ErrorCode;
382
383    use super::*;
384
385    mod std_io_reader {
386        use super::*;
387
388        #[test]
389        fn peek_one() {
390            let slice: &[u8] = &[1, 2, 3, 4, 5];
391            let mut reader = StdIoReader::new(slice);
392            let mut scratch = Vec::new();
393
394            assert_eq!(reader.peek_one().unwrap(), 1);
395
396            let byte = reader.read_one().unwrap();
397            assert_eq!(byte, 1);
398
399            assert_eq!(reader.peek_one().unwrap(), 2);
400
401            match reader.read(2, &mut scratch).unwrap() {
402                Reference::Borrowed(_) => {
403                    panic!("reader should always copy");
404                }
405                Reference::Copied(bytes) => {
406                    assert_eq!(bytes, &[2, 3]);
407                }
408            }
409
410            scratch.clear();
411
412            assert_eq!(reader.peek_one().unwrap(), 4);
413
414            scratch.resize(2, 0b0);
415            reader.read_into(&mut scratch[0..2]).unwrap();
416            assert_eq!(scratch, &[4, 5]);
417
418            scratch.clear();
419
420            assert_eq!(
421                reader.read(3, &mut scratch).err().unwrap().code(),
422                ErrorCode::UnexpectedEndOfFile
423            );
424
425            scratch.clear();
426
427            assert_eq!(
428                reader.peek_one().unwrap_err().code(),
429                ErrorCode::UnexpectedEndOfFile
430            );
431        }
432
433        #[test]
434        fn read_one() {
435            let slice: &[u8] = &[1, 2, 3, 4, 5];
436            let mut reader = StdIoReader::new(slice);
437            let mut scratch = Vec::new();
438
439            match reader.read(1, &mut scratch).unwrap() {
440                Reference::Borrowed(_) => {
441                    panic!("reader should always copy");
442                }
443                Reference::Copied(bytes) => {
444                    assert_eq!(bytes, &[1]);
445                }
446            }
447
448            scratch.clear();
449
450            match reader.read(2, &mut scratch).unwrap() {
451                Reference::Borrowed(_) => {
452                    panic!("reader should always copy");
453                }
454                Reference::Copied(bytes) => {
455                    assert_eq!(bytes, &[2, 3]);
456                }
457            }
458
459            scratch.clear();
460
461            assert_eq!(
462                reader.read(3, &mut scratch).unwrap_err().code(),
463                ErrorCode::UnexpectedEndOfFile
464            );
465        }
466
467        #[test]
468        fn read() {
469            let slice: &[u8] = &[1, 2, 3, 4, 5];
470            let mut reader = StdIoReader::new(slice);
471            let mut scratch = Vec::new();
472
473            match reader.read(1, &mut scratch).unwrap() {
474                Reference::Borrowed(_) => {
475                    panic!("reader should always copy");
476                }
477                Reference::Copied(bytes) => {
478                    assert_eq!(bytes, &[1]);
479                }
480            }
481
482            scratch.clear();
483
484            match reader.read(2, &mut scratch).unwrap() {
485                Reference::Borrowed(_) => {
486                    panic!("reader should always copy");
487                }
488                Reference::Copied(bytes) => {
489                    assert_eq!(bytes, &[2, 3]);
490                }
491            }
492
493            scratch.clear();
494
495            assert_eq!(
496                reader.read(3, &mut scratch).err().unwrap().code(),
497                ErrorCode::UnexpectedEndOfFile
498            );
499        }
500
501        #[test]
502        fn read_into() {
503            let slice: &[u8] = &[1, 2, 3, 4, 5];
504            let mut reader = StdIoReader::new(slice);
505            let mut scratch = Vec::new();
506
507            let bytes = &mut [0];
508            reader.read_into(bytes).unwrap();
509            assert_eq!(bytes, &[1]);
510
511            scratch.clear();
512
513            let bytes = &mut [0, 0];
514            reader.read_into(bytes).unwrap();
515            assert_eq!(bytes, &[2, 3]);
516
517            scratch.clear();
518
519            assert_eq!(
520                reader.read(3, &mut scratch).err().unwrap().code(),
521                ErrorCode::UnexpectedEndOfFile
522            );
523        }
524    }
525
526    mod slice_reader {
527        use super::*;
528
529        #[test]
530        fn peek_one() {
531            let slice: &[u8] = &[1, 2, 3, 4, 5];
532            let mut reader = SliceReader::new(slice);
533            let mut scratch = Vec::new();
534
535            assert_eq!(reader.peek_one().unwrap(), 1);
536
537            let byte = reader.read_one().unwrap();
538            assert_eq!(byte, 1);
539
540            assert_eq!(reader.peek_one().unwrap(), 2);
541
542            match reader.read(2, &mut scratch).unwrap() {
543                Reference::Borrowed(bytes) => {
544                    assert_eq!(bytes, &[2, 3]);
545                }
546                Reference::Copied(_) => {
547                    panic!("reader should always borrow");
548                }
549            }
550
551            scratch.clear();
552
553            assert_eq!(reader.peek_one().unwrap(), 4);
554
555            scratch.resize(2, 0b0);
556            reader.read_into(&mut scratch).unwrap();
557            assert_eq!(scratch, &[4, 5]);
558
559            scratch.clear();
560
561            assert_eq!(
562                reader.read(3, &mut scratch).err().unwrap().code(),
563                ErrorCode::UnexpectedEndOfFile
564            );
565
566            scratch.clear();
567
568            assert_eq!(
569                reader.peek_one().unwrap_err().code(),
570                ErrorCode::UnexpectedEndOfFile
571            );
572        }
573
574        #[test]
575        fn read_one() {
576            let slice: &[u8] = &[1, 2, 3, 4, 5];
577            let mut reader = SliceReader::new(slice);
578
579            assert_eq!(reader.read_one().unwrap(), 1);
580            assert_eq!(reader.read_one().unwrap(), 2);
581            assert_eq!(reader.read_one().unwrap(), 3);
582            assert_eq!(reader.read_one().unwrap(), 4);
583            assert_eq!(reader.read_one().unwrap(), 5);
584
585            assert_eq!(
586                reader.read_one().unwrap_err().code(),
587                ErrorCode::UnexpectedEndOfFile
588            );
589        }
590
591        #[test]
592        fn read() {
593            let slice: &[u8] = &[1, 2, 3, 4, 5];
594            let mut reader = SliceReader::new(slice);
595            let mut scratch = Vec::new();
596
597            match reader.read(1, &mut scratch).unwrap() {
598                Reference::Borrowed(bytes) => {
599                    assert_eq!(bytes, &[1]);
600                }
601                Reference::Copied(_) => {
602                    panic!("reader should always borrow");
603                }
604            }
605
606            scratch.clear();
607
608            match reader.read(2, &mut scratch).unwrap() {
609                Reference::Borrowed(bytes) => {
610                    assert_eq!(bytes, &[2, 3]);
611                }
612                Reference::Copied(_) => {
613                    panic!("reader should always borrow");
614                }
615            }
616
617            scratch.clear();
618
619            assert_eq!(
620                reader.read(3, &mut scratch).err().unwrap().code(),
621                ErrorCode::UnexpectedEndOfFile
622            );
623        }
624
625        #[test]
626        fn read_into() {
627            let slice: &[u8] = &[1, 2, 3, 4, 5];
628            let mut reader = SliceReader::new(slice);
629            let mut scratch = Vec::new();
630
631            let bytes = &mut [0];
632            reader.read_into(bytes).unwrap();
633            assert_eq!(bytes, &[1]);
634
635            scratch.clear();
636
637            let bytes = &mut [0, 0];
638            reader.read_into(bytes).unwrap();
639            assert_eq!(bytes, &[2, 3]);
640
641            scratch.clear();
642
643            assert_eq!(
644                reader.read(3, &mut scratch).err().unwrap().code(),
645                ErrorCode::UnexpectedEndOfFile
646            );
647        }
648    }
649}