use crate::ser::ScratchSpace;
use core::{
alloc::Layout,
borrow::{Borrow, BorrowMut},
fmt,
mem::MaybeUninit,
ops,
ptr::NonNull,
slice,
};
pub struct ScratchVec<T> {
ptr: NonNull<T>,
cap: usize,
len: usize,
}
impl<T> Drop for ScratchVec<T> {
fn drop(&mut self) {
for i in 0..self.len {
unsafe {
core::ptr::drop_in_place(self.ptr.as_ptr().add(i));
}
}
}
}
unsafe impl<T: Send> Send for ScratchVec<T> {}
unsafe impl<T: Sync> Sync for ScratchVec<T> {}
impl<T> ScratchVec<T> {
#[inline]
pub unsafe fn new<S: ScratchSpace + ?Sized>(
scratch_space: &mut S,
capacity: usize,
) -> Result<Self, S::Error> {
let layout = Layout::array::<T>(capacity).unwrap();
if layout.size() == 0 {
Ok(Self {
ptr: NonNull::dangling(),
cap: capacity,
len: 0,
})
} else {
let ptr = scratch_space.push_scratch(layout)?;
Ok(Self {
ptr: ptr.cast(),
cap: capacity,
len: 0,
})
}
}
#[inline]
pub unsafe fn free<S: ScratchSpace + ?Sized>(
self,
scratch_space: &mut S,
) -> Result<(), S::Error> {
let layout = self.layout();
if layout.size() != 0 {
let ptr = self.ptr.cast();
core::mem::drop(self);
scratch_space.pop_scratch(ptr, layout)?;
}
Ok(())
}
#[inline]
fn layout(&self) -> Layout {
Layout::array::<T>(self.cap).unwrap()
}
#[inline]
pub fn clear(&mut self) {
self.len = 0;
}
#[inline]
pub fn as_mut_ptr(&mut self) -> *mut T {
self.ptr.as_ptr()
}
#[inline]
pub fn as_mut_slice(&mut self) -> &mut [T] {
unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
}
#[inline]
pub fn as_ptr(&self) -> *const T {
self.ptr.as_ptr()
}
#[inline]
pub fn as_slice(&self) -> &[T] {
unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
}
#[inline]
pub fn capacity(&self) -> usize {
self.cap
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
if self.len + additional > self.cap {
panic!("reserve requested more capacity than the scratch vec has available");
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn extend_from_slice(&mut self, other: &[T]) {
if !other.is_empty() {
self.reserve(other.len());
unsafe {
core::ptr::copy_nonoverlapping(
other.as_ptr(),
self.as_mut_ptr().add(self.len()),
other.len(),
);
}
self.len += other.len();
}
}
#[inline]
pub fn pop(&mut self) -> Option<T> {
if self.len == 0 {
None
} else {
unsafe {
self.len -= 1;
Some(self.as_ptr().add(self.len()).read())
}
}
}
#[inline]
pub fn push(&mut self, value: T) {
unsafe {
self.reserve(1);
self.as_mut_ptr().add(self.len).write(value);
self.len += 1;
}
}
#[inline]
pub fn reserve_exact(&mut self, additional: usize) {
self.reserve(additional);
}
#[inline]
pub unsafe fn set_len(&mut self, new_len: usize) {
debug_assert!(new_len <= self.capacity());
self.len = new_len;
}
#[inline]
fn drain_range<R>(range: R, bounds: ops::RangeTo<usize>) -> ops::Range<usize>
where
R: ops::RangeBounds<usize>,
{
let len = bounds.end;
let start: ops::Bound<&usize> = range.start_bound();
let start = match start {
ops::Bound::Included(&start) => start,
ops::Bound::Excluded(start) => start
.checked_add(1)
.unwrap_or_else(|| panic!("attempted to index slice from after maximum usize")),
ops::Bound::Unbounded => 0,
};
let end: ops::Bound<&usize> = range.end_bound();
let end = match end {
ops::Bound::Included(end) => end
.checked_add(1)
.unwrap_or_else(|| panic!("attempted to index slice up to maximum usize")),
ops::Bound::Excluded(&end) => end,
ops::Bound::Unbounded => len,
};
if start > end {
panic!("slice index starts at {} but ends at {}", start, end);
}
if end > len {
panic!(
"range start index {} out of range for slice of length {}",
end, len
);
}
ops::Range { start, end }
}
#[inline]
pub fn drain<R: ops::RangeBounds<usize>>(&mut self, range: R) -> Drain<'_, T> {
let len = self.len();
let ops::Range { start, end } = Self::drain_range(range, ..len);
unsafe {
self.set_len(start);
let range_slice = slice::from_raw_parts_mut(self.as_mut_ptr().add(start), end - start);
Drain {
tail_start: end,
tail_len: len - end,
iter: range_slice.iter(),
vec: NonNull::from(self),
}
}
}
}
impl<T> ScratchVec<MaybeUninit<T>> {
#[inline]
pub fn assume_init(self) -> ScratchVec<T> {
ScratchVec {
ptr: self.ptr.cast(),
cap: self.cap,
len: self.len,
}
}
}
impl<T> AsMut<[T]> for ScratchVec<T> {
#[inline]
fn as_mut(&mut self) -> &mut [T] {
self.as_mut_slice()
}
}
impl<T> AsRef<[T]> for ScratchVec<T> {
#[inline]
fn as_ref(&self) -> &[T] {
self.as_slice()
}
}
impl<T> Borrow<[T]> for ScratchVec<T> {
#[inline]
fn borrow(&self) -> &[T] {
self.as_slice()
}
}
impl<T> BorrowMut<[T]> for ScratchVec<T> {
#[inline]
fn borrow_mut(&mut self) -> &mut [T] {
self.as_mut_slice()
}
}
impl<T: fmt::Debug> fmt::Debug for ScratchVec<T> {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.as_slice().fmt(f)
}
}
impl<T> ops::Deref for ScratchVec<T> {
type Target = [T];
#[inline]
fn deref(&self) -> &Self::Target {
self.as_slice()
}
}
impl<T> ops::DerefMut for ScratchVec<T> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
self.as_mut_slice()
}
}
impl<T, I: slice::SliceIndex<[T]>> ops::Index<I> for ScratchVec<T> {
type Output = <I as slice::SliceIndex<[T]>>::Output;
#[inline]
fn index(&self, index: I) -> &Self::Output {
&self.as_slice()[index]
}
}
impl<T, I: slice::SliceIndex<[T]>> ops::IndexMut<I> for ScratchVec<T> {
#[inline]
fn index_mut(&mut self, index: I) -> &mut Self::Output {
&mut self.as_mut_slice()[index]
}
}
pub struct Drain<'a, T: 'a> {
tail_start: usize,
tail_len: usize,
iter: slice::Iter<'a, T>,
vec: NonNull<ScratchVec<T>>,
}
impl<T: fmt::Debug> fmt::Debug for Drain<'_, T> {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("Drain").field(&self.iter.as_slice()).finish()
}
}
impl<T> Drain<'_, T> {
#[inline]
pub fn as_slice(&self) -> &[T] {
self.iter.as_slice()
}
}
impl<T> AsRef<[T]> for Drain<'_, T> {
fn as_ref(&self) -> &[T] {
self.as_slice()
}
}
impl<T> Iterator for Drain<'_, T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
self.iter
.next()
.map(|elt| unsafe { core::ptr::read(elt as *const _) })
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
}
impl<T> DoubleEndedIterator for Drain<'_, T> {
#[inline]
fn next_back(&mut self) -> Option<T> {
self.iter
.next_back()
.map(|elt| unsafe { core::ptr::read(elt as *const _) })
}
}
impl<T> Drop for Drain<'_, T> {
fn drop(&mut self) {
struct DropGuard<'r, 'a, T>(&'r mut Drain<'a, T>);
impl<'r, 'a, T> Drop for DropGuard<'r, 'a, T> {
fn drop(&mut self) {
self.0.for_each(drop);
if self.0.tail_len > 0 {
unsafe {
let source_vec = self.0.vec.as_mut();
let start = source_vec.len();
let tail = self.0.tail_start;
if tail != start {
let src = source_vec.as_ptr().add(tail);
let dst = source_vec.as_mut_ptr().add(start);
core::ptr::copy(src, dst, self.0.tail_len);
}
source_vec.set_len(start + self.0.tail_len);
}
}
}
}
while let Some(item) = self.next() {
let guard = DropGuard(self);
drop(item);
core::mem::forget(guard);
}
DropGuard(self);
}
}
impl<T> ExactSizeIterator for Drain<'_, T> {}
impl<T> core::iter::FusedIterator for Drain<'_, T> {}