diff --git a/Cargo.toml b/Cargo.toml index b2961dc..bb74a99 100644 --- a/Cargo.toml +++ b/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 } diff --git a/src/lib.rs b/src/lib.rs index 3459da5..f05bc3f 100644 --- a/src/lib.rs +++ b/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 io::Write for SmallVec { Ok(()) } } + +#[cfg(feature = "bytes")] +unsafe impl BufMut for SmallVec { + #[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(&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); + } +} diff --git a/src/tests.rs b/src/tests.rs index 925eb3b..624f44c 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1154,9 +1154,7 @@ fn collect_from_iter() { #[test] fn test_collect_with_spill() { let input = "0123456"; - let collected: SmallVec = input - .chars() - .collect(); + let collected: SmallVec = 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::(), 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; + + #[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); + } + } +}