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`.
This commit is contained in:
Stepan Koltsov
2021-08-08 18:43:53 +01:00
committed by GitHub
parent 6d42b4ccf4
commit d766a57017
4 changed files with 83 additions and 0 deletions
+44
View File
@@ -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
+13
View File
@@ -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 {
+18
View File
@@ -27,6 +27,14 @@ fn test_vec_as_mut_buf() {
assert_eq!(buf.len(), 68);
}
#[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);
@@ -103,6 +111,16 @@ fn test_mut_slice() {
assert_eq!(&v, &[0, 0, 0, 42]);
}
#[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;
+8
View File
@@ -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());
}