diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f4209e7..d4e8d30 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,7 +10,6 @@ jobs: matrix: rust-toolchain: - nightly - - nightly-2025-05-20 - nightly-2025-12-12 - stable env: diff --git a/Cargo.toml b/Cargo.toml index d835d00..3c17211 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,9 +1,9 @@ [package] name = "axio" version = "0.2.1" -edition = "2021" +edition = "2024" authors = ["Yuekai Jia "] -description = "`std::io`-like I/O traits for `no_std` environment" +description = "`std::io` for `no_std` environment" license = "GPL-3.0-or-later OR Apache-2.0 OR MulanPSL-2.0" homepage = "https://github.com/arceos-org/arceos" repository = "https://github.com/arceos-org/axio" @@ -16,7 +16,8 @@ alloc = [] default = [] [dependencies] -axerrno = "0.1" +axerrno = "0.2" +memchr = { version = "2", default-features = false } [build-dependencies] autocfg = "1" diff --git a/README.md b/README.md index 0ee809b..92e8816 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,38 @@ [![Docs.rs](https://docs.rs/axio/badge.svg)](https://docs.rs/axio) [![CI](https://github.com/arceos-org/axio/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/arceos-org/axio/actions/workflows/ci.yml) -[`std::io`][1]-like I/O traits for `no_std` environment. +[`std::io`][1] for `no_std` environment. [1]: https://doc.rust-lang.org/std/io/index.html + +### Features + +- **alloc**: + - Enables extra methods on `Read`: `read_to_end`, `read_to_string`. + - Enables extra methods on `BufRead`: `read_until`, `read_line`, `split`, `lines`. + - Enables implementations of axio traits for `alloc` types like `Vec`, `Box`, etc. + +### Differences to `std::io` + +- Error types from `axerrno` instead of `std::io::Error`. +- No `IoSlice` and `*_vectored` APIs. + +### Limitations + +- Requires nightly Rust. + +## License + +Licensed under either of + +- GNU General Public License v3.0 or later, () +- Apache License, Version 2.0, () +- Mulan Permissive Software License, Version 2, () + +at your option. + +--- + +Almost all of the code in this repository is a copy of the [Rust language codebase](https://github.com/rust-lang/rust) with minor modifications. + +For attributions, see . diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..9c4bd7d --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,18 @@ +unstable_features = true + +style_edition = "2024" + +group_imports = "StdExternalCrate" +imports_granularity = "Crate" + +normalize_comments = true +wrap_comments = true +comment_width = 100 + +condense_wildcard_suffixes = true +enum_discrim_align_threshold = 20 +use_field_init_shorthand = true + +format_strings = true +format_code_in_doc_comments = true +format_macro_matchers = true diff --git a/src/buffered/bufreader.rs b/src/buffered/bufreader.rs index b466c1e..e65ef77 100644 --- a/src/buffered/bufreader.rs +++ b/src/buffered/bufreader.rs @@ -1,9 +1,7 @@ -use crate::{BufRead, Read, Result}; - #[cfg(feature = "alloc")] use alloc::{string::String, vec::Vec}; -const DEFAULT_BUF_SIZE: usize = 1024; +use crate::{BufRead, DEFAULT_BUF_SIZE, Read, Result}; /// The `BufReader` struct adds buffering to any reader. pub struct BufReader { diff --git a/src/error.rs b/src/error.rs deleted file mode 100644 index ee2d716..0000000 --- a/src/error.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub use axerrno::AxError as Error; -pub use axerrno::AxResult as Result; diff --git a/src/impls.rs b/src/impls.rs deleted file mode 100644 index 96f83db..0000000 --- a/src/impls.rs +++ /dev/null @@ -1,54 +0,0 @@ -use crate::{prelude::*, Result}; -use core::cmp; - -impl Read for &[u8] { - #[inline] - fn read(&mut self, buf: &mut [u8]) -> Result { - let amt = cmp::min(buf.len(), self.len()); - let a = &self[..amt]; - let b = &self[amt..]; - - // First check if the amount of bytes we want to read is small: - // `copy_from_slice` will generally expand to a call to `memcpy`, and - // for a single byte the overhead is significant. - if amt == 1 { - buf[0] = a[0]; - } else { - buf[..amt].copy_from_slice(a); - } - - *self = b; - Ok(amt) - } - - #[inline] - fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { - if buf.len() > self.len() { - return axerrno::ax_err!(UnexpectedEof, "failed to fill whole buffer"); - } - let amt = buf.len(); - let a = &self[..amt]; - let b = &self[amt..]; - - // First check if the amount of bytes we want to read is small: - // `copy_from_slice` will generally expand to a call to `memcpy`, and - // for a single byte the overhead is significant. - if amt == 1 { - buf[0] = a[0]; - } else { - buf[..amt].copy_from_slice(a); - } - - *self = b; - Ok(()) - } - - #[inline] - #[cfg(feature = "alloc")] - fn read_to_end(&mut self, buf: &mut alloc::vec::Vec) -> Result { - buf.extend_from_slice(self); - let len = self.len(); - *self = &self[len..]; - Ok(len) - } -} diff --git a/src/lib.rs b/src/lib.rs index 34c1fb8..564600a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,397 +1,26 @@ -//! [`std::io`]-like I/O traits for `no_std` environment. - +#![doc = include_str!("../README.md")] #![cfg_attr(not(doc), no_std)] #![feature(doc_cfg)] #![feature(core_io_borrowed_buf)] #![cfg_attr(not(borrowedbuf_init), feature(maybe_uninit_fill))] +#![warn(missing_docs)] #[cfg(feature = "alloc")] extern crate alloc; -use core::fmt; +#[doc(no_inline)] +pub use axerrno::{AxError as Error, AxErrorKind as ErrorKind, AxResult as Result}; -mod buffered; -mod error; -mod impls; +/// Default buffer size for I/O operations. +pub const DEFAULT_BUF_SIZE: usize = 1024 * 2; +mod buffered; pub mod prelude; +mod read; +mod seek; +mod write; -pub use self::buffered::BufReader; -pub use self::error::{Error, Result}; - -#[cfg(feature = "alloc")] -use alloc::{string::String, vec::Vec}; - -use axerrno::ax_err; - -/// Default [`Read::read_to_end`] implementation with optional size hint. -/// -/// Adapted from [`std::io::default_read_to_end`]. -/// -/// [`std::io::default_read_to_end`]: https://github.com/rust-lang/rust/blob/30f168ef811aec63124eac677e14699baa9395bd/library/std/src/io/mod.rs#L409 -#[cfg(feature = "alloc")] -pub fn default_read_to_end( - r: &mut R, - buf: &mut Vec, - size_hint: Option, -) -> Result { - use core::io::BorrowedBuf; - - const DEFAULT_BUF_SIZE: usize = 1024; - - let start_len = buf.len(); - let start_cap = buf.capacity(); - // Optionally limit the maximum bytes read on each iteration. - // This adds an arbitrary fiddle factor to allow for more data than we expect. - let mut max_read_size = size_hint - .and_then(|s| { - s.checked_add(1024)? - .checked_next_multiple_of(DEFAULT_BUF_SIZE) - }) - .unwrap_or(DEFAULT_BUF_SIZE); - - const PROBE_SIZE: usize = 32; - - fn small_probe_read(r: &mut R, buf: &mut Vec) -> Result { - let mut probe = [0u8; PROBE_SIZE]; - - let n = r.read(&mut probe)?; - buf.extend_from_slice(&probe[..n]); - Ok(n) - } - - if (size_hint.is_none() || size_hint == Some(0)) && buf.capacity() - buf.len() < PROBE_SIZE { - let read = small_probe_read(r, buf)?; - - if read == 0 { - return Ok(0); - } - } - - #[cfg(borrowedbuf_init)] - let mut initialized = 0; // Extra initialized bytes from previous loop iteration - #[cfg(borrowedbuf_init)] - let mut consecutive_short_reads = 0; - - loop { - if buf.len() == buf.capacity() && buf.capacity() == start_cap { - // The buffer might be an exact fit. Let's read into a probe buffer - // and see if it returns `Ok(0)`. If so, we've avoided an - // unnecessary doubling of the capacity. But if not, append the - // probe buffer to the primary buffer and let its capacity grow. - let read = small_probe_read(r, buf)?; - - if read == 0 { - return Ok(buf.len() - start_len); - } - } - - if buf.len() == buf.capacity() { - // buf is full, need more space - if let Err(e) = buf.try_reserve(PROBE_SIZE) { - return ax_err!(NoMemory, e); - } - } - - let mut spare = buf.spare_capacity_mut(); - let buf_len = spare.len().min(max_read_size); - spare = &mut spare[..buf_len]; - let mut read_buf: BorrowedBuf<'_> = spare.into(); - - #[cfg(borrowedbuf_init)] - // SAFETY: These bytes were initialized but not filled in the previous loop - unsafe { - read_buf.set_init(initialized); - } - - let mut cursor = read_buf.unfilled(); - // Difference from `std`: We don't have a `read_buf` method that returns both data and an error, so we return early on error. - #[cfg(borrowedbuf_init)] - { - let n = r.read(cursor.ensure_init().init_mut())?; - cursor.advance(n); - } - #[cfg(not(borrowedbuf_init))] - { - // SAFETY: We do not uninitialize any part of the buffer. - let n = r.read(unsafe { cursor.as_mut().write_filled(0) })?; - assert!(n <= cursor.capacity()); - // SAFETY: We've initialized the entire buffer, and `read` can't make it uninitialized. - unsafe { - cursor.advance(n); - } - } - - #[cfg(borrowedbuf_init)] - let unfilled_but_initialized = cursor.init_mut().len(); - let bytes_read = cursor.written(); - #[cfg(borrowedbuf_init)] - let was_fully_initialized = read_buf.init_len() == buf_len; - - // SAFETY: BorrowedBuf's invariants mean this much memory is initialized. - unsafe { - let new_len = bytes_read + buf.len(); - buf.set_len(new_len); - } - - if bytes_read == 0 { - return Ok(buf.len() - start_len); - } - - #[cfg(borrowedbuf_init)] - if bytes_read < buf_len { - consecutive_short_reads += 1; - } else { - consecutive_short_reads = 0; - } - - #[cfg(borrowedbuf_init)] - { - // store how much was initialized but not filled - initialized = unfilled_but_initialized; - } - - // Use heuristics to determine the max read size if no initial size hint was provided - if size_hint.is_none() { - #[cfg(borrowedbuf_init)] - // The reader is returning short reads but it doesn't call ensure_init(). - // In that case we no longer need to restrict read sizes to avoid - // initialization costs. - // When reading from disk we usually don't get any short reads except at EOF. - // So we wait for at least 2 short reads before uncapping the read buffer; - // this helps with the Windows issue. - if !was_fully_initialized && consecutive_short_reads > 1 { - max_read_size = usize::MAX; - } - - // we have passed a larger buffer than previously and the - // reader still hasn't returned a short read - if buf_len >= max_read_size && bytes_read == buf_len { - max_read_size = max_read_size.saturating_mul(2); - } - } - } -} - -/// The `Read` trait allows for reading bytes from a source. -pub trait Read { - /// Pull some bytes from this source into the specified buffer, returning - /// how many bytes were read. - fn read(&mut self, buf: &mut [u8]) -> Result; - - /// Read all bytes until EOF in this source, placing them into `buf`. - #[cfg(feature = "alloc")] - fn read_to_end(&mut self, buf: &mut Vec) -> Result { - default_read_to_end(self, buf, None) - } - - /// Read all bytes until EOF in this source, appending them to `buf`. - #[cfg(feature = "alloc")] - fn read_to_string(&mut self, buf: &mut String) -> Result { - unsafe { append_to_string(buf, |b| self.read_to_end(b)) } - } - - /// Read the exact number of bytes required to fill `buf`. - fn read_exact(&mut self, mut buf: &mut [u8]) -> Result { - while !buf.is_empty() { - match self.read(buf) { - Ok(0) => break, - Ok(n) => { - let tmp = buf; - buf = &mut tmp[n..]; - } - Err(e) => return Err(e), - } - } - if !buf.is_empty() { - ax_err!(UnexpectedEof, "failed to fill whole buffer") - } else { - Ok(()) - } - } -} - -/// A trait for objects which are byte-oriented sinks. -pub trait Write { - /// Write a buffer into this writer, returning how many bytes were written. - fn write(&mut self, buf: &[u8]) -> Result; - - /// Flush this output stream, ensuring that all intermediately buffered - /// contents reach their destination. - fn flush(&mut self) -> Result; - - /// Attempts to write an entire buffer into this writer. - fn write_all(&mut self, mut buf: &[u8]) -> Result { - while !buf.is_empty() { - match self.write(buf) { - Ok(0) => return ax_err!(WriteZero, "failed to write whole buffer"), - Ok(n) => buf = &buf[n..], - Err(e) => return Err(e), - } - } - Ok(()) - } - - /// Writes a formatted string into this writer, returning any error - /// encountered. - fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> Result<()> { - // Create a shim which translates a Write to a fmt::Write and saves - // off I/O errors. instead of discarding them - struct Adapter<'a, T: ?Sized + 'a> { - inner: &'a mut T, - error: Result<()>, - } - - impl fmt::Write for Adapter<'_, T> { - fn write_str(&mut self, s: &str) -> fmt::Result { - match self.inner.write_all(s.as_bytes()) { - Ok(()) => Ok(()), - Err(e) => { - self.error = Err(e); - Err(fmt::Error) - } - } - } - } - - let mut output = Adapter { - inner: self, - error: Ok(()), - }; - match fmt::write(&mut output, fmt) { - Ok(()) => Ok(()), - Err(..) => { - // check if the error came from the underlying `Write` or not - if output.error.is_err() { - output.error - } else { - ax_err!(InvalidData, "formatter error") - } - } - } - } -} - -/// The `Seek` trait provides a cursor which can be moved within a stream of -/// bytes. -pub trait Seek { - /// Seek to an offset, in bytes, in a stream. - /// - /// A seek beyond the end of a stream is allowed, but behavior is defined - /// by the implementation. - /// - /// If the seek operation completed successfully, - /// this method returns the new position from the start of the stream. - /// That position can be used later with [`SeekFrom::Start`]. - fn seek(&mut self, pos: SeekFrom) -> Result; - - /// Rewind to the beginning of a stream. - /// - /// This is a convenience method, equivalent to `seek(SeekFrom::Start(0))`. - fn rewind(&mut self) -> Result<()> { - self.seek(SeekFrom::Start(0))?; - Ok(()) - } - - /// Returns the current seek position from the start of the stream. - /// - /// This is equivalent to `self.seek(SeekFrom::Current(0))`. - fn stream_position(&mut self) -> Result { - self.seek(SeekFrom::Current(0)) - } -} - -/// Enumeration of possible methods to seek within an I/O object. -/// -/// It is used by the [`Seek`] trait. -#[derive(Copy, PartialEq, Eq, Clone, Debug)] -pub enum SeekFrom { - /// Sets the offset to the provided number of bytes. - Start(u64), - - /// Sets the offset to the size of this object plus the specified number of - /// bytes. - /// - /// It is possible to seek beyond the end of an object, but it's an error to - /// seek before byte 0. - End(i64), - - /// Sets the offset to the current position plus the specified number of - /// bytes. - /// - /// It is possible to seek beyond the end of an object, but it's an error to - /// seek before byte 0. - Current(i64), -} - -/// A `BufRead` is a type of `Read`er which has an internal buffer, allowing it -/// to perform extra ways of reading. -pub trait BufRead: Read { - /// Returns the contents of the internal buffer, filling it with more data - /// from the inner reader if it is empty. - fn fill_buf(&mut self) -> Result<&[u8]>; - - /// Tells this buffer that `amt` bytes have been consumed from the buffer, - /// so they should no longer be returned in calls to `read`. - fn consume(&mut self, amt: usize); - - /// Check if the underlying `Read` has any data left to be read. - fn has_data_left(&mut self) -> Result { - self.fill_buf().map(|b| !b.is_empty()) - } - - /// Read all bytes into `buf` until the delimiter `byte` or EOF is reached. - #[cfg(feature = "alloc")] - fn read_until(&mut self, byte: u8, buf: &mut Vec) -> Result { - let mut read = 0; - loop { - let (done, used) = { - let available = match self.fill_buf() { - Ok(n) => n, - Err(Error::WouldBlock) => continue, - Err(e) => return Err(e), - }; - match available.iter().position(|&b| b == byte) { - Some(i) => { - buf.extend_from_slice(&available[..=i]); - (true, i + 1) - } - None => { - buf.extend_from_slice(available); - (false, available.len()) - } - } - }; - self.consume(used); - read += used; - if done || used == 0 { - return Ok(read); - } - } - } - - /// Read all bytes until a newline (the `0xA` byte) is reached, and append - /// them to the provided `String` buffer. - #[cfg(feature = "alloc")] - fn read_line(&mut self, buf: &mut String) -> Result { - unsafe { append_to_string(buf, |b| self.read_until(b'\n', b)) } - } -} - -#[cfg(feature = "alloc")] -unsafe fn append_to_string(buf: &mut String, f: F) -> Result -where - F: FnOnce(&mut Vec) -> Result, -{ - let old_len = buf.len(); - let buf = unsafe { buf.as_mut_vec() }; - let ret = f(buf)?; - if core::str::from_utf8(&buf[old_len..]).is_err() { - ax_err!(InvalidData, "stream did not contain valid UTF-8") - } else { - Ok(ret) - } -} +pub use self::{buffered::*, read::*, seek::*, write::*}; /// I/O poll results. #[derive(Debug, Default, Clone, Copy)] diff --git a/src/prelude.rs b/src/prelude.rs index 9f8020c..6116972 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -5,7 +5,7 @@ //! //! ``` //! # #![allow(unused_imports)] -//! use std::io::prelude::*; +//! use axio::prelude::*; //! ``` -pub use super::{BufRead, Read, Seek, Write}; +pub use crate::{BufRead, Read, Seek, Write}; diff --git a/src/read/impls.rs b/src/read/impls.rs new file mode 100644 index 0000000..3a00dca --- /dev/null +++ b/src/read/impls.rs @@ -0,0 +1,357 @@ +#[cfg(feature = "alloc")] +use alloc::{boxed::Box, collections::VecDeque, string::String, vec::Vec}; +use core::{cmp, io::BorrowedCursor}; + +use crate::{BufRead, Error, Read, Result}; + +// ============================================================================= +// Forwarding implementations + +impl Read for &mut R { + #[inline] + fn read(&mut self, buf: &mut [u8]) -> Result { + (**self).read(buf) + } + + #[inline] + fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { + (**self).read_exact(buf) + } + + #[inline] + fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> Result<()> { + (**self).read_buf(cursor) + } + + #[inline] + fn read_buf_exact(&mut self, cursor: BorrowedCursor<'_>) -> Result<()> { + (**self).read_buf_exact(cursor) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_end(&mut self, buf: &mut Vec) -> Result { + (**self).read_to_end(buf) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_string(&mut self, buf: &mut String) -> Result { + (**self).read_to_string(buf) + } +} + +impl BufRead for &mut B { + #[inline] + fn fill_buf(&mut self) -> Result<&[u8]> { + (**self).fill_buf() + } + + #[inline] + fn consume(&mut self, amt: usize) { + (**self).consume(amt) + } + + #[inline] + fn has_data_left(&mut self) -> Result { + (**self).has_data_left() + } + + #[inline] + fn skip_until(&mut self, byte: u8) -> Result { + (**self).skip_until(byte) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_until(&mut self, byte: u8, buf: &mut Vec) -> Result { + (**self).read_until(byte, buf) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_line(&mut self, buf: &mut String) -> Result { + (**self).read_line(buf) + } +} + +#[cfg(feature = "alloc")] +impl Read for Box { + #[inline] + fn read(&mut self, buf: &mut [u8]) -> Result { + (**self).read(buf) + } + + #[inline] + fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { + (**self).read_exact(buf) + } + + #[inline] + fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> Result<()> { + (**self).read_buf(cursor) + } + + #[inline] + fn read_buf_exact(&mut self, cursor: BorrowedCursor<'_>) -> Result<()> { + (**self).read_buf_exact(cursor) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_end(&mut self, buf: &mut Vec) -> Result { + (**self).read_to_end(buf) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_string(&mut self, buf: &mut String) -> Result { + (**self).read_to_string(buf) + } +} + +#[cfg(feature = "alloc")] +impl BufRead for Box { + #[inline] + fn fill_buf(&mut self) -> Result<&[u8]> { + (**self).fill_buf() + } + + #[inline] + fn consume(&mut self, amt: usize) { + (**self).consume(amt) + } + + #[inline] + fn has_data_left(&mut self) -> Result { + (**self).has_data_left() + } + + #[inline] + fn skip_until(&mut self, byte: u8) -> Result { + (**self).skip_until(byte) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_until(&mut self, byte: u8, buf: &mut Vec) -> Result { + (**self).read_until(byte, buf) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_line(&mut self, buf: &mut String) -> Result { + (**self).read_line(buf) + } +} + +// ============================================================================= +// In-memory buffer implementations + +impl Read for &[u8] { + #[inline] + fn read(&mut self, buf: &mut [u8]) -> Result { + let amt = cmp::min(buf.len(), self.len()); + let (a, b) = self.split_at(amt); + + // First check if the amount of bytes we want to read is small: + // `copy_from_slice` will generally expand to a call to `memcpy`, and + // for a single byte the overhead is significant. + if amt == 1 { + buf[0] = a[0]; + } else { + buf[..amt].copy_from_slice(a); + } + + *self = b; + Ok(amt) + } + + #[inline] + fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { + if buf.len() > self.len() { + // `read_exact` makes no promise about the content of `buf` if it + // fails so don't bother about that. + *self = &self[self.len()..]; + return Err(Error::UnexpectedEof); + } + let (a, b) = self.split_at(buf.len()); + + // First check if the amount of bytes we want to read is small: + // `copy_from_slice` will generally expand to a call to `memcpy`, and + // for a single byte the overhead is significant. + if buf.len() == 1 { + buf[0] = a[0]; + } else { + buf.copy_from_slice(a); + } + + *self = b; + Ok(()) + } + + #[inline] + fn read_buf(&mut self, mut cursor: BorrowedCursor<'_>) -> Result<()> { + let amt = cmp::min(cursor.capacity(), self.len()); + let (a, b) = self.split_at(amt); + + cursor.append(a); + + *self = b; + Ok(()) + } + + #[inline] + fn read_buf_exact(&mut self, mut cursor: BorrowedCursor<'_>) -> Result<()> { + if cursor.capacity() > self.len() { + // Append everything we can to the cursor. + cursor.append(self); + *self = &self[self.len()..]; + return Err(Error::UnexpectedEof); + } + let (a, b) = self.split_at(cursor.capacity()); + + cursor.append(a); + + *self = b; + Ok(()) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_end(&mut self, buf: &mut Vec) -> Result { + let len = self.len(); + buf.try_reserve(len).map_err(|_| Error::NoMemory)?; + buf.extend_from_slice(self); + *self = &self[len..]; + Ok(len) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_string(&mut self, buf: &mut String) -> Result { + let content = str::from_utf8(self).map_err(|_| Error::IllegalBytes)?; + let len = self.len(); + buf.try_reserve(len).map_err(|_| Error::NoMemory)?; + buf.push_str(content); + *self = &self[len..]; + Ok(len) + } +} + +impl BufRead for &[u8] { + #[inline] + fn fill_buf(&mut self) -> Result<&[u8]> { + Ok(*self) + } + + #[inline] + fn consume(&mut self, amt: usize) { + *self = &self[amt..]; + } +} + +#[cfg(feature = "alloc")] +impl Read for VecDeque { + #[inline] + fn read(&mut self, buf: &mut [u8]) -> Result { + let (ref mut front, _) = self.as_slices(); + let n = Read::read(front, buf)?; + self.drain(..n); + Ok(n) + } + + #[inline] + fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { + let (front, back) = self.as_slices(); + + // Use only the front buffer if it is big enough to fill `buf`, else use + // the back buffer too. + match buf.split_at_mut_checked(front.len()) { + None => buf.copy_from_slice(&front[..buf.len()]), + Some((buf_front, buf_back)) => match back.split_at_checked(buf_back.len()) { + Some((back, _)) => { + buf_front.copy_from_slice(front); + buf_back.copy_from_slice(back); + } + None => { + self.clear(); + return Err(Error::UnexpectedEof); + } + }, + } + + self.drain(..buf.len()); + Ok(()) + } + + #[inline] + fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> Result<()> { + let (ref mut front, _) = self.as_slices(); + let n = cmp::min(cursor.capacity(), front.len()); + Read::read_buf(front, cursor)?; + self.drain(..n); + Ok(()) + } + + #[inline] + fn read_buf_exact(&mut self, mut cursor: BorrowedCursor<'_>) -> Result<()> { + let len = cursor.capacity(); + let (front, back) = self.as_slices(); + + match front.split_at_checked(cursor.capacity()) { + Some((front, _)) => cursor.append(front), + None => { + cursor.append(front); + match back.split_at_checked(cursor.capacity()) { + Some((back, _)) => cursor.append(back), + None => { + cursor.append(back); + self.clear(); + return Err(Error::UnexpectedEof); + } + } + } + } + + self.drain(..len); + Ok(()) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_end(&mut self, buf: &mut Vec) -> Result { + // The total len is known upfront so we can reserve it in a single call. + let len = self.len(); + buf.try_reserve(len).map_err(|_| Error::NoMemory)?; + let (front, back) = self.as_slices(); + buf.extend_from_slice(front); + buf.extend_from_slice(back); + self.clear(); + Ok(len) + } + + #[inline] + #[cfg(feature = "alloc")] + fn read_to_string(&mut self, buf: &mut String) -> Result { + // SAFETY: We only append to the buffer + unsafe { super::append_to_string(buf, |buf| self.read_to_end(buf)) } + } +} + +#[cfg(feature = "alloc")] +impl BufRead for VecDeque { + /// Returns the contents of the "front" slice as returned by + /// [`as_slices`][`VecDeque::as_slices`]. If the contained byte slices of the `VecDeque` are + /// discontiguous, multiple calls to `fill_buf` will be needed to read the entire content. + #[inline] + fn fill_buf(&mut self) -> Result<&[u8]> { + let (front, _) = self.as_slices(); + Ok(front) + } + + #[inline] + fn consume(&mut self, amt: usize) { + self.drain(..amt); + } +} diff --git a/src/read/mod.rs b/src/read/mod.rs new file mode 100644 index 0000000..a486914 --- /dev/null +++ b/src/read/mod.rs @@ -0,0 +1,493 @@ +#[cfg(feature = "alloc")] +use alloc::{string::String, vec::Vec}; +use core::io::BorrowedCursor; + +use crate::{Error, Result}; + +mod impls; + +/// Default [`Read::read_exact`] implementation. +pub fn default_read_exact(this: &mut R, mut buf: &mut [u8]) -> Result<()> { + while !buf.is_empty() { + match this.read(buf) { + Ok(0) => break, + Ok(n) => { + buf = &mut buf[n..]; + } + Err(e) if e.canonicalize() == Error::Interrupted => continue, + Err(e) => return Err(e), + } + } + if !buf.is_empty() { + Err(Error::UnexpectedEof) + } else { + Ok(()) + } +} + +/// Default [`Read::read_buf`] implementation. +pub fn default_read_buf(read: F, mut cursor: BorrowedCursor<'_>) -> Result<()> +where + F: FnOnce(&mut [u8]) -> Result, +{ + #[cfg(borrowedbuf_init)] + { + let n = read(cursor.ensure_init().init_mut())?; + cursor.advance(n); + } + #[cfg(not(borrowedbuf_init))] + { + // SAFETY: We do not uninitialize any part of the buffer. + let n = read(unsafe { cursor.as_mut().write_filled(0) })?; + assert!(n <= cursor.capacity()); + // SAFETY: We've initialized the entire buffer, and `read` can't make it uninitialized. + unsafe { + cursor.advance(n); + } + } + Ok(()) +} + +/// Default [`Read::read_buf_exact`] implementation. +pub fn default_read_buf_exact( + this: &mut R, + mut cursor: BorrowedCursor<'_>, +) -> Result<()> { + while cursor.capacity() > 0 { + let prev_written = cursor.written(); + match this.read_buf(cursor.reborrow()) { + Ok(()) => {} + Err(e) if e.canonicalize() == Error::Interrupted => continue, + Err(e) => return Err(e), + } + + if cursor.written() == prev_written { + return Err(Error::UnexpectedEof); + } + } + + Ok(()) +} + +/// Default [`Read::read_to_end`] implementation with optional size hint. +#[cfg(feature = "alloc")] +pub fn default_read_to_end( + r: &mut R, + buf: &mut Vec, + size_hint: Option, +) -> Result { + use core::io::BorrowedBuf; + + use crate::DEFAULT_BUF_SIZE; + + let start_len = buf.len(); + let start_cap = buf.capacity(); + // Optionally limit the maximum bytes read on each iteration. + // This adds an arbitrary fiddle factor to allow for more data than we expect. + let mut max_read_size = size_hint + .and_then(|s| { + s.checked_add(1024)? + .checked_next_multiple_of(DEFAULT_BUF_SIZE) + }) + .unwrap_or(DEFAULT_BUF_SIZE); + + const PROBE_SIZE: usize = 32; + + fn small_probe_read(r: &mut R, buf: &mut Vec) -> Result { + let mut probe = [0u8; PROBE_SIZE]; + + loop { + match r.read(&mut probe) { + Ok(n) => { + // there is no way to recover from allocation failure here + // because the data has already been read. + buf.extend_from_slice(&probe[..n]); + return Ok(n); + } + Err(e) if e.canonicalize() == Error::Interrupted => continue, + Err(e) => return Err(e), + } + } + } + + if (size_hint.is_none() || size_hint == Some(0)) && buf.capacity() - buf.len() < PROBE_SIZE { + let read = small_probe_read(r, buf)?; + + if read == 0 { + return Ok(0); + } + } + + #[cfg(borrowedbuf_init)] + let mut initialized = 0; // Extra initialized bytes from previous loop iteration + #[cfg(borrowedbuf_init)] + let mut consecutive_short_reads = 0; + + loop { + if buf.len() == buf.capacity() && buf.capacity() == start_cap { + // The buffer might be an exact fit. Let's read into a probe buffer + // and see if it returns `Ok(0)`. If so, we've avoided an + // unnecessary doubling of the capacity. But if not, append the + // probe buffer to the primary buffer and let its capacity grow. + let read = small_probe_read(r, buf)?; + + if read == 0 { + return Ok(buf.len() - start_len); + } + } + + if buf.len() == buf.capacity() { + // buf is full, need more space + buf.try_reserve(PROBE_SIZE).map_err(|_| Error::NoMemory)?; + } + + let mut spare = buf.spare_capacity_mut(); + let buf_len = spare.len().min(max_read_size); + spare = &mut spare[..buf_len]; + let mut read_buf: BorrowedBuf<'_> = spare.into(); + + #[cfg(borrowedbuf_init)] + // SAFETY: These bytes were initialized but not filled in the previous loop + unsafe { + read_buf.set_init(initialized); + } + + let mut cursor = read_buf.unfilled(); + let result = loop { + match r.read_buf(cursor.reborrow()) { + Err(e) if e.canonicalize() == Error::Interrupted => continue, + // Do not stop now in case of error: we might have received both data + // and an error + res => break res, + } + }; + + #[cfg(borrowedbuf_init)] + let unfilled_but_initialized = cursor.init_mut().len(); + let bytes_read = cursor.written(); + #[cfg(borrowedbuf_init)] + let was_fully_initialized = read_buf.init_len() == buf_len; + + // SAFETY: BorrowedBuf's invariants mean this much memory is initialized. + unsafe { + let new_len = bytes_read + buf.len(); + buf.set_len(new_len); + } + + // Now that all data is pushed to the vector, we can fail without data loss + result?; + + if bytes_read == 0 { + return Ok(buf.len() - start_len); + } + + #[cfg(borrowedbuf_init)] + if bytes_read < buf_len { + consecutive_short_reads += 1; + } else { + consecutive_short_reads = 0; + } + + #[cfg(borrowedbuf_init)] + { + // store how much was initialized but not filled + initialized = unfilled_but_initialized; + } + + // Use heuristics to determine the max read size if no initial size hint was provided + if size_hint.is_none() { + #[cfg(borrowedbuf_init)] + // The reader is returning short reads but it doesn't call ensure_init(). + // In that case we no longer need to restrict read sizes to avoid + // initialization costs. + // When reading from disk we usually don't get any short reads except at EOF. + // So we wait for at least 2 short reads before uncapping the read buffer; + // this helps with the Windows issue. + if !was_fully_initialized && consecutive_short_reads > 1 { + max_read_size = usize::MAX; + } + + // we have passed a larger buffer than previously and the + // reader still hasn't returned a short read + if buf_len >= max_read_size && bytes_read == buf_len { + max_read_size = max_read_size.saturating_mul(2); + } + } + } +} + +#[cfg(feature = "alloc")] +pub(crate) unsafe fn append_to_string(buf: &mut String, f: F) -> Result +where + F: FnOnce(&mut Vec) -> Result, +{ + struct Guard<'a> { + buf: &'a mut Vec, + len: usize, + } + + impl Drop for Guard<'_> { + fn drop(&mut self) { + unsafe { + self.buf.set_len(self.len); + } + } + } + + let mut g = Guard { + len: buf.len(), + buf: unsafe { buf.as_mut_vec() }, + }; + let ret = f(g.buf); + + // SAFETY: the caller promises to only append data to `buf` + let appended = unsafe { g.buf.get_unchecked(g.len..) }; + if str::from_utf8(appended).is_err() { + ret.and(Err(Error::IllegalBytes)) + } else { + g.len = g.buf.len(); + ret + } +} + +/// Default [`Read::read_to_string`] implementation with optional size hint. +#[cfg(feature = "alloc")] +pub fn default_read_to_string( + r: &mut R, + buf: &mut String, + size_hint: Option, +) -> Result { + // Note that we do *not* call `r.read_to_end()` here. We are passing + // `&mut Vec` (the raw contents of `buf`) into the `read_to_end` + // method to fill it up. An arbitrary implementation could overwrite the + // entire contents of the vector, not just append to it (which is what + // we are expecting). + // + // To prevent extraneously checking the UTF-8-ness of the entire buffer + // we pass it to our hardcoded `default_read_to_end` implementation which + // we know is guaranteed to only read data into the end of the buffer. + unsafe { append_to_string(buf, |b| default_read_to_end(r, b, size_hint)) } +} + +/// The `Read` trait allows for reading bytes from a source. +/// +/// See [`std::io::Read`] for more details. +pub trait Read { + /// Pull some bytes from this source into the specified buffer, returning + /// how many bytes were read. + fn read(&mut self, buf: &mut [u8]) -> Result; + + /// Read the exact number of bytes required to fill `buf`. + fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { + default_read_exact(self, buf) + } + + /// Pull some bytes from this source into the specified buffer. + /// + /// This method makes it possible to return both data and an error but it is advised against. + fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> Result<()> { + default_read_buf(|b| self.read(b), buf) + } + + /// Reads the exact number of bytes required to fill `cursor`. + /// + /// If this function returns an error, all bytes read will be appended to `cursor`. + fn read_buf_exact(&mut self, cursor: BorrowedCursor<'_>) -> Result<()> { + default_read_buf_exact(self, cursor) + } + + /// Read all bytes until EOF in this source, placing them into `buf`. + #[cfg(feature = "alloc")] + fn read_to_end(&mut self, buf: &mut Vec) -> Result { + default_read_to_end(self, buf, None) + } + + /// Read all bytes until EOF in this source, appending them to `buf`. + #[cfg(feature = "alloc")] + fn read_to_string(&mut self, buf: &mut String) -> Result { + default_read_to_string(self, buf, None) + } + + /// Creates a "by reference" adapter for this instance of `Read`. + /// + /// The returned `adapter` also implements Read and will simply borrow this + /// current reader. + fn by_ref(&mut self) -> &mut Self + where + Self: Sized, + { + self + } +} + +/// Reads all bytes from a [reader][Read] into a new [`String`]. +/// +/// This is a convenience function for [`Read::read_to_string`]. +/// +/// See [`std::io::read_to_string`] for more details. +#[cfg(feature = "alloc")] +pub fn read_to_string(mut reader: R) -> Result { + let mut buf = String::new(); + reader.read_to_string(&mut buf)?; + Ok(buf) +} + +/// A `BufRead` is a type of `Read`er which has an internal buffer, allowing it +/// to perform extra ways of reading. +/// +/// See [`std::io::BufRead`] for more details. +pub trait BufRead: Read { + /// Returns the contents of the internal buffer, filling it with more data, via `Read` methods, + /// if empty. + fn fill_buf(&mut self) -> Result<&[u8]>; + + /// Marks the given `amount` of additional bytes from the internal buffer as having been read. + /// Subsequent calls to `read` only return bytes that have not been marked as read. + fn consume(&mut self, amount: usize); + + /// Checks if there is any data left to be `read`. + fn has_data_left(&mut self) -> Result { + self.fill_buf().map(|b| !b.is_empty()) + } + + /// Skips all bytes until the delimiter `byte` or EOF is reached. + fn skip_until(&mut self, byte: u8) -> Result { + let mut read = 0; + loop { + let (done, used) = { + let available = self.fill_buf()?; + match memchr::memchr(byte, available) { + Some(i) => (true, i + 1), + None => (false, available.len()), + } + }; + self.consume(used); + read += used; + if done || used == 0 { + return Ok(read); + } + } + } + + /// Read all bytes into `buf` until the delimiter `byte` or EOF is reached. + #[cfg(feature = "alloc")] + fn read_until(&mut self, byte: u8, buf: &mut Vec) -> Result { + let mut read = 0; + loop { + let (done, used) = { + let available = self.fill_buf()?; + match memchr::memchr(byte, available) { + Some(i) => { + buf.extend_from_slice(&available[..=i]); + (true, i + 1) + } + None => { + buf.extend_from_slice(available); + (false, available.len()) + } + } + }; + self.consume(used); + read += used; + if done || used == 0 { + return Ok(read); + } + } + } + + /// Read all bytes until a newline (the `0xA` byte) is reached, and append + /// them to the provided `String` buffer. + #[cfg(feature = "alloc")] + fn read_line(&mut self, buf: &mut String) -> Result { + unsafe { super::append_to_string(buf, |b| self.read_until(b'\n', b)) } + } + + /// Returns an iterator over the contents of this reader split on the byte + /// `byte`. + #[cfg(feature = "alloc")] + fn split(self, byte: u8) -> Split + where + Self: Sized, + { + Split { + buf: self, + delim: byte, + } + } + + /// Returns an iterator over the lines of this reader. + #[cfg(feature = "alloc")] + fn lines(self) -> Lines + where + Self: Sized, + { + Lines { buf: self } + } +} + +/// An iterator over the contents of an instance of `BufRead` split on a +/// particular byte. +/// +/// This struct is generally created by calling [`split`] on a `BufRead`. +/// Please see the documentation of [`split`] for more details. +/// +/// [`split`]: BufRead::split +#[cfg(feature = "alloc")] +#[derive(Debug)] +pub struct Split { + buf: B, + delim: u8, +} + +#[cfg(feature = "alloc")] +impl Iterator for Split { + type Item = Result>; + + fn next(&mut self) -> Option>> { + let mut buf = Vec::new(); + match self.buf.read_until(self.delim, &mut buf) { + Ok(0) => None, + Ok(_n) => { + if buf[buf.len() - 1] == self.delim { + buf.pop(); + } + Some(Ok(buf)) + } + Err(e) => Some(Err(e)), + } + } +} + +/// An iterator over the lines of an instance of `BufRead`. +/// +/// This struct is generally created by calling [`lines`] on a `BufRead`. +/// Please see the documentation of [`lines`] for more details. +/// +/// [`lines`]: BufRead::lines +#[cfg(feature = "alloc")] +#[derive(Debug)] +pub struct Lines { + buf: B, +} + +#[cfg(feature = "alloc")] +impl Iterator for Lines { + type Item = Result; + + fn next(&mut self) -> Option> { + let mut buf = String::new(); + match self.buf.read_line(&mut buf) { + Ok(0) => None, + Ok(_n) => { + if buf.ends_with('\n') { + buf.pop(); + if buf.ends_with('\r') { + buf.pop(); + } + } + Some(Ok(buf)) + } + Err(e) => Some(Err(e)), + } + } +} diff --git a/src/seek/impls.rs b/src/seek/impls.rs new file mode 100644 index 0000000..3bfb3e2 --- /dev/null +++ b/src/seek/impls.rs @@ -0,0 +1,62 @@ +#[cfg(feature = "alloc")] +use alloc::boxed::Box; + +use crate::{Result, Seek, SeekFrom}; + +// ============================================================================= +// Forwarding implementations + +impl Seek for &mut S { + #[inline] + fn seek(&mut self, pos: SeekFrom) -> Result { + (**self).seek(pos) + } + + #[inline] + fn rewind(&mut self) -> Result<()> { + (**self).rewind() + } + + #[inline] + fn stream_len(&mut self) -> Result { + (**self).stream_len() + } + + #[inline] + fn stream_position(&mut self) -> Result { + (**self).stream_position() + } + + #[inline] + fn seek_relative(&mut self, offset: i64) -> Result<()> { + (**self).seek_relative(offset) + } +} + +#[cfg(feature = "alloc")] +impl Seek for Box { + #[inline] + fn seek(&mut self, pos: SeekFrom) -> Result { + (**self).seek(pos) + } + + #[inline] + fn rewind(&mut self) -> Result<()> { + (**self).rewind() + } + + #[inline] + fn stream_len(&mut self) -> Result { + (**self).stream_len() + } + + #[inline] + fn stream_position(&mut self) -> Result { + (**self).stream_position() + } + + #[inline] + fn seek_relative(&mut self, offset: i64) -> Result<()> { + (**self).seek_relative(offset) + } +} diff --git a/src/seek/mod.rs b/src/seek/mod.rs new file mode 100644 index 0000000..f76b126 --- /dev/null +++ b/src/seek/mod.rs @@ -0,0 +1,82 @@ +use crate::Result; + +mod impls; + +/// Enumeration of possible methods to seek within an I/O object. +/// +/// It is used by the [`Seek`] trait. +#[derive(Copy, PartialEq, Eq, Clone, Debug)] +pub enum SeekFrom { + /// Sets the offset to the provided number of bytes. + Start(u64), + + /// Sets the offset to the size of this object plus the specified number of + /// bytes. + /// + /// It is possible to seek beyond the end of an object, but it's an error to + /// seek before byte 0. + End(i64), + + /// Sets the offset to the current position plus the specified number of + /// bytes. + /// + /// It is possible to seek beyond the end of an object, but it's an error to + /// seek before byte 0. + Current(i64), +} + +/// Default [`Seek::stream_len`] implementation. +pub fn default_stream_len(this: &mut T) -> Result { + let old_pos = this.stream_position()?; + let len = this.seek(SeekFrom::End(0))?; + + // Avoid seeking a third time when we were already at the end of the + // stream. The branch is usually way cheaper than a seek operation. + if old_pos != len { + this.seek(SeekFrom::Start(old_pos))?; + } + + Ok(len) +} + +/// The `Seek` trait provides a cursor which can be moved within a stream of +/// bytes. +/// +/// See [`std::io::Seek`] for more details. +pub trait Seek { + /// Seek to an offset, in bytes, in a stream. + /// + /// A seek beyond the end of a stream is allowed, but behavior is defined + /// by the implementation. + /// + /// If the seek operation completed successfully, + /// this method returns the new position from the start of the stream. + /// That position can be used later with [`SeekFrom::Start`]. + fn seek(&mut self, pos: SeekFrom) -> Result; + + /// Rewind to the beginning of a stream. + /// + /// This is a convenience method, equivalent to `seek(SeekFrom::Start(0))`. + fn rewind(&mut self) -> Result<()> { + self.seek(SeekFrom::Start(0))?; + Ok(()) + } + + /// Returns the current seek position from the start of the stream. + /// + /// This is equivalent to `self.seek(SeekFrom::Current(0))`. + fn stream_position(&mut self) -> Result { + self.seek(SeekFrom::Current(0)) + } + + /// Returns the length of this stream (in bytes). + fn stream_len(&mut self) -> Result { + default_stream_len(self) + } + + /// Seeks relative to the current position. + fn seek_relative(&mut self, offset: i64) -> Result<()> { + self.seek(SeekFrom::Current(offset))?; + Ok(()) + } +} diff --git a/src/write/impls.rs b/src/write/impls.rs new file mode 100644 index 0000000..5da75f5 --- /dev/null +++ b/src/write/impls.rs @@ -0,0 +1,144 @@ +#[cfg(feature = "alloc")] +use alloc::{boxed::Box, collections::VecDeque, vec::Vec}; +use core::{cmp, fmt, io::BorrowedCursor, mem}; + +use crate::{Error, Result, Write}; + +// ============================================================================= +// Forwarding implementations + +impl Write for &mut W { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + (**self).write(buf) + } + + #[inline] + fn flush(&mut self) -> Result<()> { + (**self).flush() + } + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> Result<()> { + (**self).write_all(buf) + } + + #[inline] + fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> Result<()> { + (**self).write_fmt(fmt) + } +} + +#[cfg(feature = "alloc")] +impl Write for Box { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + (**self).write(buf) + } + + #[inline] + fn flush(&mut self) -> Result<()> { + (**self).flush() + } + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> Result<()> { + (**self).write_all(buf) + } + + #[inline] + fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> Result<()> { + (**self).write_fmt(fmt) + } +} + +// ============================================================================= +// In-memory buffer implementations + +impl Write for &mut [u8] { + #[inline] + fn write(&mut self, data: &[u8]) -> Result { + let amt = cmp::min(data.len(), self.len()); + let (a, b) = mem::take(self).split_at_mut(amt); + a.copy_from_slice(&data[..amt]); + *self = b; + Ok(amt) + } + + #[inline] + fn flush(&mut self) -> Result<()> { + Ok(()) + } + + #[inline] + fn write_all(&mut self, data: &[u8]) -> Result<()> { + if self.write(data)? < data.len() { + Err(Error::WriteZero) + } else { + Ok(()) + } + } +} + +#[cfg(feature = "alloc")] +impl Write for Vec { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + self.extend_from_slice(buf); + Ok(buf.len()) + } + + #[inline] + fn flush(&mut self) -> Result<()> { + Ok(()) + } + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> Result<()> { + self.extend_from_slice(buf); + Ok(()) + } +} + +#[cfg(feature = "alloc")] +impl Write for VecDeque { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + self.extend(buf); + Ok(buf.len()) + } + + #[inline] + fn flush(&mut self) -> Result<()> { + Ok(()) + } + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> Result<()> { + self.extend(buf); + Ok(()) + } +} + +impl Write for BorrowedCursor<'_> { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + let amt = cmp::min(buf.len(), self.capacity()); + self.append(&buf[..amt]); + Ok(amt) + } + + #[inline] + fn flush(&mut self) -> Result<()> { + Ok(()) + } + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> Result<()> { + if self.write(buf)? < buf.len() { + Err(Error::WriteZero) + } else { + Ok(()) + } + } +} diff --git a/src/write/mod.rs b/src/write/mod.rs new file mode 100644 index 0000000..1310208 --- /dev/null +++ b/src/write/mod.rs @@ -0,0 +1,94 @@ +use core::fmt; + +use crate::{Error, Result}; + +mod impls; + +pub(crate) fn default_write_fmt( + this: &mut W, + args: fmt::Arguments<'_>, +) -> Result<()> { + // Create a shim which translates a `Write` to a `fmt::Write` and saves off + // I/O errors, instead of discarding them. + struct Adapter<'a, T: ?Sized + 'a> { + inner: &'a mut T, + error: Result<()>, + } + + impl fmt::Write for Adapter<'_, T> { + fn write_str(&mut self, s: &str) -> fmt::Result { + match self.inner.write_all(s.as_bytes()) { + Ok(()) => Ok(()), + Err(e) => { + self.error = Err(e); + Err(fmt::Error) + } + } + } + } + + let mut output = Adapter { + inner: this, + error: Ok(()), + }; + match fmt::write(&mut output, args) { + Ok(()) => Ok(()), + Err(..) => { + // Check whether the error came from the underlying `Write`. + if output.error.is_err() { + output.error + } else { + // This shouldn't happen: the underlying stream did not error, + // but somehow the formatter still errored? + panic!( + "a formatting trait implementation returned an error when the underlying \ + stream did not" + ); + } + } + } +} + +/// A trait for objects which are byte-oriented sinks. +/// +/// See [`std::io::Write`] for more details. +pub trait Write { + /// Write a buffer into this writer, returning how many bytes were written. + fn write(&mut self, buf: &[u8]) -> Result; + + /// Flush this output stream, ensuring that all intermediately buffered + /// contents reach their destination. + fn flush(&mut self) -> Result<()>; + + /// Attempts to write an entire buffer into this writer. + fn write_all(&mut self, mut buf: &[u8]) -> Result<()> { + while !buf.is_empty() { + match self.write(buf) { + Ok(0) => return Err(Error::WriteZero), + Ok(n) => buf = &buf[n..], + Err(e) if e.canonicalize() == Error::Interrupted => {} + Err(e) => return Err(e), + } + } + Ok(()) + } + + /// Writes a formatted string into this writer, returning any error + /// encountered. + fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> Result<()> { + // NOTE: std's implementation uses an internal method + // `as_statically_known_str` to optimize. + default_write_fmt(self, args) + } + + /// Creates a "by reference" adapter for this instance of `Write`. + /// + /// The returned adapter also implements `Write` and will simply borrow this + /// current writer. + fn by_ref(&mut self) -> &mut Self + where + Self: Sized, + { + self + } +}