Skip to main content

alloc/io/
util.rs

1use core::cmp;
2
3use crate::io::{
4    self, BorrowedBuf, BorrowedCursor, BufRead, Chain, Empty, IoSliceMut, Read, Repeat, Result,
5    SizeHint, Take,
6};
7use crate::slice;
8use crate::string::String;
9use crate::vec::Vec;
10
11#[stable(feature = "rust1", since = "1.0.0")]
12impl Read for Empty {
13    #[inline]
14    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
15        Ok(0)
16    }
17
18    #[inline]
19    fn read_buf(&mut self, _cursor: BorrowedCursor<'_, u8>) -> io::Result<()> {
20        Ok(())
21    }
22
23    #[inline]
24    fn read_vectored(&mut self, _bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
25        Ok(0)
26    }
27
28    #[inline]
29    fn is_read_vectored(&self) -> bool {
30        // Do not force `Chain<Empty, T>` or `Chain<T, Empty>` to use vectored
31        // reads, unless the other reader is vectored.
32        false
33    }
34
35    #[inline]
36    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
37        if !buf.is_empty() { Err(io::Error::READ_EXACT_EOF) } else { Ok(()) }
38    }
39
40    #[inline]
41    fn read_buf_exact(&mut self, cursor: BorrowedCursor<'_, u8>) -> io::Result<()> {
42        if cursor.capacity() != 0 { Err(io::Error::READ_EXACT_EOF) } else { Ok(()) }
43    }
44
45    #[inline]
46    fn read_to_end(&mut self, _buf: &mut Vec<u8>) -> io::Result<usize> {
47        Ok(0)
48    }
49
50    #[inline]
51    fn read_to_string(&mut self, _buf: &mut String) -> io::Result<usize> {
52        Ok(0)
53    }
54}
55
56#[stable(feature = "rust1", since = "1.0.0")]
57impl BufRead for Empty {
58    #[inline]
59    fn fill_buf(&mut self) -> io::Result<&[u8]> {
60        Ok(&[])
61    }
62
63    #[inline]
64    fn consume(&mut self, _n: usize) {}
65
66    #[inline]
67    fn has_data_left(&mut self) -> io::Result<bool> {
68        Ok(false)
69    }
70
71    #[inline]
72    fn read_until(&mut self, _byte: u8, _buf: &mut Vec<u8>) -> io::Result<usize> {
73        Ok(0)
74    }
75
76    #[inline]
77    fn skip_until(&mut self, _byte: u8) -> io::Result<usize> {
78        Ok(0)
79    }
80
81    #[inline]
82    fn read_line(&mut self, _buf: &mut String) -> io::Result<usize> {
83        Ok(0)
84    }
85}
86
87#[stable(feature = "rust1", since = "1.0.0")]
88impl Read for Repeat {
89    #[inline]
90    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
91        buf.fill(self.byte);
92        Ok(buf.len())
93    }
94
95    #[inline]
96    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
97        buf.fill(self.byte);
98        Ok(())
99    }
100
101    #[inline]
102    fn read_buf(&mut self, mut buf: BorrowedCursor<'_, u8>) -> io::Result<()> {
103        // SAFETY: No uninit bytes are being written.
104        unsafe { buf.as_mut() }.write_filled(self.byte);
105        // SAFETY: the entire unfilled portion of buf has been initialized.
106        unsafe { buf.advance(buf.capacity()) };
107        Ok(())
108    }
109
110    #[inline]
111    fn read_buf_exact(&mut self, buf: BorrowedCursor<'_, u8>) -> io::Result<()> {
112        self.read_buf(buf)
113    }
114
115    /// This function is not supported by `io::Repeat`, because there's no end of its data
116    fn read_to_end(&mut self, _: &mut Vec<u8>) -> io::Result<usize> {
117        Err(io::Error::from(io::ErrorKind::OutOfMemory))
118    }
119
120    /// This function is not supported by `io::Repeat`, because there's no end of its data
121    fn read_to_string(&mut self, _: &mut String) -> io::Result<usize> {
122        Err(io::Error::from(io::ErrorKind::OutOfMemory))
123    }
124
125    #[inline]
126    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
127        let mut nwritten = 0;
128        for buf in bufs {
129            nwritten += self.read(buf)?;
130        }
131        Ok(nwritten)
132    }
133
134    #[inline]
135    fn is_read_vectored(&self) -> bool {
136        true
137    }
138}
139
140#[stable(feature = "rust1", since = "1.0.0")]
141impl<T: Read, U: Read> Read for Chain<T, U> {
142    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
143        if !self.done_first {
144            match self.first.read(buf)? {
145                0 if !buf.is_empty() => self.done_first = true,
146                n => return Ok(n),
147            }
148        }
149        self.second.read(buf)
150    }
151
152    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> Result<usize> {
153        if !self.done_first {
154            match self.first.read_vectored(bufs)? {
155                0 if bufs.iter().any(|b| !b.is_empty()) => self.done_first = true,
156                n => return Ok(n),
157            }
158        }
159        self.second.read_vectored(bufs)
160    }
161
162    #[inline]
163    fn is_read_vectored(&self) -> bool {
164        self.first.is_read_vectored() || self.second.is_read_vectored()
165    }
166
167    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
168        let mut read = 0;
169        if !self.done_first {
170            read += self.first.read_to_end(buf)?;
171            self.done_first = true;
172        }
173        read += self.second.read_to_end(buf)?;
174        Ok(read)
175    }
176
177    // We don't override `read_to_string` here because an UTF-8 sequence could
178    // be split between the two parts of the chain
179
180    fn read_buf(&mut self, mut buf: BorrowedCursor<'_, u8>) -> Result<()> {
181        if buf.capacity() == 0 {
182            return Ok(());
183        }
184
185        if !self.done_first {
186            let old_len = buf.written();
187            self.first.read_buf(buf.reborrow())?;
188
189            if buf.written() != old_len {
190                return Ok(());
191            } else {
192                self.done_first = true;
193            }
194        }
195        self.second.read_buf(buf)
196    }
197}
198
199#[stable(feature = "chain_bufread", since = "1.9.0")]
200impl<T: BufRead, U: BufRead> BufRead for Chain<T, U> {
201    fn fill_buf(&mut self) -> Result<&[u8]> {
202        if !self.done_first {
203            match self.first.fill_buf()? {
204                buf if buf.is_empty() => self.done_first = true,
205                buf => return Ok(buf),
206            }
207        }
208        self.second.fill_buf()
209    }
210
211    fn consume(&mut self, amt: usize) {
212        if !self.done_first { self.first.consume(amt) } else { self.second.consume(amt) }
213    }
214
215    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> Result<usize> {
216        let mut read = 0;
217        if !self.done_first {
218            let n = self.first.read_until(byte, buf)?;
219            read += n;
220
221            match buf.last() {
222                Some(b) if *b == byte && n != 0 => return Ok(read),
223                _ => self.done_first = true,
224            }
225        }
226        read += self.second.read_until(byte, buf)?;
227        Ok(read)
228    }
229
230    // We don't override `read_line` here because an UTF-8 sequence could be
231    // split between the two parts of the chain
232}
233
234#[stable(feature = "rust1", since = "1.0.0")]
235impl<T: Read> Read for Take<T> {
236    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
237        // Don't call into inner reader at all at EOF because it may still block
238        if self.limit == 0 {
239            return Ok(0);
240        }
241
242        let max = cmp::min(buf.len() as u64, self.limit) as usize;
243        let n = self.inner.read(&mut buf[..max])?;
244        assert!(n as u64 <= self.limit, "number of read bytes exceeds limit");
245        self.limit -= n as u64;
246        Ok(n)
247    }
248
249    fn read_buf(&mut self, mut buf: BorrowedCursor<'_, u8>) -> Result<()> {
250        // Don't call into inner reader at all at EOF because it may still block
251        if self.limit == 0 {
252            return Ok(());
253        }
254
255        if self.limit < buf.capacity() as u64 {
256            // The condition above guarantees that `self.limit` fits in `usize`.
257            let limit = self.limit as usize;
258
259            let is_init = buf.is_init();
260
261            // SAFETY: no uninit data is written to ibuf
262            let mut sliced_buf = BorrowedBuf::from(unsafe { &mut buf.as_mut()[..limit] });
263
264            if is_init {
265                // SAFETY: `sliced_buf` is a subslice of `buf`, so if `buf` was initialized then
266                // `sliced_buf` is.
267                unsafe { sliced_buf.set_init() };
268            }
269
270            let result = self.inner.read_buf(sliced_buf.unfilled());
271
272            let did_init_up_to_limit = sliced_buf.is_init();
273            let filled = sliced_buf.len();
274
275            // sliced_buf must drop here
276
277            // Avoid accidentally quadratic behaviour by initializing the whole
278            // cursor if only part of it was initialized.
279            if did_init_up_to_limit && !is_init {
280                // SAFETY: No uninit data will be written.
281                let unfilled_before_advance = unsafe { buf.as_mut() };
282
283                unfilled_before_advance[limit..].write_filled(0);
284
285                // SAFETY: `unfilled_before_advance[..limit]` was initialized by `T::read_buf`, and
286                // `unfilled_before_advance[limit..]` was just initialized.
287                unsafe { buf.set_init() };
288            }
289
290            unsafe {
291                // SAFETY: filled bytes have been filled
292                buf.advance(filled);
293            }
294
295            self.limit -= filled as u64;
296
297            result
298        } else {
299            let written = buf.written();
300            let result = self.inner.read_buf(buf.reborrow());
301            self.limit -= (buf.written() - written) as u64;
302            result
303        }
304    }
305}
306
307#[stable(feature = "rust1", since = "1.0.0")]
308impl<T: BufRead> BufRead for Take<T> {
309    fn fill_buf(&mut self) -> Result<&[u8]> {
310        // Don't call into inner reader at all at EOF because it may still block
311        if self.limit == 0 {
312            return Ok(&[]);
313        }
314
315        let buf = self.inner.fill_buf()?;
316        let cap = cmp::min(buf.len() as u64, self.limit) as usize;
317        Ok(&buf[..cap])
318    }
319
320    fn consume(&mut self, amt: usize) {
321        // Don't let callers reset the limit by passing an overlarge value
322        let amt = cmp::min(amt as u64, self.limit) as usize;
323        self.limit -= amt as u64;
324        self.inner.consume(amt);
325    }
326}
327
328/// An iterator over `u8` values of a reader.
329///
330/// This struct is generally created by calling [`bytes`] on a reader.
331/// Please see the documentation of [`bytes`] for more details.
332///
333/// [`bytes`]: Read::bytes
334#[stable(feature = "rust1", since = "1.0.0")]
335#[derive(Debug)]
336pub struct Bytes<R> {
337    inner: R,
338}
339
340#[stable(feature = "rust1", since = "1.0.0")]
341impl<R: Read> Iterator for Bytes<R> {
342    type Item = Result<u8>;
343
344    // Not `#[inline]`. This function gets inlined even without it, but having
345    // the inline annotation can result in worse code generation. See #116785.
346    fn next(&mut self) -> Option<Result<u8>> {
347        SpecReadByte::spec_read_byte(&mut self.inner)
348    }
349
350    #[inline]
351    fn size_hint(&self) -> (usize, Option<usize>) {
352        SizeHint::size_hint(&self.inner)
353    }
354}
355
356/// For the specialization of `Bytes::next`.
357#[doc(hidden)]
358#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
359pub trait SpecReadByte {
360    fn spec_read_byte(&mut self) -> Option<Result<u8>>;
361}
362
363impl<R> SpecReadByte for R
364where
365    Self: Read,
366{
367    #[inline]
368    default fn spec_read_byte(&mut self) -> Option<Result<u8>> {
369        inlined_slow_read_byte(self)
370    }
371}
372
373/// Reads a single byte in a slow, generic way. This is used by the default
374/// `spec_read_byte`.
375#[inline]
376fn inlined_slow_read_byte<R: Read>(reader: &mut R) -> Option<Result<u8>> {
377    let mut byte = 0;
378    loop {
379        return match reader.read(slice::from_mut(&mut byte)) {
380            Ok(0) => None,
381            Ok(..) => Some(Ok(byte)),
382            Err(ref e) if e.is_interrupted() => continue,
383            Err(e) => Some(Err(e)),
384        };
385    }
386}
387
388// Used by `BufReader::spec_read_byte`, for which the `inline(never)` is
389// important.
390#[inline(never)]
391#[doc(hidden)]
392#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
393pub fn uninlined_slow_read_byte<R: Read>(reader: &mut R) -> Option<Result<u8>> {
394    inlined_slow_read_byte(reader)
395}
396
397#[doc(hidden)]
398#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
399pub const fn bytes<R>(inner: R) -> Bytes<R> {
400    Bytes { inner }
401}
402
403/// An iterator over the contents of an instance of `BufRead` split on a
404/// particular byte.
405///
406/// This struct is generally created by calling [`split`] on a `BufRead`.
407/// Please see the documentation of [`split`] for more details.
408///
409/// [`split`]: BufRead::split
410#[stable(feature = "rust1", since = "1.0.0")]
411#[derive(Debug)]
412#[cfg_attr(not(test), rustc_diagnostic_item = "IoSplit")]
413pub struct Split<B> {
414    buf: B,
415    delim: u8,
416}
417
418#[stable(feature = "rust1", since = "1.0.0")]
419impl<B: BufRead> Iterator for Split<B> {
420    type Item = Result<Vec<u8>>;
421
422    fn next(&mut self) -> Option<Result<Vec<u8>>> {
423        let mut buf = Vec::new();
424        match self.buf.read_until(self.delim, &mut buf) {
425            Ok(0) => None,
426            Ok(_n) => {
427                if buf[buf.len() - 1] == self.delim {
428                    buf.pop();
429                }
430                Some(Ok(buf))
431            }
432            Err(e) => Some(Err(e)),
433        }
434    }
435}
436
437#[doc(hidden)]
438#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
439pub const fn split<B>(buf: B, delim: u8) -> Split<B> {
440    Split { buf, delim }
441}
442
443/// An iterator over the lines of an instance of [`BufRead`].
444///
445/// This struct is generally created by calling [`lines`] on a [`BufRead`].
446/// Please see the documentation of [`lines`] for more details.
447///
448/// [`lines`]: BufRead::lines
449#[stable(feature = "rust1", since = "1.0.0")]
450#[derive(Debug)]
451#[cfg_attr(not(test), rustc_diagnostic_item = "IoLines")]
452pub struct Lines<B> {
453    buf: B,
454}
455
456#[stable(feature = "rust1", since = "1.0.0")]
457impl<B: BufRead> Iterator for Lines<B> {
458    type Item = Result<String>;
459
460    fn next(&mut self) -> Option<Result<String>> {
461        let mut buf = String::new();
462        match self.buf.read_line(&mut buf) {
463            Ok(0) => None,
464            Ok(_n) => {
465                if buf.ends_with('\n') {
466                    buf.pop();
467                    if buf.ends_with('\r') {
468                        buf.pop();
469                    }
470                }
471                Some(Ok(buf))
472            }
473            Err(e) => Some(Err(e)),
474        }
475    }
476}
477
478#[doc(hidden)]
479#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
480pub const fn lines<B>(buf: B) -> Lines<B> {
481    Lines { buf }
482}