use core::alloc::Layout;
use core::borrow::Borrow;
use core::marker::PhantomData;
use core::mem::{align_of, size_of, size_of_val, ManuallyDrop};
use core::ops::{Deref, DerefMut};
use core::ptr::NonNull;
use core::slice;
use alloc::alloc;
use crate::buf::{Buf, BufMut, DefaultAlignment};
use crate::mem::MaybeUninit;
use crate::pointer::{DefaultSize, Pointee, Ref, Size};
use crate::traits::{UnsizedZeroCopy, ZeroCopy};
pub struct OwnedBuf<O: Size = DefaultSize> {
data: NonNull<u8>,
len: usize,
capacity: usize,
requested: usize,
align: usize,
_marker: PhantomData<O>,
}
impl OwnedBuf {
pub const fn new() -> Self {
Self::with_alignment::<DefaultAlignment>()
}
pub fn with_capacity(capacity: usize) -> Self {
Self::with_capacity_and_alignment::<DefaultAlignment>(capacity)
}
pub const fn with_alignment<T>() -> Self
where
T: ZeroCopy,
{
let align = align_of::<T>();
Self {
data: unsafe { dangling(align) },
len: 0,
capacity: 0,
requested: align,
align,
_marker: PhantomData,
}
}
}
impl<O: Size> OwnedBuf<O> {
pub fn with_capacity_and_alignment<T>(capacity: usize) -> Self
where
T: ZeroCopy,
{
unsafe { Self::with_capacity_and_custom_alignment(capacity, align_of::<T>()) }
}
pub(crate) unsafe fn with_capacity_and_custom_alignment(capacity: usize, align: usize) -> Self where
{
if capacity == 0 {
return Self {
data: dangling(align),
len: 0,
capacity: 0,
requested: align,
align,
_marker: PhantomData,
};
}
let layout = Layout::from_size_align(capacity, align).expect("Illegal memory layout");
unsafe {
let data = alloc::alloc(layout);
if data.is_null() {
alloc::handle_alloc_error(layout);
}
Self {
data: NonNull::new_unchecked(data),
len: 0,
capacity,
requested: align,
align,
_marker: PhantomData,
}
}
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn clear(&mut self) {
self.len = 0;
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn capacity(&self) -> usize {
self.capacity
}
#[inline]
pub fn requested(&self) -> usize {
self.requested
}
#[inline]
pub fn reserve(&mut self, capacity: usize) {
let new_capacity = self.len + capacity;
self.ensure_capacity(new_capacity);
}
#[inline]
pub unsafe fn advance(&mut self, size: usize) {
self.len += size;
}
#[inline]
pub unsafe fn as_buf_mut(&mut self) -> BufMut<'_> {
BufMut::new(self.data.as_ptr())
}
#[inline]
pub fn as_ptr(&self) -> *const u8 {
self.data.as_ptr() as *const _
}
#[inline]
pub fn as_ptr_mut(&mut self) -> *mut u8 {
self.data.as_ptr()
}
#[inline]
pub fn as_slice(&self) -> &[u8] {
unsafe { slice::from_raw_parts(self.as_ptr(), self.len()) }
}
#[inline]
pub fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { slice::from_raw_parts_mut(self.as_ptr_mut(), self.len()) }
}
#[inline]
pub fn store_uninit<T>(&mut self) -> Ref<MaybeUninit<T>, O>
where
T: ZeroCopy,
{
unsafe {
self.next_offset_with_and_reserve(align_of::<T>(), size_of::<T>());
let offset = self.len;
self.data
.as_ptr()
.add(self.len)
.write_bytes(0, size_of::<T>());
self.len += size_of::<T>();
Ref::new(offset)
}
}
#[inline]
pub fn load_uninit_mut<T>(&mut self, reference: Ref<MaybeUninit<T>>) -> &mut MaybeUninit<T>
where
T: ZeroCopy,
{
let at = reference.offset();
assert!(at + size_of::<T>() <= self.len, "Length overflow");
unsafe { &mut *(self.data.as_ptr().add(at) as *mut MaybeUninit<T>) }
}
#[inline]
pub fn store<T>(&mut self, value: &T) -> Ref<T, O>
where
T: ZeroCopy,
{
self.next_offset_with_and_reserve(align_of::<T>(), size_of::<T>());
unsafe { self.store_unchecked(value) }
}
#[inline]
pub unsafe fn store_unchecked<T>(&mut self, value: &T) -> Ref<T, O>
where
T: ZeroCopy,
{
let offset = self.len;
let mut buf_mut = BufMut::new(self.data.as_ptr().wrapping_add(offset));
buf_mut.store_unaligned(value);
self.len += size_of::<T>();
Ref::new(offset)
}
#[inline]
pub fn store_unsized<T: ?Sized>(&mut self, value: &T) -> Ref<T, O>
where
T: Pointee<O, Packed = O, Metadata = usize>,
T: UnsizedZeroCopy<T, O>,
{
unsafe {
self.next_offset_with_and_reserve(T::ALIGN, value.size());
let offset = self.len;
value.store(&mut BufMut::new(self.data.as_ptr().wrapping_add(offset)));
self.len += value.size();
Ref::with_metadata(offset, value.metadata())
}
}
#[inline(always)]
pub fn store_slice<T>(&mut self, values: &[T]) -> Ref<[T], O>
where
[T]: Pointee<O, Packed = O, Metadata = usize>,
T: ZeroCopy,
{
self.store_unsized(values)
}
pub fn extend_from_slice(&mut self, bytes: &[u8]) {
self.reserve(bytes.len());
unsafe {
self.store_bytes(bytes);
}
}
pub(crate) fn fill(&mut self, byte: u8, len: usize) {
self.reserve(len);
let base = self.data.as_ptr().wrapping_add(self.len);
unsafe {
base.write_bytes(byte, len);
self.len += len;
}
}
#[inline]
pub unsafe fn store_bytes<T>(&mut self, values: &[T])
where
T: ZeroCopy,
{
let dst = self.as_ptr_mut().wrapping_add(self.len);
dst.copy_from_nonoverlapping(values.as_ptr().cast(), size_of_val(values));
self.len += size_of_val(values);
}
#[inline]
pub fn align_in_place(&mut self) {
if !unsafe { crate::buf::is_aligned_with(self.as_ptr(), self.requested) } {
let (old_layout, new_layout) = self.layouts(self.capacity);
self.alloc_new(old_layout, new_layout);
}
}
#[inline]
pub fn into_aligned(mut self) -> Self {
self.align_in_place();
self
}
#[inline]
pub fn request_align<T>(&mut self)
where
T: ZeroCopy,
{
self.requested = self.requested.max(align_of::<T>());
self.ensure_aligned_and_reserve(align_of::<T>(), size_of::<T>());
}
#[inline]
fn ensure_aligned_and_reserve(&mut self, align: usize, reserve: usize) {
let extra = crate::buf::padding_to(self.len, align);
self.reserve(extra + reserve);
unsafe {
self.data.as_ptr().add(self.len).write_bytes(0, extra);
self.len += extra;
}
}
#[inline]
pub(crate) fn next_offset_with_and_reserve(&mut self, align: usize, reserve: usize) {
self.requested = self.requested.max(align);
self.ensure_aligned_and_reserve(align, reserve);
}
#[inline]
pub fn next_offset<T>(&mut self) -> usize
where
T: ZeroCopy,
{
self.next_offset_with_and_reserve(align_of::<T>(), size_of::<T>());
self.len
}
#[inline(never)]
fn ensure_capacity(&mut self, new_capacity: usize) {
let new_capacity = new_capacity.max(self.requested);
if self.capacity >= new_capacity {
return;
}
let new_capacity = new_capacity.max((self.capacity as f32 * 1.5) as usize);
let (old_layout, new_layout) = self.layouts(new_capacity);
if old_layout.size() == 0 {
self.alloc_init(new_layout);
} else if new_layout.align() == old_layout.align() {
self.alloc_realloc(old_layout, new_layout);
} else {
self.alloc_new(old_layout, new_layout);
}
}
#[inline]
fn layouts(&self, new_capacity: usize) -> (Layout, Layout) {
let old_layout = unsafe { Layout::from_size_align_unchecked(self.capacity, self.align) };
let layout =
Layout::from_size_align(new_capacity, self.requested).expect("Proposed layout invalid");
(old_layout, layout)
}
fn alloc_init(&mut self, new_layout: Layout) {
unsafe {
let ptr = alloc::alloc(new_layout);
if ptr.is_null() {
alloc::handle_alloc_error(new_layout);
}
self.data = NonNull::new_unchecked(ptr);
self.capacity = new_layout.size();
self.align = self.requested;
}
}
fn alloc_realloc(&mut self, old_layout: Layout, new_layout: Layout) {
debug_assert_eq!(old_layout.align(), new_layout.align());
unsafe {
let ptr = alloc::realloc(self.as_ptr_mut(), old_layout, new_layout.size());
if ptr.is_null() {
alloc::handle_alloc_error(old_layout);
}
self.data = NonNull::new_unchecked(ptr);
self.capacity = new_layout.size();
}
}
#[inline(always)]
fn alloc_new(&mut self, old_layout: Layout, new_layout: Layout) {
unsafe {
let ptr = alloc::alloc(new_layout);
if ptr.is_null() {
alloc::handle_alloc_error(new_layout);
}
ptr.copy_from_nonoverlapping(self.as_ptr(), self.len);
alloc::dealloc(self.as_ptr_mut(), old_layout);
self.data = NonNull::new_unchecked(ptr);
self.capacity = new_layout.size();
self.align = self.requested;
}
}
}
unsafe impl Send for OwnedBuf {}
unsafe impl Sync for OwnedBuf {}
impl<O: Size> Deref for OwnedBuf<O> {
type Target = Buf;
#[inline]
fn deref(&self) -> &Self::Target {
Buf::new(self.as_slice())
}
}
impl<O: Size> DerefMut for OwnedBuf<O> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
Buf::new_mut(self.as_mut_slice())
}
}
impl<O: Size> AsRef<Buf> for OwnedBuf<O> {
#[inline]
fn as_ref(&self) -> &Buf {
self
}
}
impl<O: Size> AsMut<Buf> for OwnedBuf<O> {
#[inline]
fn as_mut(&mut self) -> &mut Buf {
self
}
}
impl Borrow<Buf> for OwnedBuf {
#[inline]
fn borrow(&self) -> &Buf {
self.as_ref()
}
}
impl<O: Size> Clone for OwnedBuf<O> {
fn clone(&self) -> Self {
unsafe {
let mut new = ManuallyDrop::new(Self::with_capacity_and_custom_alignment(
self.len, self.align,
));
new.as_ptr_mut()
.copy_from_nonoverlapping(self.as_ptr(), self.len);
new.requested = self.requested;
new.len = self.len;
ManuallyDrop::into_inner(new)
}
}
}
impl<O: Size> Drop for OwnedBuf<O> {
fn drop(&mut self) {
unsafe {
if self.capacity != 0 {
let layout = Layout::from_size_align_unchecked(self.capacity, self.align);
alloc::dealloc(self.as_ptr_mut(), layout);
}
}
}
}
const unsafe fn dangling(align: usize) -> NonNull<u8> {
NonNull::new_unchecked(invalid_mut(align))
}
#[allow(clippy::useless_transmute)]
const fn invalid_mut<T>(addr: usize) -> *mut T {
unsafe { core::mem::transmute(addr) }
}