diff --git a/io/src/lib.rs b/io/src/lib.rs index e4a4a79e..85b2afe9 100644 --- a/io/src/lib.rs +++ b/io/src/lib.rs @@ -36,10 +36,63 @@ pub mod io { compile_error!("At least one of std or core2 must be enabled"); #[cfg(feature = "std")] - pub use std::io::{Read, Write, sink, Cursor, Take, Error, ErrorKind, Result}; + pub use std::io::{Read, sink, Cursor, Take, Error, ErrorKind, Result}; #[cfg(not(feature = "std"))] - pub use core2::io::{Read, Write, Cursor, Take, Error, ErrorKind, Result}; + pub use core2::io::{Read, Cursor, Take, Error, ErrorKind, Result}; + /// A generic trait describing an output stream. See [`std::io::Write`] for more info. + pub trait Write { + fn write(&mut self, buf: &[u8]) -> Result; + fn flush(&mut self) -> Result<()>; + #[inline] + fn write_all(&mut self, mut buf: &[u8]) -> Result<()> { + while !buf.is_empty() { + match self.write(buf) { + Ok(0) => return Err(Error::new(ErrorKind::UnexpectedEof, "")), + Ok(len) => buf = &buf[len..], + Err(e) if e.kind() == ErrorKind::Interrupted => {} + Err(e) => return Err(e), + } + } + Ok(()) + } + } + + #[cfg(feature = "std")] + impl Write for W { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + ::write(self, buf) + } + #[inline] + fn flush(&mut self) -> Result<()> { + ::flush(self) + } + } + + #[cfg(all(feature = "alloc", not(feature = "std")))] + impl Write for alloc::vec::Vec { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + self.extend_from_slice(buf); + Ok(buf.len()) + } + #[inline] + fn flush(&mut self) -> Result<()> { Ok(()) } + } + + #[cfg(not(feature = "std"))] + impl<'a> Write for &'a mut [u8] { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + let cnt = core::cmp::min(self.len(), buf.len()); + self[..cnt].copy_from_slice(&buf[..cnt]); + *self = &mut core::mem::take(self)[cnt..]; + Ok(cnt) + } + #[inline] + fn flush(&mut self) -> Result<()> { Ok(()) } + } }