feat: impl bytes::BufMut for SmallVec (v2) (#371)

Author: BugenZhao

Cargo.toml

@@ -19,6 +19,7 @@ may_dangle = []
 extract_if = []
 
 [dependencies]
+bytes = { version = "1", optional = true, default-features = false }
 serde = { version = "1", optional = true, default-features = false }
 malloc_size_of = { version = "0.1.1", optional = true, default-features = false }
 

src/lib.rs

@@ -92,6 +92,8 @@ use core::ptr::copy;
 use core::ptr::copy_nonoverlapping;
 use core::ptr::NonNull;
 
+#[cfg(feature = "bytes")]
+use bytes::{buf::UninitSlice, BufMut};
 #[cfg(feature = "malloc_size_of")]
 use malloc_size_of::{MallocShallowSizeOf, MallocSizeOf, MallocSizeOfOps};
 #[cfg(feature = "serde")]
@@ -2220,3 +2222,71 @@ impl<const N: usize> io::Write for SmallVec<u8, N> {
         Ok(())
     }
 }
+
+#[cfg(feature = "bytes")]
+unsafe impl<const N: usize> BufMut for SmallVec<u8, N> {
+    #[inline]
+    fn remaining_mut(&self) -> usize {
+        // A vector can never have more than isize::MAX bytes
+        isize::MAX as usize - self.len()
+    }
+
+    #[inline]
+    unsafe fn advance_mut(&mut self, cnt: usize) {
+        let len = self.len();
+        let remaining = self.capacity() - len;
+
+        if remaining < cnt {
+            panic!("advance out of bounds: the len is {remaining} but advancing by {cnt}");
+        }
+
+        // Addition will not overflow since the sum is at most the capacity.
+        self.set_len(len + cnt);
+    }
+
+    #[inline]
+    fn chunk_mut(&mut self) -> &mut UninitSlice {
+        if self.capacity() == self.len() {
+            self.reserve(64); // Grow the smallvec
+        }
+
+        let cap = self.capacity();
+        let len = self.len();
+
+        let ptr = self.as_mut_ptr();
+        // SAFETY: Since `ptr` is valid for `cap` bytes, `ptr.add(len)` must be
+        // valid for `cap - len` bytes. The subtraction will not underflow since
+        // `len <= cap`.
+        unsafe { UninitSlice::from_raw_parts_mut(ptr.add(len), cap - len) }
+    }
+
+    // Specialize these methods so they can skip checking `remaining_mut`
+    // and `advance_mut`.
+    #[inline]
+    fn put<T: bytes::Buf>(&mut self, mut src: T)
+    where
+        Self: Sized,
+    {
+        // In case the src isn't contiguous, reserve upfront.
+        self.reserve(src.remaining());
+
+        while src.has_remaining() {
+            let s = src.chunk();
+            let l = s.len();
+            self.extend_from_slice(s);
+            src.advance(l);
+        }
+    }
+
+    #[inline]
+    fn put_slice(&mut self, src: &[u8]) {
+        self.extend_from_slice(src);
+    }
+
+    #[inline]
+    fn put_bytes(&mut self, val: u8, cnt: usize) {
+        // If the addition overflows, then the `resize` will fail.
+        let new_len = self.len().saturating_add(cnt);
+        self.resize(new_len, val);
+    }
+}

src/tests.rs

@@ -1154,9 +1154,7 @@ fn collect_from_iter() {
 #[test]
 fn test_collect_with_spill() {
     let input = "0123456";
-    let collected: SmallVec<char, 4> = input
-        .chars()
-        .collect();
+    let collected: SmallVec<char, 4> = input.chars().collect();
     assert_eq!(collected, &['0', '1', '2', '3', '4', '5', '6']);
 }
 
@@ -1186,3 +1184,96 @@ fn test_spare_capacity_mut() {
     assert!(spare.len() >= 1);
     assert_eq!(spare.as_ptr().cast::<u8>(), unsafe { v.as_ptr().add(3) });
 }
+
+// Adopted from `tests/test_buf_mut.rs` in the `bytes` crate.
+#[cfg(feature = "bytes")]
+mod buf_mut {
+    use bytes::BufMut as _;
+
+    type SmallVec = crate::SmallVec<u8, 8>;
+
+    #[test]
+    fn test_smallvec_as_mut_buf() {
+        let mut buf = SmallVec::with_capacity(64);
+
+        assert_eq!(buf.remaining_mut(), isize::MAX as usize);
+
+        assert!(buf.chunk_mut().len() >= 64);
+
+        buf.put(&b"zomg"[..]);
+
+        assert_eq!(&buf, b"zomg");
+
+        assert_eq!(buf.remaining_mut(), isize::MAX as usize - 4);
+        assert_eq!(buf.capacity(), 64);
+
+        for _ in 0..16 {
+            buf.put(&b"zomg"[..]);
+        }
+
+        assert_eq!(buf.len(), 68);
+    }
+
+    #[test]
+    fn test_smallvec_put_bytes() {
+        let mut buf = SmallVec::new();
+        buf.push(17);
+        buf.put_bytes(19, 2);
+        assert_eq!([17, 19, 19], &buf[..]);
+    }
+
+    #[test]
+    fn test_put_u8() {
+        let mut buf = SmallVec::with_capacity(8);
+        buf.put_u8(33);
+        assert_eq!(b"\x21", &buf[..]);
+    }
+
+    #[test]
+    fn test_put_u16() {
+        let mut buf = SmallVec::with_capacity(8);
+        buf.put_u16(8532);
+        assert_eq!(b"\x21\x54", &buf[..]);
+
+        buf.clear();
+        buf.put_u16_le(8532);
+        assert_eq!(b"\x54\x21", &buf[..]);
+    }
+
+    #[test]
+    fn test_put_int() {
+        let mut buf = SmallVec::with_capacity(8);
+        buf.put_int(0x1020304050607080, 3);
+        assert_eq!(b"\x60\x70\x80", &buf[..]);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_put_int_nbytes_overflow() {
+        let mut buf = SmallVec::with_capacity(8);
+        buf.put_int(0x1020304050607080, 9);
+    }
+
+    #[test]
+    fn test_put_int_le() {
+        let mut buf = SmallVec::with_capacity(8);
+        buf.put_int_le(0x1020304050607080, 3);
+        assert_eq!(b"\x80\x70\x60", &buf[..]);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_put_int_le_nbytes_overflow() {
+        let mut buf = SmallVec::with_capacity(8);
+        buf.put_int_le(0x1020304050607080, 9);
+    }
+
+    #[test]
+    #[should_panic(expected = "advance out of bounds: the len is 8 but advancing by 12")]
+    fn test_smallvec_advance_mut() {
+        let mut buf = SmallVec::with_capacity(8);
+        unsafe {
+            buf.advance_mut(12);
+        }
+    }
+}