io: Add BufRead trait

Add a `BufRead` trait for types that perform buffered reading.

Implement it for:
- `Take`
- `Cursor`
- `std::io::BufRead` readers
- (in no-std builds) for slice of u8s
This commit is contained in:
Tobin C. Harding 2023-12-08 11:43:51 +11:00
parent 8315760403
commit 32d68fd1fa
No known key found for this signature in database
GPG Key ID: 40BF9E4C269D6607
1 changed files with 97 additions and 0 deletions

View File

@ -52,6 +52,19 @@ pub trait Read {
fn take(&mut self, limit: u64) -> Take<Self> { Take { reader: self, remaining: limit } } fn take(&mut self, limit: u64) -> Take<Self> { Take { reader: self, remaining: limit } }
} }
/// A trait describing an input stream that uses an internal buffer when reading.
pub trait BufRead: Read {
/// Returns data read from this reader, filling the internal buffer if needed.
fn fill_buf(&mut self) -> Result<&[u8]>;
/// Marks the buffered data up to amount as consumed.
///
/// # Panics
///
/// May panic if `amount` is greater than amount of data read by `fill_buf`.
fn consume(&mut self, amount: usize);
}
pub struct Take<'a, R: Read + ?Sized> { pub struct Take<'a, R: Read + ?Sized> {
reader: &'a mut R, reader: &'a mut R,
remaining: u64, remaining: u64,
@ -67,6 +80,30 @@ impl<'a, R: Read + ?Sized> Read for Take<'a, R> {
} }
} }
// Impl copied from Rust stdlib.
impl<'a, R: BufRead + ?Sized> BufRead for Take<'a, R> {
#[inline]
fn fill_buf(&mut self) -> Result<&[u8]> {
// Don't call into inner reader at all at EOF because it may still block
if self.remaining == 0 {
return Ok(&[]);
}
let buf = self.reader.fill_buf()?;
// Cast length to a u64 instead of casting `remaining` to a `usize`
// (in case `remaining > u32::MAX` and we are on a 32 bit machine).
let cap = cmp::min(buf.len() as u64, self.remaining) as usize;
Ok(&buf[..cap])
}
#[inline]
fn consume(&mut self, amount: usize) {
assert!(amount as u64 <= self.remaining);
self.remaining -= amount as u64;
self.reader.consume(amount);
}
}
#[cfg(feature = "std")] #[cfg(feature = "std")]
impl<R: std::io::Read> Read for R { impl<R: std::io::Read> Read for R {
#[inline] #[inline]
@ -75,6 +112,15 @@ impl<R: std::io::Read> Read for R {
} }
} }
#[cfg(feature = "std")]
impl<R: std::io::BufRead + Read + ?Sized> BufRead for R {
#[inline]
fn fill_buf(&mut self) -> Result<&[u8]> { Ok(std::io::BufRead::fill_buf(self)?) }
#[inline]
fn consume(&mut self, amount: usize) { std::io::BufRead::consume(self, amount) }
}
#[cfg(not(feature = "std"))] #[cfg(not(feature = "std"))]
impl Read for &[u8] { impl Read for &[u8] {
#[inline] #[inline]
@ -86,6 +132,16 @@ impl Read for &[u8] {
} }
} }
#[cfg(not(feature = "std"))]
impl BufRead for &[u8] {
#[inline]
fn fill_buf(&mut self) -> Result<&[u8]> { Ok(self) }
// This panics if amount is out of bounds, same as the std version.
#[inline]
fn consume(&mut self, amount: usize) { *self = &self[amount..] }
}
pub struct Cursor<T> { pub struct Cursor<T> {
inner: T, inner: T,
pos: u64, pos: u64,
@ -115,6 +171,20 @@ impl<T: AsRef<[u8]>> Read for Cursor<T> {
} }
} }
impl<T: AsRef<[u8]>> BufRead for Cursor<T> {
#[inline]
fn fill_buf(&mut self) -> Result<&[u8]> {
let inner: &[u8] = self.inner.as_ref();
Ok(&inner[self.pos as usize..])
}
#[inline]
fn consume(&mut self, amount: usize) {
assert!(amount <= self.inner.as_ref().len());
self.pos += amount as u64;
}
}
/// A generic trait describing an output stream. See [`std::io::Write`] for more info. /// A generic trait describing an output stream. See [`std::io::Write`] for more info.
pub trait Write { pub trait Write {
fn write(&mut self, buf: &[u8]) -> Result<usize>; fn write(&mut self, buf: &[u8]) -> Result<usize>;
@ -202,3 +272,30 @@ impl std::io::Write for Sink {
/// Returns a sink to which all writes succeed. See [`std::io::sink`] for more info. /// Returns a sink to which all writes succeed. See [`std::io::sink`] for more info.
#[inline] #[inline]
pub fn sink() -> Sink { Sink } pub fn sink() -> Sink { Sink }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn buf_read_fill_and_consume_slice() {
let data = [0_u8, 1, 2];
let mut slice = &data[..];
let fill = BufRead::fill_buf(&mut slice).unwrap();
assert_eq!(fill.len(), 3);
assert_eq!(fill, &[0_u8, 1, 2]);
slice.consume(2);
let fill = BufRead::fill_buf(&mut slice).unwrap();
assert_eq!(fill.len(), 1);
assert_eq!(fill, &[2_u8]);
slice.consume(1);
// checks we can attempt to read from a now-empty reader.
let fill = BufRead::fill_buf(&mut slice).unwrap();
assert_eq!(fill.len(), 0);
assert_eq!(fill, &[]);
}
}