diff options
Diffstat (limited to 'src/buf')
-rw-r--r-- | src/buf/buf_impl.rs | 3 | ||||
-rw-r--r-- | src/buf/buf_mut.rs | 77 | ||||
-rw-r--r-- | src/buf/chain.rs | 27 | ||||
-rw-r--r-- | src/buf/take.rs | 12 | ||||
-rw-r--r-- | src/buf/uninit_slice.rs | 7 |
5 files changed, 114 insertions, 12 deletions
diff --git a/src/buf/buf_impl.rs b/src/buf/buf_impl.rs index 16ad8a7..a33c8a4 100644 --- a/src/buf/buf_impl.rs +++ b/src/buf/buf_impl.rs @@ -127,6 +127,9 @@ pub trait Buf { /// This function should never panic. Once the end of the buffer is reached, /// i.e., `Buf::remaining` returns 0, calls to `chunk()` should return an /// empty slice. + // The `chunk` method was previously called `bytes`. This alias makes the rename + // more easily discoverable. + #[cfg_attr(docsrs, doc(alias = "bytes"))] fn chunk(&self) -> &[u8]; /// Fills `dst` with potentially multiple slices starting at `self`'s diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs index f736727..4c2bd2c 100644 --- a/src/buf/buf_mut.rs +++ b/src/buf/buf_mut.rs @@ -33,6 +33,10 @@ pub unsafe trait BufMut { /// This value is greater than or equal to the length of the slice returned /// by `chunk_mut()`. /// + /// Writing to a `BufMut` may involve allocating more memory on the fly. + /// Implementations may fail before reaching the number of bytes indicated + /// by this method if they encounter an allocation failure. + /// /// # Examples /// /// ``` @@ -158,6 +162,12 @@ pub unsafe trait BufMut { /// `chunk_mut()` returning an empty slice implies that `remaining_mut()` will /// return 0 and `remaining_mut()` returning 0 implies that `chunk_mut()` will /// return an empty slice. + /// + /// This function may trigger an out-of-memory abort if it tries to allocate + /// memory and fails to do so. + // The `chunk_mut` method was previously called `bytes_mut`. This alias makes the + // rename more easily discoverable. + #[cfg_attr(docsrs, doc(alias = "bytes_mut"))] fn chunk_mut(&mut self) -> &mut UninitSlice; /// Transfer bytes into `self` from `src` and advance the cursor by the @@ -251,6 +261,37 @@ pub unsafe trait BufMut { } } + /// Put `cnt` bytes `val` into `self`. + /// + /// Logically equivalent to calling `self.put_u8(val)` `cnt` times, but may work faster. + /// + /// `self` must have at least `cnt` remaining capacity. + /// + /// ``` + /// use bytes::BufMut; + /// + /// let mut dst = [0; 6]; + /// + /// { + /// let mut buf = &mut dst[..]; + /// buf.put_bytes(b'a', 4); + /// + /// assert_eq!(2, buf.remaining_mut()); + /// } + /// + /// assert_eq!(b"aaaa\0\0", &dst); + /// ``` + /// + /// # Panics + /// + /// This function panics if there is not enough remaining capacity in + /// `self`. + fn put_bytes(&mut self, val: u8, cnt: usize) { + for _ in 0..cnt { + self.put_u8(val); + } + } + /// Writes an unsigned 8 bit integer to `self`. /// /// The current position is advanced by 1. @@ -693,7 +734,7 @@ pub unsafe trait BufMut { self.put_slice(&n.to_le_bytes()[0..nbytes]); } - /// Writes a signed n-byte integer to `self` in big-endian byte order. + /// Writes low `nbytes` of a signed integer to `self` in big-endian byte order. /// /// The current position is advanced by `nbytes`. /// @@ -703,19 +744,19 @@ pub unsafe trait BufMut { /// use bytes::BufMut; /// /// let mut buf = vec![]; - /// buf.put_int(0x010203, 3); + /// buf.put_int(0x0504010203, 3); /// assert_eq!(buf, b"\x01\x02\x03"); /// ``` /// /// # Panics /// /// This function panics if there is not enough remaining capacity in - /// `self`. + /// `self` or if `nbytes` is greater than 8. fn put_int(&mut self, n: i64, nbytes: usize) { self.put_slice(&n.to_be_bytes()[mem::size_of_val(&n) - nbytes..]); } - /// Writes a signed n-byte integer to `self` in little-endian byte order. + /// Writes low `nbytes` of a signed integer to `self` in little-endian byte order. /// /// The current position is advanced by `nbytes`. /// @@ -725,14 +766,14 @@ pub unsafe trait BufMut { /// use bytes::BufMut; /// /// let mut buf = vec![]; - /// buf.put_int_le(0x010203, 3); + /// buf.put_int_le(0x0504010203, 3); /// assert_eq!(buf, b"\x03\x02\x01"); /// ``` /// /// # Panics /// /// This function panics if there is not enough remaining capacity in - /// `self`. + /// `self` or if `nbytes` is greater than 8. fn put_int_le(&mut self, n: i64, nbytes: usize) { self.put_slice(&n.to_le_bytes()[0..nbytes]); } @@ -1009,12 +1050,29 @@ unsafe impl BufMut for &mut [u8] { let (_, b) = core::mem::replace(self, &mut []).split_at_mut(cnt); *self = b; } + + #[inline] + fn put_slice(&mut self, src: &[u8]) { + self[..src.len()].copy_from_slice(src); + unsafe { + self.advance_mut(src.len()); + } + } + + fn put_bytes(&mut self, val: u8, cnt: usize) { + assert!(self.remaining_mut() >= cnt); + unsafe { + ptr::write_bytes(self.as_mut_ptr(), val, cnt); + self.advance_mut(cnt); + } + } } unsafe impl BufMut for Vec<u8> { #[inline] fn remaining_mut(&self) -> usize { - usize::MAX - self.len() + // A vector can never have more than isize::MAX bytes + core::isize::MAX as usize - self.len() } #[inline] @@ -1072,6 +1130,11 @@ unsafe impl BufMut for Vec<u8> { fn put_slice(&mut self, src: &[u8]) { self.extend_from_slice(src); } + + fn put_bytes(&mut self, val: u8, cnt: usize) { + let new_len = self.len().checked_add(cnt).unwrap(); + self.resize(new_len, val); + } } // The existence of this function makes the compiler catch if the BufMut diff --git a/src/buf/chain.rs b/src/buf/chain.rs index d68bc2d..9ce5f23 100644 --- a/src/buf/chain.rs +++ b/src/buf/chain.rs @@ -1,5 +1,5 @@ use crate::buf::{IntoIter, UninitSlice}; -use crate::{Buf, BufMut}; +use crate::{Buf, BufMut, Bytes}; #[cfg(feature = "std")] use std::io::IoSlice; @@ -135,7 +135,7 @@ where U: Buf, { fn remaining(&self) -> usize { - self.a.remaining() + self.b.remaining() + self.a.remaining().checked_add(self.b.remaining()).unwrap() } fn chunk(&self) -> &[u8] { @@ -170,6 +170,24 @@ where n += self.b.chunks_vectored(&mut dst[n..]); n } + + fn copy_to_bytes(&mut self, len: usize) -> Bytes { + let a_rem = self.a.remaining(); + if a_rem >= len { + self.a.copy_to_bytes(len) + } else if a_rem == 0 { + self.b.copy_to_bytes(len) + } else { + assert!( + len - a_rem <= self.b.remaining(), + "`len` greater than remaining" + ); + let mut ret = crate::BytesMut::with_capacity(len); + ret.put(&mut self.a); + ret.put((&mut self.b).take(len - a_rem)); + ret.freeze() + } + } } unsafe impl<T, U> BufMut for Chain<T, U> @@ -178,7 +196,10 @@ where U: BufMut, { fn remaining_mut(&self) -> usize { - self.a.remaining_mut() + self.b.remaining_mut() + self.a + .remaining_mut() + .checked_add(self.b.remaining_mut()) + .unwrap() } fn chunk_mut(&mut self) -> &mut UninitSlice { diff --git a/src/buf/take.rs b/src/buf/take.rs index 1747f6e..d3cb10a 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -1,11 +1,11 @@ -use crate::Buf; +use crate::{Buf, Bytes}; use core::cmp; /// A `Buf` adapter which limits the bytes read from an underlying buffer. /// /// This struct is generally created by calling `take()` on `Buf`. See -/// documentation of [`take()`](trait.BufExt.html#method.take) for more details. +/// documentation of [`take()`](trait.Buf.html#method.take) for more details. #[derive(Debug)] pub struct Take<T> { inner: T, @@ -144,4 +144,12 @@ impl<T: Buf> Buf for Take<T> { self.inner.advance(cnt); self.limit -= cnt; } + + fn copy_to_bytes(&mut self, len: usize) -> Bytes { + assert!(len <= self.remaining(), "`len` greater than remaining"); + + let r = self.inner.copy_to_bytes(len); + self.limit -= len; + r + } } diff --git a/src/buf/uninit_slice.rs b/src/buf/uninit_slice.rs index 73f4e89..fb67c0a 100644 --- a/src/buf/uninit_slice.rs +++ b/src/buf/uninit_slice.rs @@ -40,6 +40,7 @@ impl UninitSlice { /// /// let slice = unsafe { UninitSlice::from_raw_parts_mut(ptr, len) }; /// ``` + #[inline] pub unsafe fn from_raw_parts_mut<'a>(ptr: *mut u8, len: usize) -> &'a mut UninitSlice { let maybe_init: &mut [MaybeUninit<u8>] = core::slice::from_raw_parts_mut(ptr as *mut _, len); @@ -64,6 +65,7 @@ impl UninitSlice { /// /// assert_eq!(b"boo", &data[..]); /// ``` + #[inline] pub fn write_byte(&mut self, index: usize, byte: u8) { assert!(index < self.len()); @@ -90,6 +92,7 @@ impl UninitSlice { /// /// assert_eq!(b"bar", &data[..]); /// ``` + #[inline] pub fn copy_from_slice(&mut self, src: &[u8]) { use core::ptr; @@ -116,6 +119,7 @@ impl UninitSlice { /// let mut slice = &mut data[..]; /// let ptr = BufMut::chunk_mut(&mut slice).as_mut_ptr(); /// ``` + #[inline] pub fn as_mut_ptr(&mut self) -> *mut u8 { self.0.as_mut_ptr() as *mut _ } @@ -133,6 +137,7 @@ impl UninitSlice { /// /// assert_eq!(len, 3); /// ``` + #[inline] pub fn len(&self) -> usize { self.0.len() } @@ -150,6 +155,7 @@ macro_rules! impl_index { impl Index<$t> for UninitSlice { type Output = UninitSlice; + #[inline] fn index(&self, index: $t) -> &UninitSlice { let maybe_uninit: &[MaybeUninit<u8>] = &self.0[index]; unsafe { &*(maybe_uninit as *const [MaybeUninit<u8>] as *const UninitSlice) } @@ -157,6 +163,7 @@ macro_rules! impl_index { } impl IndexMut<$t> for UninitSlice { + #[inline] fn index_mut(&mut self, index: $t) -> &mut UninitSlice { let maybe_uninit: &mut [MaybeUninit<u8>] = &mut self.0[index]; unsafe { &mut *(maybe_uninit as *mut [MaybeUninit<u8>] as *mut UninitSlice) } |