Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Override default Write methods for cursor-like types #137107

Merged
merged 2 commits into from
Mar 7, 2025
Merged

Override default Write methods for cursor-like types #137107

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
67cc82a
Inline VecDeque<u8> and BorrowedCursor methods
thaliaarchi Feb 15, 2025
41bdd2b
Override default Write methods for cursor-like types
thaliaarchi Feb 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 87 additions & 13 deletions library/std/src/io/cursor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -439,6 +439,27 @@ fn slice_write_vectored(
Ok(nwritten)
}

#[inline]
fn slice_write_all(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<()> {
let n = slice_write(pos_mut, slice, buf)?;
if n < buf.len() { Err(io::Error::WRITE_ALL_EOF) } else { Ok(()) }
}

#[inline]
fn slice_write_all_vectored(
pos_mut: &mut u64,
slice: &mut [u8],
bufs: &[IoSlice<'_>],
) -> io::Result<()> {
for buf in bufs {
let n = slice_write(pos_mut, slice, buf)?;
if n < buf.len() {
return Err(io::Error::WRITE_ALL_EOF);
}
}
Ok(())
}

/// Reserves the required space, and pads the vec with 0s if necessary.
fn reserve_and_pad<A: Allocator>(
pos_mut: &mut u64,
Expand Down Expand Up @@ -481,9 +502,12 @@ fn reserve_and_pad<A: Allocator>(
Ok(pos)
}

/// Writes the slice to the vec without allocating
/// # Safety: vec must have buf.len() spare capacity
unsafe fn vec_write_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8]) -> usize
/// Writes the slice to the vec without allocating.
///
/// # Safety
///
/// `vec` must have `buf.len()` spare capacity.
unsafe fn vec_write_all_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8]) -> usize
where
A: Allocator,
{
Expand All @@ -492,7 +516,7 @@ where
pos + buf.len()
}

/// Resizing write implementation for [`Cursor`]
/// Resizing `write_all` implementation for [`Cursor`].
///
/// Cursor is allowed to have a pre-allocated and initialised
/// vector body, but with a position of 0. This means the [`Write`]
Expand All @@ -501,7 +525,7 @@ where
/// This also allows for the vec body to be empty, but with a position of N.
/// This means that [`Write`] will pad the vec with 0 initially,
/// before writing anything from that point
fn vec_write<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize>
fn vec_write_all<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize>
where
A: Allocator,
{
Expand All @@ -512,7 +536,7 @@ where
// Safety: we have ensured that the capacity is available
// and that all bytes get written up to pos
unsafe {
pos = vec_write_unchecked(pos, vec, buf);
pos = vec_write_all_unchecked(pos, vec, buf);
if pos > vec.len() {
vec.set_len(pos);
}
Expand All @@ -523,7 +547,7 @@ where
Ok(buf_len)
}

/// Resizing write_vectored implementation for [`Cursor`]
/// Resizing `write_all_vectored` implementation for [`Cursor`].
///
/// Cursor is allowed to have a pre-allocated and initialised
/// vector body, but with a position of 0. This means the [`Write`]
Expand All @@ -532,7 +556,7 @@ where
/// This also allows for the vec body to be empty, but with a position of N.
/// This means that [`Write`] will pad the vec with 0 initially,
/// before writing anything from that point
fn vec_write_vectored<A>(
fn vec_write_all_vectored<A>(
pos_mut: &mut u64,
vec: &mut Vec<u8, A>,
bufs: &[IoSlice<'_>],
Expand All @@ -550,7 +574,7 @@ where
// and that all bytes get written up to the last pos
unsafe {
for buf in bufs {
pos = vec_write_unchecked(pos, vec, buf);
pos = vec_write_all_unchecked(pos, vec, buf);
}
if pos > vec.len() {
vec.set_len(pos);
Expand Down Expand Up @@ -579,6 +603,16 @@ impl Write for Cursor<&mut [u8]> {
true
}

#[inline]
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
slice_write_all(&mut self.pos, self.inner, buf)
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
slice_write_all_vectored(&mut self.pos, self.inner, bufs)
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand All @@ -591,18 +625,28 @@ where
A: Allocator,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
vec_write(&mut self.pos, self.inner, buf)
vec_write_all(&mut self.pos, self.inner, buf)
}

fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
vec_write_vectored(&mut self.pos, self.inner, bufs)
vec_write_all_vectored(&mut self.pos, self.inner, bufs)
}

#[inline]
fn is_write_vectored(&self) -> bool {
true
}

fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
vec_write_all(&mut self.pos, self.inner, buf)?;
Ok(())
}

fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
vec_write_all_vectored(&mut self.pos, self.inner, bufs)?;
Ok(())
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand All @@ -615,18 +659,28 @@ where
A: Allocator,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
vec_write(&mut self.pos, &mut self.inner, buf)
vec_write_all(&mut self.pos, &mut self.inner, buf)
}

fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
vec_write_vectored(&mut self.pos, &mut self.inner, bufs)
vec_write_all_vectored(&mut self.pos, &mut self.inner, bufs)
}

#[inline]
fn is_write_vectored(&self) -> bool {
true
}

fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
vec_write_all(&mut self.pos, &mut self.inner, buf)?;
Ok(())
}

fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
vec_write_all_vectored(&mut self.pos, &mut self.inner, bufs)?;
Ok(())
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand All @@ -653,6 +707,16 @@ where
true
}

#[inline]
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
slice_write_all(&mut self.pos, &mut self.inner, buf)
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
slice_write_all_vectored(&mut self.pos, &mut self.inner, bufs)
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand All @@ -676,6 +740,16 @@ impl<const N: usize> Write for Cursor<[u8; N]> {
true
}

#[inline]
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
slice_write_all(&mut self.pos, &mut self.inner, buf)
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
slice_write_all_vectored(&mut self.pos, &mut self.inner, bufs)
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand Down
60 changes: 59 additions & 1 deletion library/std/src/io/impls.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -455,7 +455,17 @@ impl Write for &mut [u8] {

#[inline]
fn write_all(&mut self, data: &[u8]) -> io::Result<()> {
if self.write(data)? == data.len() { Ok(()) } else { Err(io::Error::WRITE_ALL_EOF) }
if self.write(data)? < data.len() { Err(io::Error::WRITE_ALL_EOF) } else { Ok(()) }
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
for buf in bufs {
if self.write(buf)? < buf.len() {
return Err(io::Error::WRITE_ALL_EOF);
}
}
Ok(())
}

#[inline]
Expand Down Expand Up @@ -495,6 +505,12 @@ impl<A: Allocator> Write for Vec<u8, A> {
Ok(())
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
self.write_vectored(bufs)?;
Ok(())
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand All @@ -515,6 +531,7 @@ impl<A: Allocator> Read for VecDeque<u8, A> {
Ok(n)
}

#[inline]
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
let (front, back) = self.as_slices();

Expand Down Expand Up @@ -547,6 +564,7 @@ impl<A: Allocator> Read for VecDeque<u8, A> {
Ok(())
}

#[inline]
fn read_buf_exact(&mut self, mut cursor: BorrowedCursor<'_>) -> io::Result<()> {
let len = cursor.capacity();
let (front, back) = self.as_slices();
Expand Down Expand Up @@ -638,6 +656,12 @@ impl<A: Allocator> Write for VecDeque<u8, A> {
Ok(())
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
self.write_vectored(bufs)?;
Ok(())
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand All @@ -646,12 +670,46 @@ impl<A: Allocator> Write for VecDeque<u8, A> {

#[unstable(feature = "read_buf", issue = "78485")]
impl<'a> io::Write for core::io::BorrowedCursor<'a> {
#[inline]
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let amt = cmp::min(buf.len(), self.capacity());
self.append(&buf[..amt]);
Ok(amt)
}

#[inline]
fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
let mut nwritten = 0;
for buf in bufs {
let n = self.write(buf)?;
nwritten += n;
if n < buf.len() {
break;
}
}
Ok(nwritten)
}

#[inline]
fn is_write_vectored(&self) -> bool {
true
}

#[inline]
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
if self.write(buf)? < buf.len() { Err(io::Error::WRITE_ALL_EOF) } else { Ok(()) }
}

#[inline]
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
for buf in bufs {
if self.write(buf)? < buf.len() {
return Err(io::Error::WRITE_ALL_EOF);
}
}
Ok(())
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
Ok(())
Expand Down
Loading