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::{self, SliceIndex};
use alloc::alloc;
use crate::buf::{self, Buf, DefaultAlignment, Padder, StoreBuf};
use crate::endian::{ByteOrder, Native};
use crate::error::Error;
use crate::mem::MaybeUninit;
use crate::pointer::{DefaultSize, Ref, Size};
use crate::traits::{UnsizedZeroCopy, ZeroCopy};
pub struct OwnedBuf<E = Native, O = DefaultSize>
where
E: ByteOrder,
O: Size,
{
data: NonNull<u8>,
len: usize,
capacity: usize,
requested: usize,
align: usize,
_marker: PhantomData<(E, O)>,
}
impl Default for OwnedBuf {
#[inline]
fn default() -> Self {
Self::new()
}
}
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 {
let align = align_of::<T>();
Self {
data: unsafe { dangling(align) },
len: 0,
capacity: 0,
requested: align,
align,
_marker: PhantomData,
}
}
pub fn with_capacity_and_alignment<T>(capacity: usize) -> Self {
unsafe { Self::with_capacity_and_custom_alignment(capacity, align_of::<T>()) }
}
}
impl<E, O> OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
#[inline]
pub fn with_size<U: Size>(self) -> OwnedBuf<E, U> {
let this = ManuallyDrop::new(self);
OwnedBuf {
data: this.data,
len: this.len,
capacity: this.capacity,
requested: this.requested,
align: this.align,
_marker: PhantomData,
}
}
#[inline]
pub fn with_byte_order<U>(self) -> OwnedBuf<U, O>
where
U: ByteOrder,
{
let this = ManuallyDrop::new(self);
OwnedBuf {
data: this.data,
len: this.len,
capacity: this.capacity,
requested: this.requested,
align: this.align,
_marker: PhantomData,
}
}
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(crate) fn as_ptr(&self) -> *const u8 {
self.data.as_ptr() as *const _
}
#[inline]
pub(crate) fn as_mut_ptr(&mut self) -> *mut u8 {
self.data.as_ptr()
}
#[inline]
#[cfg(test)]
pub(crate) fn as_nonnull(&mut self) -> NonNull<u8> {
self.data
}
#[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_mut_ptr(), self.len()) }
}
#[inline]
pub fn store_uninit<T>(&mut self) -> Ref<MaybeUninit<T>, E, 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, U, I>(
&mut self,
reference: Ref<MaybeUninit<T>, U, I>,
) -> &mut MaybeUninit<T>
where
T: ZeroCopy,
U: ByteOrder,
I: Size,
{
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, E, 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, E, O>
where
T: ZeroCopy,
{
let offset = self.len;
let ptr = NonNull::new_unchecked(self.data.as_ptr().add(offset));
buf::store_unaligned(ptr, value);
self.len += size_of::<T>();
Ref::new(offset)
}
#[inline]
pub fn store_unsized<T>(&mut self, value: &T) -> Ref<T, E, O>
where
T: ?Sized + UnsizedZeroCopy,
{
unsafe {
let size = size_of_val(value);
self.next_offset_with_and_reserve(T::ALIGN, size);
let offset = self.len;
let ptr = NonNull::new_unchecked(self.data.as_ptr().add(offset));
ptr.as_ptr().copy_from_nonoverlapping(value.as_ptr(), size);
if T::PADDED {
let mut padder = Padder::new(ptr);
value.pad(&mut padder);
padder.remaining_unsized(value);
}
self.len += size;
Ref::with_metadata(offset, value.metadata())
}
}
#[inline(always)]
pub fn store_slice<T>(&mut self, values: &[T]) -> Ref<[T], E, O>
where
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);
unsafe {
let ptr = self.data.as_ptr().add(self.len);
ptr.write_bytes(byte, len);
self.len += len;
}
}
#[inline]
pub(crate) unsafe fn store_bytes<T>(&mut self, values: &[T])
where
T: ZeroCopy,
{
let dst = self.as_mut_ptr().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 !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 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 = 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 {
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_mut_ptr(), 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_mut_ptr(), 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<E, O> Deref for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
type Target = Buf;
#[inline]
fn deref(&self) -> &Self::Target {
Buf::new(self.as_slice())
}
}
impl<E, O> DerefMut for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
Buf::new_mut(self.as_mut_slice())
}
}
impl<E, O> AsRef<Buf> for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
#[inline]
fn as_ref(&self) -> &Buf {
self
}
}
impl<E, O> AsMut<Buf> for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
#[inline]
fn as_mut(&mut self) -> &mut Buf {
self
}
}
impl<E, O> Borrow<Buf> for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
#[inline]
fn borrow(&self) -> &Buf {
self.as_ref()
}
}
impl<E, O> Clone for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
fn clone(&self) -> Self {
unsafe {
let mut new = ManuallyDrop::new(Self::with_capacity_and_custom_alignment(
self.len, self.align,
));
new.as_mut_ptr()
.copy_from_nonoverlapping(self.as_ptr(), self.len);
new.requested = self.requested;
new.len = self.len;
ManuallyDrop::into_inner(new)
}
}
}
impl<E, O> Drop for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
fn drop(&mut self) {
unsafe {
if self.capacity != 0 {
let layout = Layout::from_size_align_unchecked(self.capacity, self.align);
alloc::dealloc(self.data.as_ptr(), 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) }
}
impl<E, O> StoreBuf for OwnedBuf<E, O>
where
E: ByteOrder,
O: Size,
{
type ByteOrder = E;
type Size = O;
#[inline]
fn len(&self) -> usize {
OwnedBuf::len(self)
}
#[inline]
fn truncate(&mut self, len: usize) {
if self.len > len {
self.len = len;
}
}
#[inline]
fn store_unsized<T>(&mut self, value: &T) -> Ref<T, Self::ByteOrder, Self::Size>
where
T: ?Sized + UnsizedZeroCopy,
{
OwnedBuf::store_unsized(self, value)
}
#[inline]
fn store<T>(&mut self, value: &T) -> Ref<T, Self::ByteOrder, Self::Size>
where
T: ZeroCopy,
{
OwnedBuf::store(self, value)
}
#[inline]
fn swap<T>(
&mut self,
a: Ref<T, Self::ByteOrder, Self::Size>,
b: Ref<T, Self::ByteOrder, Self::Size>,
) -> Result<(), Error>
where
T: ZeroCopy,
{
Buf::swap(self, a, b)
}
#[inline]
fn align_in_place(&mut self) {
OwnedBuf::align_in_place(self);
}
#[inline]
fn next_offset<T>(&mut self) -> usize {
OwnedBuf::next_offset::<T>(self)
}
#[inline]
fn next_offset_with_and_reserve(&mut self, align: usize, reserve: usize) {
OwnedBuf::next_offset_with_and_reserve(self, align, reserve)
}
#[inline]
fn fill(&mut self, byte: u8, len: usize) {
OwnedBuf::fill(self, byte, len);
}
#[inline]
fn get<I>(&self, index: I) -> Option<&I::Output>
where
I: SliceIndex<[u8]>,
{
Buf::get(self, index)
}
#[inline]
fn get_mut<I>(&mut self, index: I) -> Option<&mut I::Output>
where
I: SliceIndex<[u8]>,
{
Buf::get_mut(self, index)
}
#[inline]
fn as_buf(&self) -> &Buf {
self
}
#[inline]
fn as_mut_buf(&mut self) -> &mut Buf {
self
}
}