use std::fmt;
use std::hash::{Hash, Hasher};
use std::mem::ManuallyDrop;
use std::ops::Deref;
use std::ops::DerefMut;
use std::ptr::NonNull;
use crate::{HeapBytesGrowable, InlineBytes23, INLINE_MASK};
pub union CompactBytes {
heap: ManuallyDrop<HeapBytesGrowable>,
inline: InlineBytes23,
}
unsafe impl Send for CompactBytes {}
unsafe impl Sync for CompactBytes {}
static_assertions::assert_eq_align!(
InlineBytes23,
HeapBytesGrowable,
CompactBytes,
Vec<u8>,
usize
);
static_assertions::assert_eq_size!(InlineBytes23, HeapBytesGrowable, CompactBytes, Vec<u8>);
static_assertions::const_assert_eq!(std::mem::size_of::<CompactBytes>(), 24);
impl CompactBytes {
pub const MAX_INLINE: usize = 23;
pub const MIN_HEAP: usize = HeapBytesGrowable::MIN_ALLOCATION_SIZE;
pub const MAX_HEAP: usize = HeapBytesGrowable::MAX_ALLOCATION_SIZE;
#[inline]
pub fn new(slice: &[u8]) -> Self {
if slice.len() <= Self::MAX_INLINE {
let inline = unsafe { InlineBytes23::new(slice) };
CompactBytes { inline }
} else {
let heap = ManuallyDrop::new(HeapBytesGrowable::new(slice));
CompactBytes { heap }
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
if capacity <= Self::MAX_INLINE {
let inline = InlineBytes23::empty();
CompactBytes { inline }
} else {
let heap = ManuallyDrop::new(HeapBytesGrowable::with_capacity(capacity));
CompactBytes { heap }
}
}
#[inline]
pub const fn empty() -> Self {
let inline = InlineBytes23::empty();
CompactBytes { inline }
}
#[inline]
pub unsafe fn from_raw_parts(ptr: *mut u8, length: usize, capacity: usize) -> Self {
let heap = HeapBytesGrowable {
ptr: NonNull::new_unchecked(ptr),
len: length,
cap: capacity,
};
let heap = ManuallyDrop::new(heap);
CompactBytes { heap }
}
#[inline]
pub fn as_slice(&self) -> &[u8] {
let pointer = self.as_ptr();
let length = self.len();
unsafe { core::slice::from_raw_parts(pointer, length) }
}
#[inline]
pub fn as_mut_slice(&mut self) -> &mut [u8] {
let pointer = self.as_mut_ptr();
let length = self.len();
unsafe { core::slice::from_raw_parts_mut(pointer, length) }
}
#[inline(always)]
pub fn len(&self) -> usize {
let (mut length, heap_length) = unsafe { (self.inline.len(), self.heap.len) };
if self.spilled() {
length = heap_length;
}
length
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[inline(always)]
pub fn capacity(&self) -> usize {
let (mut capacity, heap_capacity) = unsafe { (Self::MAX_INLINE, self.heap.cap) };
if self.spilled() {
capacity = heap_capacity;
}
capacity
}
#[inline]
pub fn push(&mut self, byte: u8) {
self.extend_from_slice(&[byte]);
}
#[inline(always)]
pub fn extend_from_slice(&mut self, slice: &[u8]) {
self.reserve(slice.len());
let (ptr, len, cap) = self.as_mut_triple();
let push_ptr = unsafe { ptr.add(len) };
debug_assert!((cap - len) >= slice.len(), "failed to reserve enough space");
unsafe { std::ptr::copy_nonoverlapping(slice.as_ptr(), push_ptr, slice.len()) };
unsafe { self.set_len(len + slice.len()) };
}
#[inline]
pub fn clear(&mut self) {
self.truncate(0);
}
#[inline]
pub fn truncate(&mut self, new_len: usize) {
if new_len >= self.len() {
return;
}
unsafe { self.set_len(new_len) }
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
let len = self.len();
let needed_capacity = len
.checked_add(additional)
.expect("attempt to reserve more than usize::MAX");
if self.capacity() >= needed_capacity {
return;
}
realloc(self, len, additional);
#[cold]
fn realloc(this: &mut CompactBytes, len: usize, additional: usize) {
if !this.spilled() {
let heap = HeapBytesGrowable::with_additional(this.as_slice(), additional);
*this = CompactBytes {
heap: ManuallyDrop::new(heap),
};
} else {
let heap_row = unsafe { &mut this.heap };
let amortized_capacity = HeapBytesGrowable::amortized_growth(len, additional);
if heap_row.realloc(amortized_capacity).is_err() {
let heap = HeapBytesGrowable::with_additional(this.as_slice(), additional);
let heap = ManuallyDrop::new(heap);
*this = CompactBytes { heap };
}
}
}
}
#[inline]
pub fn into_vec(self) -> Vec<u8> {
if self.spilled() {
let heap = unsafe { &self.heap };
let vec = unsafe { Vec::from_raw_parts(heap.ptr.as_ptr(), heap.len, heap.cap) };
std::mem::forget(self);
vec
} else {
self.as_slice().to_vec()
}
}
#[inline(always)]
pub fn spilled(&self) -> bool {
unsafe { self.inline.data < INLINE_MASK }
}
#[inline]
unsafe fn set_len(&mut self, new_len: usize) {
if self.spilled() {
self.heap.set_len(new_len);
} else {
self.inline.set_len(new_len);
}
}
#[inline(always)]
fn as_ptr(&self) -> *const u8 {
let mut pointer = self as *const Self as *const u8;
if self.spilled() {
pointer = unsafe { self.heap.ptr }.as_ptr()
}
pointer
}
#[inline(always)]
fn as_mut_ptr(&mut self) -> *mut u8 {
let mut pointer = self as *mut Self as *mut u8;
if self.spilled() {
pointer = unsafe { self.heap.ptr }.as_ptr()
}
pointer
}
#[inline(always)]
fn as_mut_triple(&mut self) -> (*mut u8, usize, usize) {
let ptr = self.as_mut_ptr();
let len = self.len();
let cap = self.capacity();
(ptr, len, cap)
}
}
impl Default for CompactBytes {
#[inline]
fn default() -> Self {
CompactBytes::new(&[])
}
}
impl Deref for CompactBytes {
type Target = [u8];
#[inline]
fn deref(&self) -> &[u8] {
self.as_slice()
}
}
impl DerefMut for CompactBytes {
#[inline]
fn deref_mut(&mut self) -> &mut [u8] {
self.as_mut_slice()
}
}
impl AsRef<[u8]> for CompactBytes {
#[inline]
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
impl<T: AsRef<[u8]>> PartialEq<T> for CompactBytes {
#[inline]
fn eq(&self, other: &T) -> bool {
self.as_slice() == other.as_ref()
}
}
impl Eq for CompactBytes {}
impl Hash for CompactBytes {
fn hash<H: Hasher>(&self, state: &mut H) {
self.as_slice().hash(state)
}
}
impl fmt::Debug for CompactBytes {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}", self.as_slice())
}
}
impl Drop for CompactBytes {
#[inline]
fn drop(&mut self) {
#[cold]
fn outlined_drop(this: &mut CompactBytes) {
let heap = unsafe { &mut this.heap };
heap.dealloc();
}
if self.spilled() {
outlined_drop(self);
}
}
}
impl Clone for CompactBytes {
#[inline]
fn clone(&self) -> Self {
#[cold]
fn outlined_clone(this: &CompactBytes) -> CompactBytes {
CompactBytes::new(this.as_slice())
}
if self.spilled() {
outlined_clone(self)
} else {
let inline = unsafe { &self.inline };
CompactBytes { inline: *inline }
}
}
#[inline]
fn clone_from(&mut self, source: &Self) {
self.clear();
self.extend_from_slice(source.as_slice());
}
}
impl From<Vec<u8>> for CompactBytes {
#[inline]
fn from(mut value: Vec<u8>) -> Self {
if value.is_empty() {
let inline = InlineBytes23::empty();
return CompactBytes { inline };
}
let (ptr, len, cap) = (value.as_mut_ptr(), value.len(), value.capacity());
let ptr = unsafe { NonNull::new_unchecked(ptr) };
std::mem::forget(value);
let heap = HeapBytesGrowable { ptr, len, cap };
CompactBytes {
heap: ManuallyDrop::new(heap),
}
}
}
impl serde::Serialize for CompactBytes {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
self.as_slice().serialize(serializer)
}
}
impl<'de> serde::Deserialize<'de> for CompactBytes {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
deserialize_compact_bytes(deserializer)
}
}
fn deserialize_compact_bytes<'de: 'a, 'a, D: serde::Deserializer<'de>>(
deserializer: D,
) -> Result<CompactBytes, D::Error> {
struct CompactBytesVisitor;
impl<'a> serde::de::Visitor<'a> for CompactBytesVisitor {
type Value = CompactBytes;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("bytes")
}
fn visit_seq<A: serde::de::SeqAccess<'a>>(
self,
mut seq: A,
) -> Result<Self::Value, A::Error> {
let mut bytes = CompactBytes::default();
if let Some(capacity_hint) = seq.size_hint() {
bytes.reserve(capacity_hint);
}
while let Some(elem) = seq.next_element::<u8>()? {
bytes.push(elem)
}
Ok(bytes)
}
fn visit_borrowed_bytes<E: serde::de::Error>(self, v: &'a [u8]) -> Result<Self::Value, E> {
Ok(CompactBytes::new(v))
}
fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
Ok(CompactBytes::new(v))
}
}
deserializer.deserialize_bytes(CompactBytesVisitor)
}
#[cfg(test)]
mod test {
use proptest::prelude::*;
use test_case::test_case;
use test_strategy::proptest;
use super::{CompactBytes, HeapBytesGrowable};
#[test]
fn test_bitcode() {
let obj = CompactBytes::new(b"hello world");
let encoded = bitcode::serialize(&obj).unwrap();
let decoded: CompactBytes = bitcode::deserialize(&encoded).unwrap();
assert_eq!(obj.as_slice(), decoded.as_slice());
}
#[test]
fn test_empty() {
let obj = const { CompactBytes::empty() };
assert_eq!(obj.as_slice(), [0u8; 0].as_slice());
assert!(obj.is_empty());
assert!(!obj.spilled())
}
#[test]
#[cfg_attr(miri, ignore)]
fn test_discriminant() {
let mut buf = vec![0u8; 32];
let heap = HeapBytesGrowable {
ptr: unsafe { std::ptr::NonNull::new_unchecked(buf.as_mut_ptr()) },
len: 0,
cap: usize::MAX >> 1,
};
let repr = CompactBytes {
heap: std::mem::ManuallyDrop::new(heap),
};
assert!(repr.spilled());
std::mem::forget(repr);
let bad_heap = HeapBytesGrowable {
ptr: unsafe { std::ptr::NonNull::new_unchecked(buf.as_mut_ptr()) },
len: 0,
cap: usize::MAX,
};
let repr = CompactBytes {
heap: std::mem::ManuallyDrop::new(bad_heap),
};
assert!(!repr.spilled());
std::mem::forget(repr);
}
#[test_case(&[], 0 ; "empty")]
#[test_case(b"hello world", 11 ; "short")]
#[test_case(b"can fit 23 bytes inline", 23 ; "max_inline")]
#[test_case(b"24 bytes and will spill!", 24 ; "first_spill")]
#[test_case(b"i am very large and will spill to the heap", 42 ; "heap")]
fn smoketest_row(slice: &[u8], expected_len: usize) {
let repr = CompactBytes::new(slice);
assert_eq!(repr.len(), expected_len);
assert_eq!(repr.as_slice(), slice);
assert_eq!(repr.spilled(), expected_len > CompactBytes::MAX_INLINE);
}
#[test_case(&[], &[] ; "empty_empty")]
#[test_case(&[], &[1, 2, 3, 4] ; "empty_inline")]
#[test_case(&[], b"once extended I will end up on the heap" ; "empty_heap")]
#[test_case(&[1, 2], &[3, 4] ; "inline_inline")]
#[test_case(&[1, 2, 3, 4], b"i am some more bytes, i will be on the heap, woohoo!" ; "inline_heap")]
#[test_case(b"this row will start on the heap because it's large", b"and this will keep it on the heap" ; "heap_heap")]
fn smoketest_extend(initial: &[u8], other: &[u8]) {
let mut repr = CompactBytes::new(initial);
repr.extend_from_slice(other);
let mut control = initial.to_vec();
control.extend_from_slice(other);
assert_eq!(repr.len(), control.len());
assert_eq!(repr.as_slice(), control.as_slice());
}
#[test_case(&[] ; "empty")]
#[test_case(b"i am smol" ; "inline")]
#[test_case(b"i am large and will end up on the heap" ; "heap")]
fn smoketest_clear(initial: &[u8]) {
let mut repr = CompactBytes::new(initial);
let capacity = repr.capacity();
assert_eq!(repr.as_slice(), initial);
repr.clear();
assert!(repr.as_slice().is_empty());
assert_eq!(repr.len(), 0);
assert_eq!(repr.capacity(), capacity);
}
#[test_case(&[] ; "empty")]
#[test_case(b"smol" ; "inline")]
#[test_case(b"large large large large large large" ; "heap")]
fn smoketest_clone(initial: &[u8]) {
let repr_a = CompactBytes::new(initial);
let repr_b = repr_a.clone();
assert_eq!(repr_a.len(), repr_b.len());
assert_eq!(repr_a.capacity(), repr_b.capacity());
assert_eq!(repr_a.as_slice(), repr_b.as_slice());
}
#[test_case(&[], &[], false ; "empty_empty")]
#[test_case(&[], b"hello", false ; "empty_inline")]
#[test_case(&[], b"I am long and will be on the heap", true ; "empty_heap")]
#[test_case(b"short", &[], false ; "inline_empty")]
#[test_case(b"hello", b"world", false ; "inline_inline")]
#[test_case(b"i am short", b"I am long and will be on the heap", true ; "inline_heap")]
fn smoketest_clone_from(a: &[u8], b: &[u8], should_reallocate: bool) {
let mut a = CompactBytes::new(a);
let a_capacity = a.capacity();
let a_pointer = a.as_slice().as_ptr();
let b = CompactBytes::new(b);
a.clone_from(&b);
assert_eq!(a.capacity() != a_capacity, should_reallocate);
assert_eq!(a.as_slice().as_ptr() != a_pointer, should_reallocate);
}
#[test_case(vec![] ; "empty")]
#[test_case(vec![0, 1, 2, 3, 4] ; "inline")]
#[test_case(b"I am long and will be on the heap, yada yada yada".to_vec() ; "heap")]
fn smoketest_from_vec(initial: Vec<u8>) {
let control = initial.clone();
let pointer = initial.as_ptr();
let repr = CompactBytes::from(initial);
assert_eq!(control.len(), repr.len());
assert_eq!(control.as_slice(), repr.as_slice());
assert_eq!(repr.spilled(), !control.is_empty());
assert_eq!(repr.as_ptr() == pointer, !control.is_empty());
}
#[test]
fn test_cloning_inlines() {
let mut c = CompactBytes::with_capacity(48);
c.push(42);
assert_eq!(c.as_slice(), &[42]);
assert_eq!(c.capacity(), 48);
assert!(c.spilled());
let clone = c.clone();
assert_eq!(clone.as_slice(), &[42]);
assert_eq!(clone.capacity(), CompactBytes::MAX_INLINE);
assert!(!clone.spilled());
}
#[test]
fn test_cloning_drops_excess_capacity() {
let mut c = CompactBytes::with_capacity(48);
c.extend_from_slice(&[42; 32]);
assert_eq!(c.as_slice(), &[42; 32]);
assert_eq!(c.capacity(), 48);
assert_eq!(c.len(), 32);
assert!(c.spilled());
let clone = c.clone();
assert_eq!(clone.as_slice(), &[42; 32]);
assert_eq!(clone.capacity(), 32);
assert_eq!(clone.capacity(), clone.len());
assert!(clone.spilled());
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_row(initial: Vec<u8>) {
let repr = CompactBytes::new(&initial);
prop_assert_eq!(repr.as_slice(), initial.as_slice());
prop_assert_eq!(repr.len(), initial.len());
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_extend(initial: Vec<u8>, other: Vec<u8>) {
let mut repr = CompactBytes::new(&initial);
repr.extend_from_slice(other.as_slice());
let mut control = initial;
control.extend_from_slice(other.as_slice());
prop_assert_eq!(repr.as_slice(), control.as_slice());
prop_assert_eq!(repr.len(), control.len());
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_clear(initial: Vec<u8>) {
let mut repr = CompactBytes::new(&initial);
let capacity = repr.capacity();
repr.clear();
assert!(repr.as_slice().is_empty());
assert_eq!(repr.len(), 0);
assert_eq!(repr.capacity(), capacity);
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_clear_then_extend(initial: Vec<u8>, a: Vec<u8>) {
let mut repr = CompactBytes::new(&initial);
let capacity = repr.capacity();
let pointer = repr.as_slice().as_ptr();
repr.clear();
assert!(repr.as_slice().is_empty());
assert_eq!(repr.len(), 0);
assert_eq!(repr.capacity(), capacity);
repr.extend_from_slice(&a);
assert_eq!(repr.as_slice(), &a);
assert_eq!(repr.len(), a.len());
if a.len() < capacity {
assert_eq!(repr.capacity(), capacity);
assert_eq!(repr.as_slice().as_ptr(), pointer);
}
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_clone(initial: Vec<u8>) {
let repr_a = CompactBytes::new(&initial);
let repr_b = repr_a.clone();
assert_eq!(repr_a.len(), repr_b.len());
assert_eq!(repr_a.capacity(), repr_b.capacity());
assert_eq!(repr_a.as_slice(), repr_b.as_slice());
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_from_vec(initial: Vec<u8>) {
let control = initial.clone();
let pointer = initial.as_ptr();
let repr = CompactBytes::from(initial);
assert_eq!(control.len(), repr.len());
assert_eq!(control.as_slice(), repr.as_slice());
assert_eq!(repr.spilled(), !control.is_empty());
assert_eq!(repr.as_ptr() == pointer, !control.is_empty());
}
#[proptest]
#[cfg_attr(miri, ignore)]
fn proptest_serde(initial: Vec<u8>) {
let repr = CompactBytes::new(&initial);
let (repr_json, ctrl_json) = match (
serde_json::to_string(&repr),
serde_json::to_string(&initial),
) {
(Ok(r), Ok(c)) => (r, c),
(Err(_), Err(_)) => return Ok(()),
(r, c) => panic!("Got mismatched results when serializing {r:?}, {c:?}"),
};
prop_assert_eq!(&repr_json, &ctrl_json);
let (repr_rnd_trip, ctrl_rnd_trip): (CompactBytes, Vec<u8>) = match (
serde_json::from_str(&repr_json),
serde_json::from_str(&ctrl_json),
) {
(Ok(r), Ok(c)) => (r, c),
(Err(_), Err(_)) => return Ok(()),
(r, c) => panic!("Got mismatched results {r:?}, {c:?}"),
};
prop_assert_eq!(&repr, &repr_rnd_trip);
prop_assert_eq!(repr_rnd_trip, ctrl_rnd_trip);
}
#[test]
fn test_heap_bytes_capacity() {
let heap = HeapBytesGrowable::with_capacity(1);
drop(heap);
}
}