axio/utils/
chain.rs

1#[cfg(feature = "alloc")]
2use alloc::vec::Vec;
3use core::io::BorrowedCursor;
4
5use crate::{BufRead, IoBuf, Read, Result};
6
7/// Adapter to chain together two readers.
8///
9/// This struct is generally created by calling [`chain`] on a reader.
10/// Please see the documentation of [`chain`] for more details.
11///
12/// See [`std::io::Chain`] for more details.
13///
14/// [`chain`]: Read::chain
15#[derive(Debug)]
16pub struct Chain<T, U> {
17    first: T,
18    second: U,
19    done_first: bool,
20}
21
22impl<T, U> Chain<T, U> {
23    pub(crate) fn new(first: T, second: U) -> Self {
24        Chain {
25            first,
26            second,
27            done_first: false,
28        }
29    }
30
31    /// Consumes the `Chain`, returning the wrapped readers.
32    pub fn into_inner(self) -> (T, U) {
33        (self.first, self.second)
34    }
35
36    /// Gets references to the underlying readers in this `Chain`.
37    ///
38    /// Care should be taken to avoid modifying the internal I/O state of the
39    /// underlying readers as doing so may corrupt the internal state of this
40    /// `Chain`.
41    pub fn get_ref(&self) -> (&T, &U) {
42        (&self.first, &self.second)
43    }
44
45    /// Gets mutable references to the underlying readers in this `Chain`.
46    ///
47    /// Care should be taken to avoid modifying the internal I/O state of the
48    /// underlying readers as doing so may corrupt the internal state of this
49    /// `Chain`.
50    pub fn get_mut(&mut self) -> (&mut T, &mut U) {
51        (&mut self.first, &mut self.second)
52    }
53}
54
55impl<T: Read, U: Read> Read for Chain<T, U> {
56    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
57        if !self.done_first {
58            match self.first.read(buf)? {
59                0 if !buf.is_empty() => self.done_first = true,
60                n => return Ok(n),
61            }
62        }
63        self.second.read(buf)
64    }
65
66    fn read_buf(&mut self, mut buf: BorrowedCursor<'_>) -> Result<()> {
67        if buf.capacity() == 0 {
68            return Ok(());
69        }
70
71        if !self.done_first {
72            let old_len = buf.written();
73            self.first.read_buf(buf.reborrow())?;
74
75            if buf.written() != old_len {
76                return Ok(());
77            } else {
78                self.done_first = true;
79            }
80        }
81        self.second.read_buf(buf)
82    }
83
84    #[cfg(feature = "alloc")]
85    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
86        let mut read = 0;
87        if !self.done_first {
88            read += self.first.read_to_end(buf)?;
89            self.done_first = true;
90        }
91        read += self.second.read_to_end(buf)?;
92        Ok(read)
93    }
94
95    // We don't override `read_to_string` here because an UTF-8 sequence could
96    // be split between the two parts of the chain
97}
98
99impl<T: BufRead, U: BufRead> BufRead for Chain<T, U> {
100    fn fill_buf(&mut self) -> Result<&[u8]> {
101        if !self.done_first {
102            match self.first.fill_buf()? {
103                [] => self.done_first = true,
104                buf => return Ok(buf),
105            }
106        }
107        self.second.fill_buf()
108    }
109
110    fn consume(&mut self, amt: usize) {
111        if !self.done_first {
112            self.first.consume(amt)
113        } else {
114            self.second.consume(amt)
115        }
116    }
117
118    #[cfg(feature = "alloc")]
119    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> Result<usize> {
120        let mut read = 0;
121        if !self.done_first {
122            let n = self.first.read_until(byte, buf)?;
123            read += n;
124
125            match buf.last() {
126                Some(b) if *b == byte && n != 0 => return Ok(read),
127                _ => self.done_first = true,
128            }
129        }
130        read += self.second.read_until(byte, buf)?;
131        Ok(read)
132    }
133
134    // We don't override `read_line` here because an UTF-8 sequence could be
135    // split between the two parts of the chain
136}
137
138impl<T: IoBuf, U: IoBuf> IoBuf for Chain<T, U> {
139    #[inline]
140    fn remaining(&self) -> usize {
141        self.first.remaining() + self.second.remaining()
142    }
143}
axio/utils/chain.rs

axio/utils/
chain.rs
