scuffle_bytes_util/zero_copy/
reader.rs

1//! Zero-copy reader types.
2
3use std::io::{self, Read};
4
5use crate::BytesCow;
6
7/// A trait for zero-copy readers.
8pub trait ZeroCopyReader<'a> {
9    /// Attempts to read a specified number of bytes from the reader without copying.
10    ///
11    /// This function does not guarantee that no copying will occur.
12    /// Some implementations can't avoid copying.
13    fn try_read(&mut self, size: usize) -> Result<BytesCow<'a>, io::Error>;
14
15    /// Returns a standard [`io::Read`] interface for the reader.
16    fn as_std(&mut self) -> impl io::Read;
17
18    /// Limits the number of bytes that can be read from the reader.
19    fn take(self, limit: usize) -> Take<Self>
20    where
21        Self: Sized,
22    {
23        Take::new(self, limit)
24    }
25
26    /// Reads all remaining bytes from the reader.
27    fn try_read_to_end(&mut self) -> Result<BytesCow<'a>, io::Error> {
28        let mut buf = Vec::new();
29        self.as_std().read_to_end(&mut buf)?;
30        Ok(BytesCow::from_vec(buf))
31    }
32}
33
34impl<'a, T: ZeroCopyReader<'a>> ZeroCopyReader<'a> for &mut T {
35    fn try_read(&mut self, size: usize) -> Result<BytesCow<'a>, io::Error> {
36        (*self).try_read(size)
37    }
38
39    fn as_std(&mut self) -> impl io::Read {
40        (*self).as_std()
41    }
42}
43
44/// A zero-copy reader that wraps a [`bytes::Buf`].
45pub struct BytesBuf<B>(B);
46
47impl<B: bytes::Buf> From<B> for BytesBuf<B> {
48    fn from(buf: B) -> Self {
49        Self(buf)
50    }
51}
52
53impl<'a, B: bytes::Buf> ZeroCopyReader<'a> for BytesBuf<B> {
54    fn try_read(&mut self, size: usize) -> Result<BytesCow<'a>, io::Error> {
55        if self.0.remaining() < size {
56            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Not enough data"));
57        }
58
59        Ok(BytesCow::from_bytes(self.0.copy_to_bytes(size)))
60    }
61
62    fn as_std(&mut self) -> impl io::Read {
63        bytes::Buf::reader(&mut self.0)
64    }
65}
66
67/// A zero-copy reader that wraps a [`std::io::Read`].
68///
69/// This implementation is not zero-copy and will always copy the data into a new buffer.
70/// It is not possible to implement zero-copy reading for [`std::io::Read`]
71/// because it does not provide a way to access the underlying buffer directly.
72pub struct IoRead<R>(R);
73
74impl<R: io::Read> From<R> for IoRead<R> {
75    fn from(reader: R) -> Self {
76        Self(reader)
77    }
78}
79
80impl<'a, R: io::Read> ZeroCopyReader<'a> for IoRead<R> {
81    fn try_read(&mut self, size: usize) -> Result<BytesCow<'a>, io::Error> {
82        let mut buf = vec![0; size];
83        self.0.read_exact(&mut buf)?;
84        Ok(BytesCow::from_vec(buf))
85    }
86
87    fn as_std(&mut self) -> impl io::Read {
88        &mut self.0
89    }
90}
91
92/// A zero-copy reader that wraps a byte slice (`&[u8]`).
93pub struct Slice<'a>(io::Cursor<&'a [u8]>);
94
95impl<'a> From<&'a [u8]> for Slice<'a> {
96    fn from(slice: &'a [u8]) -> Self {
97        Self(io::Cursor::new(slice))
98    }
99}
100
101impl<'a> ZeroCopyReader<'a> for Slice<'a> {
102    fn try_read(&mut self, size: usize) -> Result<BytesCow<'a>, io::Error> {
103        let start = self.0.position() as usize;
104        let end = start + size;
105
106        if end > self.0.get_ref().len() {
107            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Not enough data"));
108        }
109
110        let slice = &self.0.get_ref()[start..end];
111        self.0.set_position(end as u64);
112        Ok(BytesCow::from_slice(slice))
113    }
114
115    fn as_std(&mut self) -> impl io::Read {
116        &mut self.0
117    }
118}
119
120/// Zero-copy reader that limits the number of bytes read.
121///
122/// This is equivalent to [`std::io::Take`], but it implements [`ZeroCopyReader`].
123pub struct Take<R> {
124    inner: R,
125    limit: usize,
126}
127
128struct TakeAsStd<'a, R> {
129    inner: R,
130    limit: &'a mut usize,
131}
132
133impl<R: io::Read> io::Read for TakeAsStd<'_, R> {
134    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
135        if *self.limit == 0 {
136            return Ok(0);
137        }
138
139        let size = std::cmp::min(buf.len(), *self.limit);
140        let n = self.inner.read(&mut buf[..size])?;
141        *self.limit -= n;
142        Ok(n)
143    }
144}
145
146impl<R> Take<R> {
147    /// Creates a new [`Take`] reader that limits the number of bytes read.
148    pub fn new(inner: R, limit: usize) -> Self {
149        Self { inner, limit }
150    }
151
152    /// Returns the underlying reader.
153    pub fn into_inner(self) -> R {
154        self.inner
155    }
156}
157
158impl<'a, R> ZeroCopyReader<'a> for Take<R>
159where
160    R: ZeroCopyReader<'a>,
161{
162    fn as_std(&mut self) -> impl io::Read {
163        TakeAsStd {
164            inner: self.inner.as_std(),
165            limit: &mut self.limit,
166        }
167    }
168
169    fn try_read(&mut self, size: usize) -> Result<BytesCow<'a>, io::Error> {
170        let size = std::cmp::min(size, self.limit);
171        let result = self.inner.try_read(size)?;
172        self.limit -= result.len();
173        Ok(result)
174    }
175}
176
177#[cfg(test)]
178#[cfg_attr(all(test, coverage_nightly), coverage(off))]
179mod tests {
180    use super::{Slice, ZeroCopyReader};
181
182    #[test]
183    fn take_and_read_to_end() {
184        let data = b"Hello, world!";
185        let mut reader = Slice::from(&data[..]).take(5);
186        assert_eq!(reader.try_read_to_end().unwrap(), b"Hello");
187    }
188}
scuffle_bytes_util/zero_copy/reader.rs

scuffle_bytes_util/zero_copy/
reader.rs
