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-09 02:43:53 +09:00
committed by GitHub
parent fa9cbf1258
commit 2697fa7a9d
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 {