From 2697fa7a9dd871d8d64d1959b7ac46227b7b23b4 Mon Sep 17 00:00:00 2001
From: Stepan Koltsov <stepan.koltsov@gmail.com>
Date: Sun, 8 Aug 2021 18:43:53 +0100
Subject: [PATCH] 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`.
---
 src/buf/buf_mut.rs    | 44 +++++++++++++++++++++++++++++++++++++++++++
 src/bytes_mut.rs      | 13 +++++++++++++
 tests/test_buf_mut.rs | 18 ++++++++++++++++++
 tests/test_bytes.rs   |  8 ++++++++
 4 files changed, 83 insertions(+)

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
@@ -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;
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());
+}