nuclei/io/util/
buf_read_ext.rs

1//! [`AsyncBufReadExt`] and [`AsyncSeekExt`]: the reads every buffered reader
2//! gets (up to a byte, a line, line by line) and the seeks every seekable
3//! stream gets, as futures.
4
5use core::future::Future;
6use core::pin::Pin;
7use core::task::{Context, Poll, ready};
8use std::io::{self, SeekFrom};
9
10use futures_io::{AsyncBufRead, AsyncSeek};
11
12/// Reads on any [`AsyncBufRead`]: up to a delimiter, a line, the lines or
13/// the segments one by one.
14pub trait AsyncBufReadExt: AsyncBufRead {
15    /// Reads up to and including `byte` (or to the end), appending to `buf`;
16    /// answers how many bytes it appended (0 at the end).
17    fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut Vec<u8>) -> ReadUntil<'a, Self>
18    where
19        Self: Unpin,
20    {
21        ReadUntil {
22            reader: self,
23            byte,
24            buf,
25            read: 0,
26        }
27    }
28
29    /// Reads up to and including a newline (or to the end), appending to
30    /// `buf` if what it read is UTF-8 (an
31    /// [`InvalidData`](io::ErrorKind::InvalidData) error and `buf` as it was
32    /// otherwise); answers how many bytes it appended (0 at the end).
33    fn read_line<'a>(&'a mut self, buf: &'a mut String) -> ReadLine<'a, Self>
34    where
35        Self: Unpin,
36    {
37        ReadLine {
38            reader: self,
39            buf,
40            bytes: Vec::new(),
41            read: 0,
42        }
43    }
44
45    /// The lines, without their `\n` or `\r\n`, one by one through
46    /// [`Lines::next_line`].
47    fn lines(self) -> Lines<Self>
48    where
49        Self: Sized,
50    {
51        Lines {
52            reader: self,
53            bytes: Vec::new(),
54            read: 0,
55        }
56    }
57
58    /// The segments between `byte`s, without it, one by one through
59    /// [`Split::next_segment`].
60    fn split(self, byte: u8) -> Split<Self>
61    where
62        Self: Sized,
63    {
64        Split {
65            reader: self,
66            byte,
67            buf: Vec::new(),
68            read: 0,
69        }
70    }
71
72    /// The reader's buffer, filled if it was empty (empty at the end).
73    fn fill_buf(&mut self) -> FillBuf<'_, Self>
74    where
75        Self: Unpin,
76    {
77        FillBuf { reader: Some(self) }
78    }
79
80    /// Marks `amount` bytes of the buffer read.
81    fn consume(&mut self, amount: usize)
82    where
83        Self: Unpin,
84    {
85        Pin::new(self).consume(amount);
86    }
87}
88
89impl<R: AsyncBufRead + ?Sized> AsyncBufReadExt for R {}
90
91/// Reads up to and including `byte`, appending to `buf` and counting into
92/// `read`; ready at the byte or the end.
93fn read_until<R: AsyncBufRead + ?Sized>(
94    mut reader: Pin<&mut R>,
95    cx: &mut Context<'_>,
96    byte: u8,
97    buf: &mut Vec<u8>,
98    read: &mut usize,
99) -> Poll<io::Result<usize>> {
100    loop {
101        let (done, used) = {
102            let available = ready!(reader.as_mut().poll_fill_buf(cx))?;
103            match memchr(byte, available) {
104                Some(i) => {
105                    buf.extend_from_slice(&available[..=i]);
106                    (true, i + 1)
107                }
108                None => {
109                    buf.extend_from_slice(available);
110                    (available.is_empty(), available.len())
111                }
112            }
113        };
114        reader.as_mut().consume(used);
115        *read += used;
116        if done {
117            return Poll::Ready(Ok(core::mem::take(read)));
118        }
119    }
120}
121
122fn memchr(byte: u8, haystack: &[u8]) -> Option<usize> {
123    haystack.iter().position(|&b| b == byte)
124}
125
126/// The future of [`read_until`](AsyncBufReadExt::read_until).
127#[derive(Debug)]
128#[must_use = "futures do nothing unless polled"]
129pub struct ReadUntil<'a, R: ?Sized> {
130    reader: &'a mut R,
131    byte: u8,
132    buf: &'a mut Vec<u8>,
133    read: usize,
134}
135
136impl<R: AsyncBufRead + Unpin + ?Sized> Future for ReadUntil<'_, R> {
137    type Output = io::Result<usize>;
138
139    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
140        let this = self.get_mut();
141        read_until(
142            Pin::new(&mut *this.reader),
143            cx,
144            this.byte,
145            this.buf,
146            &mut this.read,
147        )
148    }
149}
150
151/// The future of [`read_line`](AsyncBufReadExt::read_line).
152#[derive(Debug)]
153#[must_use = "futures do nothing unless polled"]
154pub struct ReadLine<'a, R: ?Sized> {
155    reader: &'a mut R,
156    buf: &'a mut String,
157    bytes: Vec<u8>,
158    read: usize,
159}
160
161impl<R: AsyncBufRead + Unpin + ?Sized> Future for ReadLine<'_, R> {
162    type Output = io::Result<usize>;
163
164    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
165        let this = self.get_mut();
166        let n = ready!(read_until(
167            Pin::new(&mut *this.reader),
168            cx,
169            b'\n',
170            &mut this.bytes,
171            &mut this.read
172        ))?;
173        let bytes = core::mem::take(&mut this.bytes);
174        match String::from_utf8(bytes) {
175            Ok(line) => {
176                this.buf.push_str(&line);
177                Poll::Ready(Ok(n))
178            }
179            Err(_) => Poll::Ready(Err(io::Error::new(
180                io::ErrorKind::InvalidData,
181                "the line is not UTF-8",
182            ))),
183        }
184    }
185}
186
187/// The lines of a buffered reader; made by
188/// [`lines`](AsyncBufReadExt::lines).
189#[derive(Debug)]
190pub struct Lines<R> {
191    reader: R,
192    bytes: Vec<u8>,
193    read: usize,
194}
195
196impl<R: AsyncBufRead + Unpin> Lines<R> {
197    /// The next line without its `\n` or `\r\n`, `None` at the end.
198    pub async fn next_line(&mut self) -> io::Result<Option<String>> {
199        core::future::poll_fn(|cx| self.poll_next_line(cx)).await
200    }
201
202    /// Polls for the next line.
203    pub fn poll_next_line(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<Option<String>>> {
204        let n = ready!(read_until(
205            Pin::new(&mut self.reader),
206            cx,
207            b'\n',
208            &mut self.bytes,
209            &mut self.read
210        ))?;
211        if n == 0 && self.bytes.is_empty() {
212            return Poll::Ready(Ok(None));
213        }
214        let mut bytes = core::mem::take(&mut self.bytes);
215        if bytes.last() == Some(&b'\n') {
216            bytes.pop();
217            if bytes.last() == Some(&b'\r') {
218                bytes.pop();
219            }
220        }
221        match String::from_utf8(bytes) {
222            Ok(line) => Poll::Ready(Ok(Some(line))),
223            Err(_) => Poll::Ready(Err(io::Error::new(
224                io::ErrorKind::InvalidData,
225                "the line is not UTF-8",
226            ))),
227        }
228    }
229
230    /// The reader, taken back.
231    pub fn into_inner(self) -> R {
232        self.reader
233    }
234}
235
236/// The segments of a buffered reader between a byte; made by
237/// [`split`](AsyncBufReadExt::split).
238#[derive(Debug)]
239pub struct Split<R> {
240    reader: R,
241    byte: u8,
242    buf: Vec<u8>,
243    read: usize,
244}
245
246impl<R: AsyncBufRead + Unpin> Split<R> {
247    /// The next segment without its delimiter, `None` at the end.
248    pub async fn next_segment(&mut self) -> io::Result<Option<Vec<u8>>> {
249        core::future::poll_fn(|cx| self.poll_next_segment(cx)).await
250    }
251
252    /// Polls for the next segment.
253    pub fn poll_next_segment(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<Option<Vec<u8>>>> {
254        let n = ready!(read_until(
255            Pin::new(&mut self.reader),
256            cx,
257            self.byte,
258            &mut self.buf,
259            &mut self.read
260        ))?;
261        if n == 0 && self.buf.is_empty() {
262            return Poll::Ready(Ok(None));
263        }
264        let mut segment = core::mem::take(&mut self.buf);
265        if segment.last() == Some(&self.byte) {
266            segment.pop();
267        }
268        Poll::Ready(Ok(Some(segment)))
269    }
270
271    /// The reader, taken back.
272    pub fn into_inner(self) -> R {
273        self.reader
274    }
275}
276
277/// The future of [`fill_buf`](AsyncBufReadExt::fill_buf).
278#[derive(Debug)]
279#[must_use = "futures do nothing unless polled"]
280pub struct FillBuf<'a, R: ?Sized> {
281    reader: Option<&'a mut R>,
282}
283
284impl<'a, R: AsyncBufRead + Unpin + ?Sized> Future for FillBuf<'a, R> {
285    type Output = io::Result<&'a [u8]>;
286
287    #[expect(
288        clippy::panic,
289        reason = "polled after it completed, which the `Future` contract rules out"
290    )]
291    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&'a [u8]>> {
292        let this = self.get_mut();
293        let Some(reader) = this.reader.take() else {
294            panic!("`FillBuf` polled after it completed");
295        };
296        match Pin::new(&mut *reader).poll_fill_buf(cx) {
297            Poll::Ready(Ok(slice)) => {
298                // SAFETY: `reader` is the `&'a mut R` taken out of the future,
299                // which never uses it again, so the buffer it lends may live
300                // for `'a`. Only the borrow checker's handling of a borrow
301                // returned on one arm and dropped on the other needs this.
302                let slice = unsafe { core::mem::transmute::<&[u8], &'a [u8]>(slice) };
303                Poll::Ready(Ok(slice))
304            }
305            Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
306            Poll::Pending => {
307                this.reader = Some(reader);
308                Poll::Pending
309            }
310        }
311    }
312}
313
314/// Seeks on any [`AsyncSeek`].
315pub trait AsyncSeekExt: AsyncSeek {
316    /// Moves the stream's position, answering the new one from the start.
317    fn seek(&mut self, pos: SeekFrom) -> Seek<'_, Self>
318    where
319        Self: Unpin,
320    {
321        Seek { stream: self, pos }
322    }
323
324    /// Moves the stream's position to the start.
325    fn rewind(&mut self) -> Seek<'_, Self>
326    where
327        Self: Unpin,
328    {
329        self.seek(SeekFrom::Start(0))
330    }
331
332    /// The stream's position from the start.
333    fn stream_position(&mut self) -> Seek<'_, Self>
334    where
335        Self: Unpin,
336    {
337        self.seek(SeekFrom::Current(0))
338    }
339}
340
341impl<S: AsyncSeek + ?Sized> AsyncSeekExt for S {}
342
343/// The future of [`seek`](AsyncSeekExt::seek) and the seeks built on it.
344#[derive(Debug)]
345#[must_use = "futures do nothing unless polled"]
346pub struct Seek<'a, S: ?Sized> {
347    stream: &'a mut S,
348    pos: SeekFrom,
349}
350
351impl<S: AsyncSeek + Unpin + ?Sized> Future for Seek<'_, S> {
352    type Output = io::Result<u64>;
353
354    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
355        let this = self.get_mut();
356        Pin::new(&mut *this.stream).poll_seek(cx, this.pos)
357    }
358}
nuclei/io/util/buf_read_ext.rs

nuclei/io/util/
buf_read_ext.rs
