From 74e0578bf95c9c45b10a466f3feff99854be0795 Mon Sep 17 00:00:00 2001 From: Sam Rijs Date: Mon, 12 Jun 2023 01:22:39 +0200 Subject: [PATCH 1/8] add specialized `Buf::chunks_vectored` for `Take` --- src/buf/take.rs | 26 +++++++++++++++++++++++ tests/test_take.rs | 52 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/src/buf/take.rs b/src/buf/take.rs index a16a434ee..253b3ce2d 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -152,4 +152,30 @@ impl Buf for Take { self.limit -= len; r } + + #[cfg(feature = "std")] + fn chunks_vectored<'a>(&'a self, dst: &mut [std::io::IoSlice<'a>]) -> usize { + let cnt = self.inner.chunks_vectored(dst); + let mut len = 0; + for (n, io) in dst[0..cnt].iter_mut().enumerate() { + let max = self.limit - len; + if max == 0 { + return n; + } + if io.len() > max { + // In this case, `IoSlice` is longer than our max, so we need to truncate it to the max. + // + // We need to work around the fact here that even though `IoSlice<'a>` has the correct + // lifetime, its `Deref` impl strips it. So we need to reassamble the slice to add the + // correct lifetime that allows us to call `IoSlice::<'a>::new` with it. + // + // TODO: remove `unsafe` as soon as `IoSlice::as_bytes` is available (rust-lang/rust#111277) + let buf = unsafe { std::slice::from_raw_parts::<'a, u8>(io.as_ptr(), max) }; + *io = std::io::IoSlice::new(buf); + return n + 1; + } + len += io.len(); + } + cnt + } } diff --git a/tests/test_take.rs b/tests/test_take.rs index 51df91d14..0c0159be1 100644 --- a/tests/test_take.rs +++ b/tests/test_take.rs @@ -30,3 +30,55 @@ fn take_copy_to_bytes_panics() { let abcd = Bytes::copy_from_slice(b"abcd"); abcd.take(2).copy_to_bytes(3); } + +#[cfg(feature = "std")] +#[test] +fn take_chunks_vectored() { + fn chain() -> impl Buf { + Bytes::from([1, 2, 3].to_vec()).chain(Bytes::from([4, 5, 6].to_vec())) + } + + { + let mut dst = [std::io::IoSlice::new(&[]); 2]; + let take = chain().take(0); + assert_eq!(take.chunks_vectored(&mut dst), 0); + } + + { + let mut dst = [std::io::IoSlice::new(&[]); 2]; + let take = chain().take(1); + assert_eq!(take.chunks_vectored(&mut dst), 1); + assert_eq!(&*dst[0], &[1]); + } + + { + let mut dst = [std::io::IoSlice::new(&[]); 2]; + let take = chain().take(3); + assert_eq!(take.chunks_vectored(&mut dst), 1); + assert_eq!(&*dst[0], &[1, 2, 3]); + } + + { + let mut dst = [std::io::IoSlice::new(&[]); 2]; + let take = chain().take(4); + assert_eq!(take.chunks_vectored(&mut dst), 2); + assert_eq!(&*dst[0], &[1, 2, 3]); + assert_eq!(&*dst[1], &[4]); + } + + { + let mut dst = [std::io::IoSlice::new(&[]); 2]; + let take = chain().take(6); + assert_eq!(take.chunks_vectored(&mut dst), 2); + assert_eq!(&*dst[0], &[1, 2, 3]); + assert_eq!(&*dst[1], &[4, 5, 6]); + } + + { + let mut dst = [std::io::IoSlice::new(&[]); 2]; + let take = chain().take(7); + assert_eq!(take.chunks_vectored(&mut dst), 2); + assert_eq!(&*dst[0], &[1, 2, 3]); + assert_eq!(&*dst[1], &[4, 5, 6]); + } +} From 598f0e5fe5362d5a018f99ce2267e54912f30f0e Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Fri, 10 Jan 2025 15:05:15 +0100 Subject: [PATCH 2/8] tweaks --- src/buf/take.rs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/buf/take.rs b/src/buf/take.rs index 253b3ce2d..0eff707b7 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -157,19 +157,13 @@ impl Buf for Take { fn chunks_vectored<'a>(&'a self, dst: &mut [std::io::IoSlice<'a>]) -> usize { let cnt = self.inner.chunks_vectored(dst); let mut len = 0; - for (n, io) in dst[0..cnt].iter_mut().enumerate() { + for (n, io) in dst[..cnt].iter_mut().enumerate() { let max = self.limit - len; - if max == 0 { - return n; - } - if io.len() > max { + if io.len() >= max { // In this case, `IoSlice` is longer than our max, so we need to truncate it to the max. // - // We need to work around the fact here that even though `IoSlice<'a>` has the correct - // lifetime, its `Deref` impl strips it. So we need to reassamble the slice to add the - // correct lifetime that allows us to call `IoSlice::<'a>::new` with it. - // // TODO: remove `unsafe` as soon as `IoSlice::as_bytes` is available (rust-lang/rust#111277) + // SAFETY: This reimplements the safe function `IoSlice::as_bytes`. let buf = unsafe { std::slice::from_raw_parts::<'a, u8>(io.as_ptr(), max) }; *io = std::io::IoSlice::new(buf); return n + 1; From 118f67930157c45ef8a79993abdedf7cc22cc396 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Fri, 10 Jan 2025 15:07:35 +0100 Subject: [PATCH 3/8] Update test now that chunks_vectored is fixed --- tests/test_buf.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_buf.rs b/tests/test_buf.rs index 5a5ac7e80..099016e24 100644 --- a/tests/test_buf.rs +++ b/tests/test_buf.rs @@ -404,7 +404,7 @@ mod chain_limited_slices { Buf::take(Buf::chain(Buf::chain(a, b), Buf::chain(c, d)), buf.len()) } - buf_tests!(make_input, /* `Limit` does not forward `chucks_vectored */ false); + buf_tests!(make_input, true); } #[allow(unused_allocation)] // This is intentional. From a9ec05d7e194780906e13d8f053439e27c70f8cf Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 13 Jan 2025 09:57:49 +0100 Subject: [PATCH 4/8] Do not write to extra slots in `dst` --- src/buf/take.rs | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/src/buf/take.rs b/src/buf/take.rs index 0eff707b7..a9d88ec61 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -2,6 +2,9 @@ use crate::{Buf, Bytes}; use core::cmp; +#[cfg(feature = "std")] +use std::io::IoSlice; + /// A `Buf` adapter which limits the bytes read from an underlying buffer. /// /// This struct is generally created by calling `take()` on `Buf`. See @@ -154,21 +157,20 @@ impl Buf for Take { } #[cfg(feature = "std")] - fn chunks_vectored<'a>(&'a self, dst: &mut [std::io::IoSlice<'a>]) -> usize { - let cnt = self.inner.chunks_vectored(dst); - let mut len = 0; - for (n, io) in dst[..cnt].iter_mut().enumerate() { - let max = self.limit - len; - if io.len() >= max { - // In this case, `IoSlice` is longer than our max, so we need to truncate it to the max. - // - // TODO: remove `unsafe` as soon as `IoSlice::as_bytes` is available (rust-lang/rust#111277) - // SAFETY: This reimplements the safe function `IoSlice::as_bytes`. - let buf = unsafe { std::slice::from_raw_parts::<'a, u8>(io.as_ptr(), max) }; - *io = std::io::IoSlice::new(buf); - return n + 1; + fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { + let mut slices = [IoSlice::new(&[]); 16]; + + let cnt = self.inner.chunks_vectored(&mut slices[..dst.len().min(16)]); + let mut limit = self.limit; + for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter().copied()).enumerate() { + if let Some(buf) = slice.get(..limit) { + // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV. + let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(buf) }; + *dst = IoSlice::new(buf); + return i + 1; } - len += io.len(); + *dst = slice; + limit -= slice.len(); } cnt } From 6000c62c231a67bacc399da69c3f75b1e77b6e57 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 13 Jan 2025 10:07:20 +0100 Subject: [PATCH 5/8] IoSlice is not Copy on MSRV --- src/buf/take.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/buf/take.rs b/src/buf/take.rs index a9d88ec61..7cec140fb 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -162,14 +162,14 @@ impl Buf for Take { let cnt = self.inner.chunks_vectored(&mut slices[..dst.len().min(16)]); let mut limit = self.limit; - for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter().copied()).enumerate() { + for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter()).enumerate() { if let Some(buf) = slice.get(..limit) { // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV. let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(buf) }; *dst = IoSlice::new(buf); return i + 1; } - *dst = slice; + *dst = *slice; limit -= slice.len(); } cnt From 2310cd35ea6e44f72c64d94667c0ff20a349da17 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 13 Jan 2025 10:10:03 +0100 Subject: [PATCH 6/8] Handle limit == 0 --- src/buf/take.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/buf/take.rs b/src/buf/take.rs index 7cec140fb..92e8ea061 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -158,6 +158,10 @@ impl Buf for Take { #[cfg(feature = "std")] fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { + if self.limit == 0 { + return 0; + } + let mut slices = [IoSlice::new(&[]); 16]; let cnt = self.inner.chunks_vectored(&mut slices[..dst.len().min(16)]); From 94e58437f89d6a4bfbac1ba10a07dcbebf64f5b7 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 13 Jan 2025 10:13:56 +0100 Subject: [PATCH 7/8] IoSlice is not Copy on MSRV 2 --- src/buf/take.rs | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/src/buf/take.rs b/src/buf/take.rs index 92e8ea061..7c83a6a74 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -162,9 +162,29 @@ impl Buf for Take { return 0; } - let mut slices = [IoSlice::new(&[]); 16]; - - let cnt = self.inner.chunks_vectored(&mut slices[..dst.len().min(16)]); + const LEN: usize = 16; + let mut slices: [IoSlice<'a>; LEN] = [ + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + IoSlice::new(&[]), + ]; + + let cnt = self + .inner + .chunks_vectored(&mut slices[..dst.len().min(LEN)]); let mut limit = self.limit; for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter()).enumerate() { if let Some(buf) = slice.get(..limit) { From bae9dc70827031ff752a0ab70d11bd65aeeec89c Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Mon, 13 Jan 2025 10:19:08 +0100 Subject: [PATCH 8/8] It still isn't Copy --- src/buf/take.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/buf/take.rs b/src/buf/take.rs index 7c83a6a74..fc4e39dda 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -192,9 +192,12 @@ impl Buf for Take { let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(buf) }; *dst = IoSlice::new(buf); return i + 1; + } else { + // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV. + let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(slice) }; + *dst = IoSlice::new(buf); + limit -= slice.len(); } - *dst = *slice; - limit -= slice.len(); } cnt }