1use std::ops::Deref;
4
5use crate::error::{Error, Result};
6
7#[derive(Copy, Clone, Eq, PartialEq, Debug)]
9pub enum Reference<'b, 'c, T>
10where
11 T: ?Sized + 'static,
12{
13 Borrowed(&'b T),
15 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
33pub trait Read<'r> {
39 fn peek_one(&mut self) -> Result<u8>;
41
42 fn skip_one(&mut self) -> Result<()> {
44 match self.read_one() {
45 Ok(_) => Ok(()),
46 Err(err) => Err(err),
47 }
48 }
49
50 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 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 fn read<'s>(
90 &'s mut self,
91 len: usize,
92 scratch: &'s mut Vec<u8>,
93 ) -> Result<Reference<'r, 's, [u8]>>;
94
95 fn read_into(&mut self, buf: &mut [u8]) -> Result<()>;
97}
98
99pub struct StdIoReader<R> {
103 reader: R,
104 peeked: Option<u8>,
105}
106
107impl<R> StdIoReader<R> {
108 pub fn new(reader: R) -> Self {
110 Self {
111 reader,
112 peeked: None,
113 }
114 }
115
116 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 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
207pub struct SliceReader<'r> {
211 slice: &'r [u8],
212 pos: usize,
213}
214
215impl<'r> SliceReader<'r> {
216 pub fn new(slice: &'r [u8]) -> Self {
218 Self { slice, pos: 0 }
219 }
220
221 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
267pub trait Write {
273 fn write(&mut self, buf: &[u8]) -> Result<usize>;
275
276 fn flush(&mut self) -> Result<()>;
279}
280
281pub struct MutSliceWriter<'w> {
285 slice: &'w mut [u8],
286 pos: usize,
287}
288
289impl<'w> MutSliceWriter<'w> {
290 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
317pub struct VecWriter<'w> {
321 vec: &'w mut Vec<u8>,
322}
323
324impl<'w> VecWriter<'w> {
325 pub fn new(vec: &'w mut Vec<u8>) -> Self {
327 Self { vec }
328 }
329
330 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
347pub struct StdIoWriter<W> {
351 writer: W,
352}
353
354impl<W> StdIoWriter<W> {
355 pub fn new(writer: W) -> Self {
357 Self { writer }
358 }
359
360 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}