#![doc(hidden)]
use crate::{
alloc::{
handle_alloc_error,
AllocErr,
AllocInit::{self, Uninitialized, Zeroed},
AllocRef,
Global,
Layout,
MemoryBlock,
ReallocPlacement::{self, InPlace, MayMove},
},
boxed::Box,
collections::TryReserveError::{self, AllocError, CapacityOverflow},
};
use core::{
cmp,
mem::{self, MaybeUninit},
ops::Drop,
ptr::{self, NonNull, Unique},
slice,
};
#[allow(missing_debug_implementations)]
pub struct RawVec<T, A: AllocRef = Global> {
ptr: Unique<T>,
cap: usize,
alloc: A,
}
impl<T> RawVec<T, Global> {
pub const NEW: Self = Self::new();
pub const fn new() -> Self {
Self::new_in(Global)
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self::with_capacity_in(capacity, Global)
}
#[inline]
pub fn with_capacity_zeroed(capacity: usize) -> Self {
Self::with_capacity_zeroed_in(capacity, Global)
}
#[inline]
pub unsafe fn from_raw_parts(ptr: *mut T, capacity: usize) -> Self {
Self::from_raw_parts_in(ptr, capacity, Global)
}
pub fn from_box(mut slice: Box<[T]>) -> Self {
unsafe {
let result = Self::from_raw_parts(slice.as_mut_ptr(), slice.len());
mem::forget(slice);
result
}
}
}
impl<T, A: AllocRef> RawVec<T, A> {
pub const fn new_in(alloc: A) -> Self {
Self {
ptr: Unique::dangling(),
cap: 0,
alloc,
}
}
#[inline]
pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
Self::allocate_in(capacity, Uninitialized, alloc)
}
#[inline]
pub fn try_with_capacity_zeroed_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
Self::allocate_in(capacity, Zeroed, alloc)
}
#[inline]
pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
match Self::try_with_capacity_in(capacity, alloc) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(v) => v,
}
}
#[inline]
pub fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self {
match Self::try_with_capacity_zeroed_in(capacity, alloc) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(v) => v,
}
}
fn allocate_in(
capacity: usize,
init: AllocInit,
mut alloc: A,
) -> Result<Self, TryReserveError> {
if mem::size_of::<T>() == 0 {
Ok(Self::new_in(alloc))
} else {
let layout = Layout::array::<T>(capacity)?;
alloc_guard(layout.size())?;
let memory = alloc
.alloc(layout, init)
.map_err(|_| AllocError { layout })?;
Ok(Self {
ptr: Unique::new(memory.ptr.cast().as_ptr()).unwrap(),
cap: Self::capacity_from_bytes(memory.size),
alloc,
})
}
}
#[inline]
pub unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, a: A) -> Self {
Self {
ptr: Unique::new_unchecked(ptr),
cap: capacity,
alloc: a,
}
}
pub fn ptr(&self) -> *mut T {
self.ptr.as_ptr()
}
#[inline]
pub fn capacity(&self) -> usize {
if mem::size_of::<T>() == 0 {
usize::MAX
} else {
self.cap
}
}
pub fn alloc(&self) -> &A {
&self.alloc
}
pub fn alloc_mut(&mut self) -> &mut A {
&mut self.alloc
}
fn current_memory(&self) -> Option<(NonNull<u8>, Layout)> {
if mem::size_of::<T>() == 0 || self.cap == 0 {
None
} else {
unsafe {
let align = mem::align_of::<T>();
let size = mem::size_of::<T>() * self.cap;
let layout = Layout::from_size_align_unchecked(size, align);
Some((self.ptr.cast().into(), layout))
}
}
}
#[inline(never)]
#[cold]
pub fn double(&mut self) {
match self.grow(Double, MayMove, Uninitialized) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(()) => { }
}
}
#[inline(never)]
#[cold]
pub fn double_in_place(&mut self) -> bool {
self.grow(Double, InPlace, Uninitialized).is_ok()
}
pub fn reserve(&mut self, used_capacity: usize, needed_extra_capacity: usize) {
match self.try_reserve(used_capacity, needed_extra_capacity) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(()) => { }
}
}
pub fn try_reserve(
&mut self,
used_capacity: usize,
needed_extra_capacity: usize,
) -> Result<(), TryReserveError> {
if self.needs_to_grow(used_capacity, needed_extra_capacity) {
self.grow(
Amortized {
used_capacity,
needed_extra_capacity,
},
MayMove,
Uninitialized,
)
} else {
Ok(())
}
}
pub fn reserve_in_place(&mut self, used_capacity: usize, needed_extra_capacity: usize) -> bool {
if self.needs_to_grow(used_capacity, needed_extra_capacity) {
self.grow(
Amortized {
used_capacity,
needed_extra_capacity,
},
InPlace,
Uninitialized,
)
.is_ok()
} else {
true
}
}
pub fn reserve_exact(&mut self, used_capacity: usize, needed_extra_capacity: usize) {
match self.try_reserve_exact(used_capacity, needed_extra_capacity) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(()) => { }
}
}
pub fn try_reserve_exact(
&mut self,
used_capacity: usize,
needed_extra_capacity: usize,
) -> Result<(), TryReserveError> {
if self.needs_to_grow(used_capacity, needed_extra_capacity) {
self.grow(
Exact {
used_capacity,
needed_extra_capacity,
},
MayMove,
Uninitialized,
)
} else {
Ok(())
}
}
pub fn shrink_to_fit(&mut self, amount: usize) {
match self.try_shrink_to_fit(amount) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(result) => result,
}
}
pub fn try_shrink_to_fit(&mut self, amount: usize) -> Result<(), TryReserveError> {
self.shrink(amount, MayMove)
}
}
#[derive(Copy, Clone)]
enum Strategy {
Double,
Amortized {
used_capacity: usize,
needed_extra_capacity: usize,
},
Exact {
used_capacity: usize,
needed_extra_capacity: usize,
},
}
use Strategy::{Amortized, Double, Exact};
impl<T, A: AllocRef> RawVec<T, A> {
fn needs_to_grow(&self, used_capacity: usize, needed_extra_capacity: usize) -> bool {
needed_extra_capacity > self.capacity().wrapping_sub(used_capacity)
}
fn capacity_from_bytes(excess: usize) -> usize {
debug_assert_ne!(mem::size_of::<T>(), 0);
excess / mem::size_of::<T>()
}
fn set_memory(&mut self, memory: MemoryBlock) {
self.ptr = Unique::new(memory.ptr.cast().as_ptr()).unwrap();
self.cap = Self::capacity_from_bytes(memory.size);
}
fn grow(
&mut self,
strategy: Strategy,
placement: ReallocPlacement,
init: AllocInit,
) -> Result<(), TryReserveError> {
let elem_size = mem::size_of::<T>();
if elem_size == 0 {
return Err(CapacityOverflow);
}
let new_layout = match strategy {
Double => unsafe {
let cap = if self.cap == 0 {
if elem_size > usize::MAX / 8 { 1 } else { 4 }
} else {
self.cap * 2
};
Layout::from_size_align_unchecked(cap * elem_size, mem::align_of::<T>())
},
Amortized {
used_capacity,
needed_extra_capacity,
} => {
let required_cap = used_capacity
.checked_add(needed_extra_capacity)
.ok_or(CapacityOverflow)?;
let double_cap = self.cap * 2;
let cap = cmp::max(double_cap, required_cap);
Layout::array::<T>(cap).map_err(|_| CapacityOverflow)?
}
Exact {
used_capacity,
needed_extra_capacity,
} => {
let cap = used_capacity
.checked_add(needed_extra_capacity)
.ok_or(CapacityOverflow)?;
Layout::array::<T>(cap).map_err(|_| CapacityOverflow)?
}
};
let memory = if let Some((ptr, old_layout)) = self.current_memory() {
debug_assert_eq!(old_layout.align(), new_layout.align());
unsafe {
self.alloc
.grow(ptr, old_layout, new_layout.size(), placement, init)
.map_err(|_| AllocError { layout: new_layout })?
}
} else {
match placement {
MayMove => self.alloc.alloc(new_layout, init),
InPlace => Err(AllocErr),
}
.map_err(|_| AllocError { layout: new_layout })?
};
self.set_memory(memory);
Ok(())
}
fn shrink(
&mut self,
amount: usize,
placement: ReallocPlacement,
) -> Result<(), TryReserveError> {
assert!(
amount <= self.capacity(),
"Tried to shrink to a larger capacity"
);
let (ptr, layout) = if let Some(mem) = self.current_memory() {
mem
} else {
return Ok(());
};
let new_size = amount * mem::size_of::<T>();
let memory = unsafe {
self.alloc
.shrink(ptr, layout, new_size, placement)
.map_err(|_| AllocError {
layout: Layout::from_size_align_unchecked(new_size, layout.align()),
})?
};
self.set_memory(memory);
Ok(())
}
pub unsafe fn into_box(self, len: usize) -> Box<[MaybeUninit<T>], A> {
debug_assert!(
len <= self.capacity(),
"`len` must be smaller than or equal to `self.capacity()`"
);
let slice = slice::from_raw_parts_mut(self.ptr() as *mut MaybeUninit<T>, len);
let output = Box::from_raw_in(slice, ptr::read(self.alloc()));
mem::forget(self);
output
}
}
unsafe impl<#[may_dangle] T, A: AllocRef> Drop for RawVec<T, A> {
fn drop(&mut self) {
if let Some((ptr, layout)) = self.current_memory() {
unsafe { self.alloc.dealloc(ptr, layout) }
}
}
}
#[inline]
fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> {
if mem::size_of::<usize>() < 8 && alloc_size > isize::MAX as usize {
Err(CapacityOverflow)
} else {
Ok(())
}
}
fn capacity_overflow() -> ! {
panic!("capacity overflow");
}