Skip to content

Commit

Permalink
BufMut::put_bytes(self, val, cnt)
Browse files Browse the repository at this point in the history
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`.
  • Loading branch information
stepancheg committed Aug 8, 2021
1 parent ab8e3c0 commit 81e557a
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 0 deletions.
44 changes: 44 additions & 0 deletions src/buf/buf_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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
Expand Down
13 changes: 13 additions & 0 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
18 changes: 18 additions & 0 deletions tests/test_buf_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -75,6 +83,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;
Expand Down
8 changes: 8 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}

0 comments on commit 81e557a

Please sign in to comment.