diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/buf/buf_mut.rs | 44 | ||||
-rw-r--r-- | src/bytes_mut.rs | 13 |
2 files changed, 57 insertions, 0 deletions
diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs index b946658..4c2bd2c 100644 --- a/src/buf/buf_mut.rs +++ b/src/buf/buf_mut.rs @@ -261,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. @@ -1027,6 +1058,14 @@ unsafe impl BufMut for &mut [u8] { 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> { @@ -1091,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/bytes_mut.rs b/src/bytes_mut.rs index 8e42079..b754ed1 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1010,6 +1010,19 @@ unsafe impl BufMut for BytesMut { fn put_slice(&mut self, src: &[u8]) { self.extend_from_slice(src); } + + fn put_bytes(&mut self, val: u8, cnt: usize) { + self.reserve(cnt); + unsafe { + let dst = self.uninit_slice(); + // Reserved above + debug_assert!(dst.len() >= cnt); + + ptr::write_bytes(dst.as_mut_ptr(), val, cnt); + + self.advance_mut(cnt); + } + } } impl AsRef<[u8]> for BytesMut { |