feat: add BufWriter and extend buffered I/O traits and impls

Cargo.toml

@@ -17,3 +17,4 @@ default = ["alloc"]
 
 [dependencies]
 axerrno = "0.2"
+bytemuck = { version = "1.23", features = ["derive"] }

src/buffered/mod.rs

@@ -1,3 +1,6 @@
 mod reader;
+mod writer;
 
-pub use self::reader::BufReader;
+pub use self::{reader::BufReader, writer::BufWriter};
+
+const DEFAULT_BUF_SIZE: usize = 1024;

src/buffered/reader.rs

@@ -1,12 +1,11 @@
 use core::mem::MaybeUninit;
 
 use crate::{BufRead, Read, Result};
+use super::DEFAULT_BUF_SIZE;
 
 #[cfg(feature = "alloc")]
 use alloc::{string::String, vec::Vec};
 
-const DEFAULT_BUF_SIZE: usize = 1024;
-
 /// The `BufReader<R>` struct adds buffering to any reader.
 pub struct BufReader<R> {
     inner: R,

src/buffered/writer.rs

@@ -0,0 +1,113 @@
+use core::{
+    mem::{ManuallyDrop, MaybeUninit},
+    ptr,
+};
+
+use super::DEFAULT_BUF_SIZE;
+use crate::{BufWrite, Result, Write};
+
+/// The `BufWriter<W>` struct adds buffering to any writer.
+pub struct BufWriter<W: Write> {
+    inner: W,
+    pos: usize,
+    buf: [MaybeUninit<u8>; DEFAULT_BUF_SIZE],
+}
+
+impl<W: Write> BufWriter<W> {
+    /// Creates a new `BufWriter<W>` with a default buffer capacity (1 KB).
+    pub const fn new(inner: W) -> BufWriter<W> {
+        Self {
+            inner,
+            pos: 0,
+            buf: [const { MaybeUninit::uninit() }; DEFAULT_BUF_SIZE],
+        }
+    }
+
+    /// Gets a reference to the underlying writer.
+    pub const fn get_ref(&self) -> &W {
+        &self.inner
+    }
+
+    /// Gets a mutable reference to the underlying writer.
+    pub fn get_mut(&mut self) -> &mut W {
+        &mut self.inner
+    }
+
+    /// Returns a reference to the internally buffered data.
+    pub fn buffer(&self) -> &[u8] {
+        unsafe { self.buf[..self.pos].assume_init_ref() }
+    }
+
+    /// Returns the number of bytes the internal buffer can hold at once.
+    pub const fn capacity(&self) -> usize {
+        DEFAULT_BUF_SIZE
+    }
+
+    /// Returns the remaining spare capacity in the internal buffer.
+    pub const fn spare_capacity(&self) -> usize {
+        self.capacity() - self.pos
+    }
+
+    /// Unwraps this `BufWriter<W>`, returning the underlying writer.
+    ///
+    /// Any buffered data will be flushed before returning.
+    pub fn into_inner(self) -> Result<W> {
+        // Prevent Drop from running while we manually extract the inner writer.
+        let mut this = ManuallyDrop::new(self);
+        this.flush_buf()?;
+        Ok(unsafe { ptr::read(&this.inner) })
+    }
+}
+
+impl<W: Write> Write for BufWriter<W> {
+    /// Writes a buffer into this writer, returning how many bytes were written.
+    fn write(&mut self, buf: &[u8]) -> Result<usize> {
+        if self.spare_capacity() < buf.len() {
+            self.flush_buf()?;
+        }
+        let written = buf.len().min(self.spare_capacity());
+        unsafe {
+            self.buf[self.pos..self.pos + written]
+                .assume_init_mut()
+                .copy_from_slice(&buf[..written]);
+        }
+        self.pos += written;
+        Ok(written)
+    }
+
+    /// Flushes this writer, ensuring that all intermediately buffered contents reach their destination.
+    fn flush(&mut self) -> Result<()> {
+        self.flush_buf()?;
+        self.inner.flush()
+    }
+}
+
+impl<W: Write> BufWrite for BufWriter<W> {
+    /// Flushes the internal buffer to the underlying writer.
+    fn flush_buf(&mut self) -> Result<()> {
+        if self.pos > 0 {
+            self.inner
+                .write_all(unsafe { self.buf[..self.pos].assume_init_ref() })?;
+            self.pos = 0;
+        }
+        Ok(())
+    }
+
+    /// Skips a number of bytes in the internal buffer, flushing if necessary.
+    fn skip_some(&mut self, len: usize) -> Result<()> {
+        let mut sparce = self.spare_capacity();
+        if sparce < len {
+            self.flush_buf()?;
+            sparce = self.spare_capacity();
+        }
+        self.pos += len.min(sparce);
+        Ok(())
+    }
+}
+
+/// Drops the `BufWriter`, flushing the internal buffer.
+impl<W: Write> Drop for BufWriter<W> {
+    fn drop(&mut self) {
+        let _ = self.flush_buf();
+    }
+}

src/impls.rs

@@ -1,10 +1,12 @@
+#[cfg(feature = "alloc")]
+use alloc::{string::String, vec::Vec};
 use core::{cmp, mem};
 
 use axerrno::ax_bail;
 
 use crate::{
+    BufRead, BufWrite, Read, Result, Seek, SeekFrom, Write,
     buf::{Buf, BufMut},
-    Read, Result, Write,
 };
 
 impl Read for &[u8] {
@@ -105,3 +107,101 @@ impl BufMut for &mut [u8] {
         Ok(filled)
     }
 }
+
+impl<R: Read + ?Sized> Read for &mut R {
+    #[inline]
+    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
+        (**self).read(buf)
+    }
+
+    #[inline]
+    fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
+        (**self).read_exact(buf)
+    }
+
+    #[cfg(feature = "alloc")]
+    #[inline]
+    fn read_to_end(&mut self, buf: &mut alloc::vec::Vec<u8>) -> Result<usize> {
+        (**self).read_to_end(buf)
+    }
+
+    #[cfg(feature = "alloc")]
+    #[inline]
+    fn read_to_string(&mut self, buf: &mut alloc::string::String) -> Result<usize> {
+        (**self).read_to_string(buf)
+    }
+}
+
+impl<W: Write + ?Sized> Write for &mut W {
+    #[inline]
+    fn write(&mut self, buf: &[u8]) -> Result<usize> {
+        (**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: core::fmt::Arguments<'_>) -> Result<()> {
+        (**self).write_fmt(fmt)
+    }
+}
+
+impl<S: Seek + ?Sized> Seek for &mut S {
+    fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
+        (**self).seek(pos)
+    }
+
+    fn rewind(&mut self) -> Result<()> {
+        (**self).rewind()
+    }
+
+    fn stream_position(&mut self) -> Result<u64> {
+        (**self).stream_position()
+    }
+}
+
+impl<B: BufRead + ?Sized> BufRead for &mut B {
+    fn fill_buf(&mut self) -> Result<&[u8]> {
+        (**self).fill_buf()
+    }
+
+    fn consume(&mut self, amt: usize) {
+        (**self).consume(amt)
+    }
+
+    fn has_data_left(&mut self) -> Result<bool> {
+        (**self).has_data_left()
+    }
+
+    #[cfg(feature = "alloc")]
+    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> Result<usize> {
+        (**self).read_until(byte, buf)
+    }
+
+    #[cfg(feature = "alloc")]
+    fn read_line(&mut self, buf: &mut String) -> Result<usize> {
+        (**self).read_line(buf)
+    }
+}
+
+impl<B: BufWrite + ?Sized> BufWrite for &mut B {
+    fn flush_buf(&mut self) -> Result<()> {
+        (**self).flush_buf()
+    }
+
+    fn skip_some(&mut self, len: usize) -> Result<()> {
+        (**self).skip_some(len)
+    }
+
+    fn write_val<T: bytemuck::Pod>(&mut self, val: &T) -> Result<()> {
+        (**self).write_val(val)
+    }
+}

src/lib.rs

@@ -18,13 +18,13 @@ pub mod prelude;
 
 pub use self::{
     buf::{Buf, BufMut},
-    buffered::BufReader,
+    buffered::{BufReader, BufWriter},
     error::{Error, Result},
 };
 
 #[cfg(feature = "alloc")]
-use alloc::{string::String, vec::Vec};
-use axerrno::ax_bail;
+use alloc::{ffi::CString, string::String, vec::Vec};
+use axerrno::{ax_bail, AxError};
 
 const DEFAULT_BUF_SIZE: usize = 1024;
 
@@ -349,6 +349,35 @@ pub trait BufRead: Read {
         }
     }
 
+    /// Attempts to decode `T` directly from the currently buffered bytes.
+    ///
+    /// Returns `Ok(None)` if fewer than `size_of::<T>()` bytes are available.
+    #[cfg(feature = "alloc")]
+    fn read_val<T: bytemuck::Pod>(&mut self) -> Result<Option<T>> {
+        let size = core::mem::size_of::<T>();
+        let buf = self.fill_buf()?;
+        if buf.len() < size {
+            return Ok(None);
+        }
+        let val = bytemuck::try_pod_read_unaligned::<T>(&buf[..size])
+            .map_err(|_| AxError::InvalidData)?;
+        self.consume(size);
+        Ok(Some(val))
+    }
+
+    /// Reads a C-style string, stopping at the first NUL or EOF.
+    #[cfg(feature = "alloc")]
+    fn read_cstring(&mut self) -> Result<(CString, usize)> {
+        let mut buf = Vec::new();
+        let n = self.read_until(0, &mut buf)?;
+        let cstring = if buf.ends_with(&[0]) {
+            CString::from_vec_with_nul(buf).map_err(|_| AxError::InvalidData)?
+        } else {
+            CString::new(buf).map_err(|_| AxError::InvalidData)?
+        };
+        Ok((cstring, n))
+    }
+
     /// Read all bytes until a newline (the `0xA` byte) is reached, and append
     /// them to the provided `String` buffer.
     #[cfg(feature = "alloc")]
@@ -372,6 +401,21 @@ where
     }
 }
 
+// A `BufWrite` is a type of `Write`r which has an internal buffer, allowing it
+// to perform extra ways of writing.
+pub trait BufWrite: Write {
+    /// Flushes the internal buffer to the underlying writer.
+    fn flush_buf(&mut self) -> Result<()>;
+
+    /// Skips writing `len` bytes by advancing the internal buffer.
+    fn skip_some(&mut self, len: usize) -> Result<()>;
+
+    /// Writes the bytes of a plain-old-data value into this writer.
+    fn write_val<T: bytemuck::Pod>(&mut self, val: &T) -> Result<()> {
+        self.write_all(bytemuck::bytes_of(val))
+    }
+}
+
 // FIXME: This is reserved for smooth migration of ArceOS.
 /// I/O poll results.
 #[derive(Debug, Default, Clone, Copy)]

src/prelude.rs

@@ -8,4 +8,4 @@
 //! use std::io::prelude::*;
 //! ```
 
-pub use super::{BufRead, Read, Seek, Write};
+pub use super::{BufRead, BufWrite, Read, Seek, Write};