297 lines
No EOL
8.6 KiB
Rust
297 lines
No EOL
8.6 KiB
Rust
use std::{alloc::Layout, cmp, marker::PhantomData, mem};
|
||
|
||
use std::{alloc, ptr};
|
||
use unique::Unique;
|
||
|
||
/// A [`Vec`] with its elements aligned to a runtime alignment value.
|
||
pub struct AVec<T> {
|
||
buf: Unique<u8>,
|
||
cap: usize,
|
||
len: usize,
|
||
align: usize,
|
||
_marker: PhantomData<T>,
|
||
}
|
||
|
||
impl<T> AVec<T> {
|
||
// Tiny Vecs are dumb. Skip to:
|
||
// - 8 if the element size is 1, because any heap allocators are likely
|
||
// to round up a request of less than 8 bytes to at least 8 bytes.
|
||
// - 4 if elements are moderate-sized (<= 1 KiB).
|
||
// - 1 otherwise, to avoid wasting too much space for very short Vecs.
|
||
//
|
||
// Taken from Rust's standard library RawVec
|
||
pub(crate) const MIN_NON_ZERO_CAP: usize = if mem::size_of::<T>() == 1 {
|
||
8
|
||
} else if mem::size_of::<T>() <= 1024 {
|
||
4
|
||
} else {
|
||
1
|
||
};
|
||
|
||
#[inline]
|
||
pub fn new(alignment: usize) -> Self {
|
||
debug_assert!(mem::size_of::<T>() > 0, "ZSTs not yet supported");
|
||
|
||
Self {
|
||
buf: Unique::dangling(),
|
||
cap: 0,
|
||
len: 0,
|
||
align: alignment,
|
||
_marker: PhantomData
|
||
}
|
||
}
|
||
|
||
/// Constructs a new, empty `AVec` with at least the specified capacity.
|
||
///
|
||
/// The aligned vector will be able to hold at least `capacity` elements without reallocating.
|
||
/// This method may allocate for more elements than `capacity`. If `capacity` is zero,
|
||
/// the vector will not allocate.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if the capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
pub fn with_capacity(alignment: usize, capacity: usize) -> Self {
|
||
let mut s = Self::new(alignment);
|
||
|
||
if capacity > 0 {
|
||
unsafe {
|
||
s.grow_amortized(0, capacity);
|
||
}
|
||
}
|
||
|
||
s
|
||
}
|
||
|
||
/// Calculates the size of the 'slot' for a single **aligned** item.
|
||
#[inline(always)]
|
||
fn slot_size(&self) -> usize {
|
||
let a = self.align - 1;
|
||
(mem::align_of::<T>() + (a)) & !a
|
||
}
|
||
|
||
/// # Panics
|
||
///
|
||
/// Panics if the new capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
unsafe fn grow_amortized(&mut self, len: usize, additional: usize) {
|
||
debug_assert!(additional > 0);
|
||
|
||
let required_cap = len.checked_add(additional)
|
||
.expect("Capacity overflow");
|
||
|
||
let cap = cmp::max(self.cap * 2, required_cap);
|
||
let cap = cmp::max(Self::MIN_NON_ZERO_CAP, cap);
|
||
|
||
let new_layout = Layout::from_size_align_unchecked(cap * self.slot_size(), self.align);
|
||
|
||
let ptr = alloc::alloc(new_layout);
|
||
self.buf = Unique::new_unchecked(ptr);
|
||
self.cap = cap;
|
||
}
|
||
|
||
/// # Panics
|
||
///
|
||
/// Panics if the new capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
unsafe fn grow_exact(&mut self, len: usize, additional: usize) {
|
||
debug_assert!(additional > 0);
|
||
|
||
let cap = len.checked_add(additional)
|
||
.expect("Capacity overflow");
|
||
|
||
let new_layout = Layout::from_size_align_unchecked(cap * self.slot_size(), self.align);
|
||
|
||
let ptr = alloc::alloc(new_layout);
|
||
self.buf = Unique::new_unchecked(ptr);
|
||
self.cap = cap;
|
||
}
|
||
|
||
/// Reserves capacity for at least `additional` more elements.
|
||
///
|
||
/// The collection may reserve more space to speculatively avoid frequent reallocations.
|
||
/// After calling `reserve`, capacity will be greater than or equal to
|
||
/// `self.len() + additional`. Does nothing if capacity is already sufficient.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if the new capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
pub fn reserve(&mut self, additional: usize) {
|
||
debug_assert!(additional > 0);
|
||
|
||
let remaining = self.capacity().wrapping_sub(self.len);
|
||
|
||
if additional > remaining {
|
||
unsafe { self.grow_amortized(self.len, additional) };
|
||
}
|
||
}
|
||
|
||
/// Reserves capacity for `additional` more elements.
|
||
///
|
||
/// Unlike [`reserve`], this will not over-allocate to speculatively avoid frequent
|
||
/// reallocations. After calling `reserve_exact`, capacity will be equal to
|
||
/// `self.len() + additional`. Does nothing if the capacity is already sufficient.
|
||
///
|
||
/// Prefer [`reserve`] if future insertions are expected.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if the new capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
pub fn reserve_exact(&mut self, additional: usize) {
|
||
let remaining = self.capacity().wrapping_sub(self.len);
|
||
|
||
if additional > remaining {
|
||
unsafe { self.grow_exact(self.len, additional) };
|
||
}
|
||
}
|
||
|
||
/// Appends an element to the back of the collection.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if the new capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
pub fn push(&mut self, val: T) {
|
||
if self.len == self.cap {
|
||
self.reserve(self.slot_size());
|
||
}
|
||
|
||
unsafe {
|
||
// SAFETY: the length is ensured to be less than the capacity.
|
||
self.set_at_unchecked(self.len, val);
|
||
}
|
||
|
||
self.len += 1;
|
||
}
|
||
|
||
/// Sets an element at position `idx` within the vector to `val`.
|
||
///
|
||
/// # Unsafe
|
||
///
|
||
/// If `self.len > idx`, bytes past the length of the vector will be written to, potentially
|
||
/// also writing past the capacity of the vector.
|
||
#[inline(always)]
|
||
unsafe fn set_at_unchecked(&mut self, idx: usize, val: T) {
|
||
let ptr = self.buf
|
||
.as_ptr()
|
||
.add(idx * self.slot_size());
|
||
|
||
std::ptr::write(ptr.cast::<T>(), val);
|
||
}
|
||
|
||
/// Sets an element at position `idx` within the vector to `val`.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if `idx >= self.len`.
|
||
#[inline(always)]
|
||
pub fn set_at(&mut self, idx: usize, val: T) {
|
||
assert!(self.len > idx);
|
||
|
||
unsafe {
|
||
self.set_at_unchecked(idx, val);
|
||
}
|
||
}
|
||
|
||
/// Shortens the vector, keeping the first `len` elements and dropping the rest.
|
||
///
|
||
/// If `len` is greater or equal to the vector’s current length, this has no effect.
|
||
#[inline]
|
||
pub fn truncate(&mut self, len: usize) {
|
||
if len > self.len {
|
||
return;
|
||
}
|
||
|
||
unsafe {
|
||
// drop each element past the new length
|
||
for i in len..self.len {
|
||
let ptr = self.buf.as_ptr()
|
||
.add(i * self.slot_size())
|
||
.cast::<T>();
|
||
|
||
ptr::drop_in_place(ptr);
|
||
}
|
||
}
|
||
|
||
self.len = len;
|
||
}
|
||
|
||
#[inline(always)]
|
||
pub fn as_ptr(&self) -> *const u8 {
|
||
self.buf.as_ptr()
|
||
}
|
||
|
||
#[inline(always)]
|
||
pub fn as_mut_ptr(&self) -> *mut u8 {
|
||
self.buf.as_ptr()
|
||
}
|
||
|
||
/// Returns the alignment of the elements in the vector.
|
||
#[inline(always)]
|
||
pub fn align(&self) -> usize {
|
||
self.align
|
||
}
|
||
|
||
/// Returns the length of the vector.
|
||
#[inline(always)]
|
||
pub fn len(&self) -> usize {
|
||
self.len
|
||
}
|
||
|
||
#[inline(always)]
|
||
pub fn is_empty(&self) -> bool {
|
||
self.len == 0
|
||
}
|
||
|
||
/// Returns the capacity of the vector.
|
||
///
|
||
/// The capacity is the amount of elements that the vector can store without reallocating.
|
||
#[inline(always)]
|
||
pub fn capacity(&self) -> usize {
|
||
self.cap
|
||
}
|
||
}
|
||
|
||
impl<T: Clone> AVec<T> {
|
||
/// Resized the `AVec` in-place so that `len` is equal to `new_len`.
|
||
///
|
||
/// If `new_len` is greater than `len`, the `AVec` is extended by the difference, and
|
||
/// each additional slot is filled with `value`. If `new_len` is less than `len`,
|
||
/// the `AVec` will be truncated by to be `new_len`
|
||
///
|
||
/// This method requires `T` to implement [`Clone`] in order to clone the passed value.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if the new capacity exceeds `usize::MAX` bytes.
|
||
#[inline]
|
||
pub fn resize(&mut self, new_len: usize, value: T) {
|
||
if new_len > self.len {
|
||
self.reserve(new_len - self.len);
|
||
|
||
unsafe {
|
||
let mut ptr = self.buf
|
||
.as_ptr().add(self.len * self.slot_size());
|
||
|
||
// write all elements besides the last one
|
||
for _ in 1..new_len {
|
||
std::ptr::write(ptr.cast::<T>(), value.clone());
|
||
ptr = ptr.add(self.slot_size());
|
||
self.len += 1;
|
||
}
|
||
|
||
if new_len > 0 {
|
||
// the last element can be written without cloning
|
||
std::ptr::write(ptr.cast::<T>(), value.clone());
|
||
self.len += 1;
|
||
}
|
||
|
||
self.len = new_len;
|
||
}
|
||
} else {
|
||
self.truncate(new_len);
|
||
}
|
||
}
|
||
} |