use std::{cmp, mem, num::NonZeroUsize, ptr, slice, sync::atomic::Ordering::Relaxed};
use crate::{BytePageSize, info::Info, info::Kind, stext::StorageVTable, stvec};
#[cfg(target_endian = "little")]
#[repr(C)]
pub(crate) struct Storage {
pub(crate) offset: NonZeroUsize,
pub(crate) ptr: *mut u8,
pub(crate) len: usize,
}
#[cfg(target_endian = "big")]
#[repr(C)]
pub(crate) struct Storage {
pub(crate) len: usize,
pub(crate) ptr: *mut u8,
pub(crate) offset: NonZeroUsize,
}
const KIND_STEXT: usize = 0b00;
const KIND_VEC: usize = 0b01;
const KIND_INLINE: usize = 0b10;
const KIND_STATIC: usize = 0b11;
const KIND_MASK: usize = 0b11;
const KIND_OFFSET_BITS: usize = 2;
pub(crate) const MIN_CAPACITY: usize = 128 - crate::METADATA_SIZE;
const INLINE_LEN_MASK: usize = 0b1111_1100;
#[cfg(target_endian = "little")]
const INLINE_DATA_OFFSET: isize = 1;
#[cfg(target_endian = "big")]
const INLINE_DATA_OFFSET: isize = 0;
#[cfg(target_pointer_width = "64")]
pub(crate) const INLINE_CAP: usize = 3 * 8 - 1;
#[cfg(target_pointer_width = "32")]
pub(crate) const INLINE_CAP: usize = 3 * 4 - 1;
const PTR_INLINE: NonZeroUsize = NonZeroUsize::new(KIND_INLINE).unwrap();
const PTR_STATIC: NonZeroUsize = NonZeroUsize::new(KIND_STATIC).unwrap();
const DEFAUILT_OFFSET: NonZeroUsize =
NonZeroUsize::new((stvec::METADATA_SIZE << KIND_OFFSET_BITS) ^ KIND_VEC).unwrap();
impl Storage {
#[inline]
pub(crate) const fn empty() -> Storage {
Storage {
ptr: ptr::null_mut(),
len: 0,
offset: PTR_INLINE,
}
}
#[inline]
pub(crate) const fn from_static(bytes: &'static [u8]) -> Storage {
let ptr = bytes.as_ptr().cast_mut();
Storage {
ptr,
len: bytes.len(),
offset: PTR_STATIC,
}
}
#[inline]
pub(crate) fn from_stext(
addr: *const u8,
len: usize,
vtable: *const StorageVTable,
) -> Storage {
Storage {
len,
ptr: addr.cast_mut(),
offset: unsafe { NonZeroUsize::new_unchecked(vtable.expose_provenance()) },
}
}
#[inline]
pub(crate) fn from_slice(src: &[u8]) -> Storage {
if src.len() <= INLINE_CAP {
unsafe { Storage::from_ptr_inline(src.as_ptr(), src.len()) }
} else {
Storage::from_slice_with_capacity(src.len(), src)
}
}
#[inline]
fn from_slice_with_capacity(cap: usize, src: &[u8]) -> Storage {
unsafe {
let shared = stvec::SharedVec::create(BytePageSize::Unset, cap, src);
Storage {
len: src.len(),
ptr: shared.as_ptr().add(1).cast::<u8>(),
offset: DEFAUILT_OFFSET,
}
}
}
pub(crate) unsafe fn from_ptr_inline(src: *const u8, len: usize) -> Storage {
let mut st = Storage {
ptr: ptr::null_mut(),
len: 0,
offset: PTR_INLINE,
};
let dst = st.inline_ptr();
ptr::copy(src, dst, len);
st.set_inline_len(len);
st
}
pub(crate) fn as_ref(&self) -> &[u8] {
unsafe {
match self.kind() {
KIND_INLINE => {
slice::from_raw_parts(self.inline_ptr_ro(), self.inline_len())
}
KIND_STEXT => slice::from_raw_parts(self.as_ptr(), self.len()),
_ => slice::from_raw_parts(self.ptr, self.len),
}
}
}
pub(crate) unsafe fn as_ptr(&self) -> *const u8 {
unsafe {
match self.kind() {
KIND_INLINE => self.inline_ptr_ro(),
KIND_STEXT => ((*self.st_vtable()).as_ptr)(self.st_addr(), self.st_len()),
_ => self.ptr,
}
}
}
pub(crate) fn len(&self) -> usize {
match self.kind() {
KIND_STEXT => unsafe {
((*self.st_vtable()).len)(self.st_addr(), self.st_len())
},
KIND_INLINE => self.inline_len(),
_ => self.len,
}
}
pub(crate) fn get_u8(&mut self) -> u8 {
unsafe {
let ret = match self.kind() {
KIND_INLINE => {
assert!(self.inline_len() >= 1);
*self.inline_ptr_ro()
}
KIND_STEXT => {
let vt = &*self.st_vtable();
let len = (vt.len)(self.st_addr(), self.st_len());
assert!(len >= 1);
*(vt.as_ptr)(self.st_addr(), self.st_len())
}
_ => {
assert!(self.len >= 1);
*self.ptr
}
};
self.set_start(1);
ret
}
}
#[inline]
unsafe fn inline_ptr(&mut self) -> *mut u8 {
(ptr::from_mut::<Storage>(self).cast::<u8>()).offset(INLINE_DATA_OFFSET)
}
#[inline]
unsafe fn inline_ptr_ro(&self) -> *const u8 {
ptr::from_ref::<Storage>(self)
.cast::<u8>()
.offset(INLINE_DATA_OFFSET)
}
#[inline]
fn inline_len(&self) -> usize {
(self.offset.get() & INLINE_LEN_MASK) >> KIND_OFFSET_BITS
}
#[inline]
pub(crate) fn is_empty(&self) -> bool {
self.len() == 0
}
#[inline]
pub(crate) fn capacity(&self) -> usize {
let kind = self.kind();
match kind {
KIND_VEC => unsafe { (*self.shared_vec()).capacity() },
KIND_STEXT => unsafe {
((*self.st_vtable()).len)(self.st_addr(), self.st_len())
},
KIND_INLINE => INLINE_CAP,
_ => self.len,
}
}
pub(crate) fn split_off(&mut self, at: usize, create_inline: bool) -> Storage {
let other = unsafe {
if create_inline && self.len() - at <= INLINE_CAP {
Storage::from_ptr_inline(self.as_ptr().add(at), self.len() - at)
} else {
let mut other = self.shallow_clone();
other.set_start(at);
other
}
};
unsafe {
if create_inline && at <= INLINE_CAP {
*self = Storage::from_ptr_inline(self.as_ptr(), at);
} else {
self.set_end(at);
}
}
other
}
pub(crate) fn split_to(&mut self, at: usize) -> Storage {
let other = unsafe {
if at <= INLINE_CAP {
Storage::from_ptr_inline(self.as_ptr(), at)
} else {
let mut other = self.shallow_clone();
other.set_end(at);
other
}
};
unsafe {
self.set_start(at);
}
other
}
pub(crate) fn truncate(&mut self, len: usize) {
unsafe {
if len <= self.len() {
if len < INLINE_CAP {
*self = Storage::from_ptr_inline(self.as_ptr(), len);
} else {
self.set_len(len);
}
}
}
}
pub(crate) fn trimdown(&mut self) {
let kind = self.kind();
if !(kind == KIND_INLINE || kind == KIND_STATIC) {
if self.len() <= INLINE_CAP {
*self = unsafe { Storage::from_ptr_inline(self.as_ptr(), self.len()) };
} else if self.capacity() - self.len() >= 64 {
*self = Storage::from_slice_with_capacity(self.len(), self.as_ref());
}
}
}
#[inline]
pub(crate) unsafe fn set_len(&mut self, len: usize) {
let kind = self.kind();
match kind {
KIND_VEC => {
assert!(len <= self.capacity());
self.len = len;
}
KIND_STEXT => {
if len != self.len() {
assert!(len <= self.capacity());
*self = Storage::from_slice(&self.as_ref()[..len]);
}
}
KIND_INLINE => self.set_inline_len(len),
_ => {
assert!(len <= self.len);
self.len = len;
}
}
}
#[inline]
fn set_inline_len(&mut self, len: usize) {
debug_assert!(len <= INLINE_CAP);
self.offset = unsafe {
NonZeroUsize::new_unchecked(
self.offset.get() & !INLINE_LEN_MASK | (len << KIND_OFFSET_BITS),
)
};
}
pub(crate) unsafe fn set_start(&mut self, start: usize) {
if start == 0 {
return;
}
match self.kind() {
KIND_VEC => {
let shared = self.shared_vec();
let offset = (self.offset.get() >> KIND_OFFSET_BITS) + start;
self.ptr = (shared.cast::<u8>()).add(offset);
if self.len >= start {
self.len -= start;
} else {
self.len = 0;
}
self.offset =
NonZeroUsize::new_unchecked((offset << KIND_OFFSET_BITS) ^ KIND_VEC);
}
KIND_INLINE => {
assert!(start <= INLINE_CAP);
let len = self.inline_len();
if len <= start {
self.set_inline_len(0);
} else {
let new_len = len - start;
let dst = self.inline_ptr();
let src = (dst.cast_const()).add(start);
ptr::copy(src, dst, new_len);
self.set_inline_len(new_len);
}
}
KIND_STEXT => {
assert!(start <= self.capacity());
*self = Storage::from_slice(&self.as_ref()[start..]);
}
_ => {
assert!(start <= self.len);
self.len -= start;
self.ptr = self.ptr.add(start);
}
}
}
pub(crate) unsafe fn set_end(&mut self, end: usize) {
match self.kind() {
KIND_VEC => {
self.len = cmp::min(self.len, end);
}
KIND_INLINE => {
assert!(end <= INLINE_CAP);
let new_len = cmp::min(self.inline_len(), end);
self.set_inline_len(new_len);
}
KIND_STEXT => {
let len = cmp::min(self.len(), end);
*self = Storage::from_slice(&self.as_ref()[..len]);
}
_ => {
assert!(end <= self.len);
self.len = end;
}
}
}
#[inline]
unsafe fn shallow_clone(&self) -> Storage {
let kind = self.kind();
if kind == KIND_INLINE || kind == KIND_STATIC {
let mut inner: mem::MaybeUninit<Storage> = mem::MaybeUninit::uninit();
ptr::copy_nonoverlapping(self, inner.as_mut_ptr(), 1);
inner.assume_init()
} else if kind == KIND_VEC {
let shared = self.shared_vec();
let ref_cnt = (*shared).ref_count.fetch_add(1, Relaxed);
if ref_cnt == u32::MAX {
stvec::abort();
}
Storage { ..*self }
} else {
if let Some((addr, len)) =
((*self.st_vtable()).clone)(self.st_addr(), self.st_len())
{
Storage {
len,
ptr: addr.cast_mut(),
offset: self.offset,
}
} else {
Storage::from_slice(self.as_ref())
}
}
}
#[inline]
pub(crate) fn is_inline(&self) -> bool {
self.kind() == KIND_INLINE
}
#[inline]
fn shared_vec(&self) -> *mut stvec::SharedVec {
let offset = self.offset.get() >> KIND_OFFSET_BITS;
#[allow(clippy::cast_ptr_alignment)]
unsafe {
self.ptr.sub(offset).cast::<stvec::SharedVec>()
}
}
#[inline]
fn st_vtable(&self) -> *const StorageVTable {
ptr::with_exposed_provenance::<StorageVTable>(self.offset.get())
}
#[inline]
fn st_addr(&self) -> *const u8 {
self.ptr.cast_const()
}
#[inline]
fn st_len(&self) -> usize {
self.len
}
#[inline]
fn kind(&self) -> usize {
#[cfg(target_endian = "little")]
#[inline]
fn imp(ptr: usize) -> usize {
ptr & KIND_MASK
}
#[cfg(target_endian = "big")]
#[inline]
fn imp(arc: usize) -> usize {
unsafe {
let p: *const u8 = arc as *const u8;
*p & KIND_MASK
}
}
imp(self.offset.get())
}
pub(crate) fn info(&self) -> Info {
let kind = self.kind();
let (id, refs, capacity) = unsafe {
if kind == KIND_VEC {
let ptr = self.shared_vec();
(
ptr as usize,
(*ptr).ref_count.load(Relaxed),
(*ptr).offset as usize
+ (*ptr).len as usize
+ (*ptr).remaining as usize,
)
} else {
(0, 0, 0)
}
};
Info {
id,
refs,
capacity,
kind: Kind::from_raw(kind),
}
}
}
unsafe impl Send for Storage {}
unsafe impl Sync for Storage {}
impl Clone for Storage {
fn clone(&self) -> Storage {
unsafe { self.shallow_clone() }
}
}
impl Drop for Storage {
fn drop(&mut self) {
match self.kind() {
KIND_VEC => {
stvec::release_shared_vec(self.shared_vec());
}
KIND_STEXT => unsafe {
((*self.st_vtable()).drop)(self.st_addr(), self.st_len());
},
_ => {}
}
}
}
impl Kind {
fn from_raw(n: usize) -> Kind {
match n {
KIND_STEXT => Kind::StExt,
KIND_INLINE => Kind::Inline,
KIND_STATIC => Kind::Static,
_ => Kind::Vec,
}
}
}
#[cfg(test)]
mod tests {
use crate::*;
const LONG: &[u8] =
b"mary had a little lamb, little lamb, little lamb, little lamb, little lamb, little lamb \
mary had a little lamb, little lamb, little lamb, little lamb, little lamb, little lamb \
mary had a little lamb, little lamb, little lamb, little lamb, little lamb, little lamb";
#[test]
fn trimdown() {
let mut b = Bytes::from(LONG.to_vec());
assert_eq!(b.storage.capacity(), 263);
unsafe { b.storage.set_len(68) };
assert_eq!(b.len(), 68);
assert_eq!(&b[..], &LONG[..68]);
assert_eq!(b.storage.capacity(), 263);
b.trimdown();
assert_eq!(&b[..], &LONG[..68]);
assert_eq!(b.storage.capacity(), 68);
unsafe { b.storage.set_len(16) };
assert_eq!(&b[..], &LONG[..16]);
b.trimdown();
assert!(b.is_inline());
}
#[test]
#[allow(clippy::unnecessary_fallible_conversions)]
fn bytes_mut() {
let bv = BytesMut::copy_from_slice(LONG);
assert_eq!(bv.capacity(), 263);
assert_eq!(bv.len(), 263);
assert_eq!(bv.as_ref().len(), 263);
assert_eq!(bv.as_ref(), LONG);
assert_eq!(&bv[..], LONG);
let sl: &[u8] = &[];
let bv = BytesMut::copy_from_slice(sl);
assert_eq!(bv.capacity(), 0);
assert_eq!(bv.len(), 0);
assert_eq!(bv.as_ref().len(), 0);
assert_eq!(bv.as_ref(), sl);
assert_eq!(&bv[..], sl);
let mut bv = BytesMut::copy_from_slice(&b"hello"[..]);
assert_eq!(bv.capacity(), 5);
bv.reserve_capacity(128);
assert_eq!(bv.capacity(), 128);
assert_eq!(bv.len(), 5);
assert_eq!(bv.as_ref(), &b"hello"[..]);
let mut bv = BytesMut::copy_from_slice(&b"hello"[..]);
assert_eq!(bv.capacity(), 5);
assert_eq!(bv.len(), 5);
assert_eq!(bv.as_ref().len(), 5);
assert_eq!(bv.as_ref()[0], b"h"[0]);
assert_eq!(bv.remaining_mut(), 0);
bv.reserve(1);
assert_eq!(bv.remaining_mut(), 1);
bv.put_u8(b" "[0]);
assert_eq!(bv.as_ref(), &b"hello "[..]);
assert_eq!(bv.remaining_mut(), 0);
bv.reserve(5);
assert_eq!(bv.remaining_mut(), 5);
bv.put("world");
assert_eq!(bv, "hello world");
bv.advance_to(6);
assert_eq!(bv, "world");
assert_eq!(bv.remaining_mut(), 0);
let bv = BytesMut::copy_from_slice(&b"hello world"[..]);
let b = Bytes::from(bv);
assert_eq!(b, "hello world");
let mut bv = BytesMut::with_capacity(0);
bv.extend_from_slice(b"hello world.");
bv.extend_from_slice(b"hello world.");
bv.extend_from_slice(b"hello world.");
bv.extend_from_slice(b"hello world.");
let p1 = unsafe { bv.storage.as_ptr() as usize };
bv.advance(48);
assert!(bv.is_empty());
assert_eq!(bv.capacity(), 0);
bv.reserve(48);
assert!(bv.is_empty());
assert_eq!(bv.capacity(), 48);
let p2 = unsafe { bv.storage.as_ptr() as usize };
assert_eq!(p1, p2);
}
}