diff options
author | 2021-08-08 18:43:53 +0100 | |
---|---|---|
committer | 2021-08-09 02:43:53 +0900 | |
commit | 2697fa7a9dd871d8d64d1959b7ac46227b7b23b4 (patch) | |
tree | c588b368a944cc77af3781c4b711b344501205d6 | |
parent | fa9cbf1258cea7812d4009639ed8700b987d8d4e (diff) | |
download | bytes-2697fa7a9dd871d8d64d1959b7ac46227b7b23b4.tar.gz bytes-2697fa7a9dd871d8d64d1959b7ac46227b7b23b4.tar.zst bytes-2697fa7a9dd871d8d64d1959b7ac46227b7b23b4.zip |
BufMut::put_bytes(self, val, cnt) (#487)
Equivalent to
```
for _ in 0..cnt {
self.put_u8(val);
}
```
but may work faster.
Name and signature is chosen to be consistent with `ptr::write_bytes`.
Include three specializations:
* `Vec<u8>`
* `&mut [u8]`
* `BytesMut`
`BytesMut` and `&mut [u8]` specializations use `ptr::write`, `Vec<u8>`
specialization uses `Vec::resize`.
-rw-r--r-- | src/buf/buf_mut.rs | 44 | ||||
-rw-r--r-- | src/bytes_mut.rs | 13 | ||||
-rw-r--r-- | tests/test_buf_mut.rs | 18 | ||||
-rw-r--r-- | tests/test_bytes.rs | 8 |
4 files changed, 83 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 { diff --git a/tests/test_buf_mut.rs b/tests/test_buf_mut.rs index 78a76bc..53f4e86 100644 --- a/tests/test_buf_mut.rs +++ b/tests/test_buf_mut.rs @@ -28,6 +28,14 @@ fn test_vec_as_mut_buf() { } #[test] +fn test_vec_put_bytes() { + let mut buf = Vec::new(); + buf.push(17); + buf.put_bytes(19, 2); + assert_eq!([17, 19, 19], &buf[..]); +} + +#[test] fn test_put_u8() { let mut buf = Vec::with_capacity(8); buf.put_u8(33); @@ -104,6 +112,16 @@ fn test_mut_slice() { } #[test] +fn test_slice_put_bytes() { + let mut v = [0, 0, 0, 0]; + let mut s = &mut v[..]; + s.put_u8(17); + s.put_bytes(19, 2); + assert_eq!(1, s.remaining_mut()); + assert_eq!(&[17, 19, 19, 0], &v[..]); +} + +#[test] fn test_deref_bufmut_forwards() { struct Special; diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 2d22fdd..0914155 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -986,3 +986,11 @@ fn bytes_with_capacity_but_empty() { let vec = Vec::with_capacity(1); let _ = Bytes::from(vec); } + +#[test] +fn bytes_put_bytes() { + let mut bytes = BytesMut::new(); + bytes.put_u8(17); + bytes.put_bytes(19, 2); + assert_eq!([17, 19, 19], bytes.as_ref()); +} |