use core::{alloc::Layout, marker::Destruct, ops, ptr::NonNull};
use super::{AllocError, Allocator, ConstAllocator};
#[doc(hidden)]
pub struct ComptimeVec<T: ~const Destruct> {
ptr: NonNull<T>,
len: usize,
capacity: usize,
allocator: ConstAllocator,
}
impl<T: ~const Clone + ~const Destruct> const Clone for ComptimeVec<T> {
fn clone(&self) -> Self {
#[inline]
const fn clone_shim<T: ~const Clone>(x: &T) -> T {
x.clone()
}
self.map(clone_shim)
}
}
impl<T: ~const Destruct> const Drop for ComptimeVec<T> {
fn drop(&mut self) {
if core::mem::needs_drop::<T>() {
while self.pop().is_some() {}
}
unsafe {
self.allocator
.deallocate(self.ptr.cast(), layout_array::<T>(self.capacity));
}
}
}
impl<T: ~const Destruct> ComptimeVec<T> {
pub const fn new_in(allocator: ConstAllocator) -> Self {
Self::with_capacity_in(0, allocator)
}
pub const fn with_capacity_in(capacity: usize, allocator: ConstAllocator) -> Self {
Self {
ptr: unwrap_alloc_error(allocator.allocate(layout_array::<T>(capacity))).cast(),
len: 0,
capacity,
allocator,
}
}
pub const fn repeat_in(allocator: ConstAllocator, x: T, n: usize) -> Self
where
T: Copy,
{
let mut this = Self::with_capacity_in(n, allocator);
while this.len() < n {
this.push(x);
}
this
}
pub const fn allocator(&self) -> &ConstAllocator {
&self.allocator
}
pub const fn push(&mut self, x: T) {
unsafe {
self.reserve(1);
self.ptr.as_ptr().wrapping_add(self.len).write(x)
}
self.len += 1;
}
pub const fn pop(&mut self) -> Option<T> {
unsafe {
if let Some(i) = self.len.checked_sub(1) {
self.len = i;
Some(self.ptr.as_ptr().wrapping_add(i).read())
} else {
None
}
}
}
const fn reserve(&mut self, additional: usize) {
if self.capacity - self.len >= additional {
return;
}
let mut new_cap = self.capacity.checked_add(2).expect("capacity overflow");
while new_cap - self.len < additional {
new_cap = new_cap.checked_mul(2).expect("capacity overflow");
}
unsafe {
self.ptr = unwrap_alloc_error(self.allocator.grow(
self.ptr.cast(),
layout_array::<T>(self.capacity),
layout_array::<T>(new_cap),
))
.cast();
self.capacity = new_cap;
}
}
pub const fn map<F: ~const FnMut(&T) -> U + ~const Destruct, U: ~const Destruct>(
&self,
mut f: F,
) -> ComptimeVec<U> {
let mut out = ComptimeVec::with_capacity_in(self.len, self.allocator.clone());
let mut i = 0;
while i < self.len() {
out.push(f(&self[i]));
i += 1;
}
out
}
pub const fn clear(&mut self)
where
T: ~const Destruct,
{
if core::mem::needs_drop::<T>() {
while self.pop().is_some() {}
} else {
self.len = 0;
}
}
#[inline]
pub const fn as_slice(&self) -> &[T] {
unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
}
#[inline]
pub const fn as_mut_slice(&mut self) -> &mut [T] {
unsafe { core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
}
pub const fn to_array<const LEN: usize>(&self) -> [T; LEN]
where
T: Copy,
{
assert!(self.len() == LEN);
unsafe { *(self.ptr.as_ptr() as *const [T; LEN]) }
}
}
impl<T: ~const Destruct> const ops::Deref for ComptimeVec<T> {
type Target = [T];
fn deref(&self) -> &Self::Target {
self.as_slice()
}
}
impl<T: ~const Destruct> const ops::DerefMut for ComptimeVec<T> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.as_mut_slice()
}
}
#[allow(unused_macros)]
macro_rules! vec_position {
($vec:expr, |$item:ident| $predicate:expr) => {{
let mut i = 0;
loop {
if i >= $vec.len() {
break None;
}
let $item = &$vec[i];
if $predicate {
break Some(i);
}
i += 1;
}
}};
}
const fn unwrap_alloc_error<T: ~const Destruct>(x: Result<T, AllocError>) -> T {
match x {
Ok(x) => x,
Err(AllocError) => panic!("compile-time heap allocation failed"),
}
}
const fn layout_array<T>(len: usize) -> Layout {
let singular = Layout::new::<T>();
let Some(size) = singular.size().checked_mul(len) else { panic!("size overflow") };
let Ok(layout) = Layout::from_size_align(size, singular.align()) else { unreachable!() };
layout
}
#[cfg(test)]
mod tests {
use super::*;
use quickcheck::TestResult;
use quickcheck_macros::quickcheck;
#[test]
fn as_slice() {
#[allow(dead_code)] const fn array(allocator: &ConstAllocator) {
let mut x = ComptimeVec::new_in(allocator.clone());
x.push(1);
x.push(2);
x.push(3);
x.push(4);
assert!(matches!(x.pop(), Some(4)));
let slice = x.as_slice();
assert!(matches!(slice, [1, 2, 3]));
}
const _: () = ConstAllocator::with(array);
}
#[test]
fn map() {
const fn array(allocator: &ConstAllocator) -> [i32; 3] {
let mut x = ComptimeVec::new_in(allocator.clone());
x.push(1);
x.push(2);
x.push(3);
const fn transform(x: &i32) -> i32 {
*x + 1
}
let y = x.map(transform);
y.to_array()
}
const OUT: [i32; 3] = ConstAllocator::with(array);
assert_eq!(OUT, [2, 3, 4]);
}
#[test]
fn to_array() {
const fn array(allocator: &ConstAllocator) -> [u32; 3] {
let mut v = ComptimeVec::new_in(allocator.clone());
v.push(1);
v.push(2);
v.push(3);
v.to_array()
}
const OUT: [u32; 3] = ConstAllocator::with(array);
assert_eq!(OUT, [1, 2, 3]);
}
#[test]
fn get_mut() {
const fn val(allocator: &ConstAllocator) -> u32 {
let mut v = ComptimeVec::new_in(allocator.clone());
v.push(1);
v.push(2);
v.push(3);
v[1] += 2;
v[1]
}
const OUT: u32 = ConstAllocator::with(val);
assert_eq!(OUT, 4);
}
#[test]
fn const_vec_position() {
const fn pos(allocator: &ConstAllocator) -> [Option<usize>; 2] {
let mut v = ComptimeVec::new_in(allocator.clone());
v.push(42);
v.push(43);
v.push(44);
[
vec_position!(v, |i| *i == 43),
vec_position!(v, |i| *i == 50),
]
}
const OUT: [Option<usize>; 2] = ConstAllocator::with(pos);
assert_eq!(OUT, [Some(1), None]);
}
#[test]
fn drop_on_clear() {
#[allow(dead_code)] const fn array(allocator: &ConstAllocator) {
let mut x = ComptimeVec::new_in(allocator.clone());
x.push(allocator.clone());
x.push(allocator.clone());
x.push(allocator.clone());
}
const _: () = ConstAllocator::with(array);
}
#[quickcheck]
fn vec_position(values: Vec<u8>, expected_index: usize) -> TestResult {
let needle = if values.is_empty() {
42
} else {
values[expected_index % values.len()]
};
let got = vec_position!(values, |i| *i == needle);
let expected = values.iter().position(|i| *i == needle);
assert_eq!(got, expected);
TestResult::passed()
}
}