protobuf/coded_input_stream/
buf_read_or_reader.rs

1//! `BufRead` pointer or `BufReader` owned.
2
3use std::cmp;
4use std::fmt;
5use std::io;
6use std::io::BufRead;
7use std::io::BufReader;
8use std::io::Read;
9use std::mem::MaybeUninit;
10
11use crate::misc::maybe_uninit_write_slice;
12
13/// Helper type to simplify `BufReadIter` implementation.
14pub(crate) enum BufReadOrReader<'a> {
15    BufReader(BufReader<&'a mut dyn Read>),
16    BufRead(&'a mut dyn BufRead),
17}
18
19impl<'a> fmt::Debug for BufReadOrReader<'a> {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        match self {
22            BufReadOrReader::BufReader(..) => write!(f, "BufReader(...)"),
23            BufReadOrReader::BufRead(..) => write!(f, "BufRead(...)"),
24        }
25    }
26}
27
28impl<'a> Read for BufReadOrReader<'a> {
29    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
30        match self {
31            BufReadOrReader::BufReader(r) => r.read(buf),
32            BufReadOrReader::BufRead(r) => r.read(buf),
33        }
34    }
35
36    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize, io::Error> {
37        match self {
38            BufReadOrReader::BufReader(r) => r.read_to_end(buf),
39            BufReadOrReader::BufRead(r) => r.read_to_end(buf),
40        }
41    }
42
43    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), io::Error> {
44        match self {
45            BufReadOrReader::BufReader(r) => r.read_exact(buf),
46            BufReadOrReader::BufRead(r) => r.read_exact(buf),
47        }
48    }
49}
50
51impl<'a> BufReadOrReader<'a> {
52    /// Similar to `read_exact` but reads into `MaybeUninit`.
53    pub(crate) fn read_exact_uninit(
54        &mut self,
55        buf: &mut [MaybeUninit<u8>],
56    ) -> Result<(), io::Error> {
57        let mut pos = 0;
58        while pos != buf.len() {
59            let fill_buf = match self {
60                BufReadOrReader::BufReader(r) => r.fill_buf()?,
61                BufReadOrReader::BufRead(r) => r.fill_buf()?,
62            };
63            if fill_buf.is_empty() {
64                return Err(io::Error::new(
65                    io::ErrorKind::UnexpectedEof,
66                    "Unexpected end of file",
67                ));
68            }
69            let consume = cmp::min(fill_buf.len(), buf.len() - pos);
70            maybe_uninit_write_slice(&mut buf[pos..pos + consume], &fill_buf[..consume]);
71            match self {
72                BufReadOrReader::BufReader(r) => r.consume(consume),
73                BufReadOrReader::BufRead(r) => r.consume(consume),
74            }
75            pos += consume;
76        }
77        Ok(())
78    }
79
80    pub(crate) fn skip_bytes(&mut self, count: usize) -> Result<(), io::Error> {
81        let mut rem = count;
82        while rem != 0 {
83            let buf = self.fill_buf()?;
84            if buf.is_empty() {
85                return Err(io::Error::new(
86                    io::ErrorKind::UnexpectedEof,
87                    "Unexpected end of file",
88                ));
89            }
90            let consume = cmp::min(buf.len(), rem);
91            self.consume(consume);
92            rem -= consume;
93        }
94        Ok(())
95    }
96}
97
98impl<'a> BufRead for BufReadOrReader<'a> {
99    fn fill_buf(&mut self) -> Result<&[u8], io::Error> {
100        match self {
101            BufReadOrReader::BufReader(r) => r.fill_buf(),
102            BufReadOrReader::BufRead(r) => r.fill_buf(),
103        }
104    }
105
106    fn consume(&mut self, amt: usize) {
107        match self {
108            BufReadOrReader::BufReader(r) => r.consume(amt),
109            BufReadOrReader::BufRead(r) => r.consume(amt),
110        }
111    }
112}