use std::alloc::Layout;
use std::cell::UnsafeCell;
use std::fmt::Debug;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, Ordering};
use std::{ptr, slice};
use crate::{Item, Utf32String};
const BUCKETS: u32 = u32::BITS - SKIP_BUCKET;
const MAX_ENTRIES: u32 = u32::MAX - SKIP;
pub(crate) struct Vec<T> {
inflight: AtomicU64,
buckets: [Bucket<T>; BUCKETS as usize],
columns: u32,
}
impl<T> Vec<T> {
pub fn with_capacity(capacity: u32, columns: u32) -> Vec<T> {
assert_ne!(columns, 0, "there must be atleast one matcher column");
let init = match capacity {
0 => 0,
n => Location::of(n).bucket,
};
let mut buckets = [ptr::null_mut(); BUCKETS as usize];
for (i, bucket) in buckets[..=init as usize].iter_mut().enumerate() {
let len = Location::bucket_len(i as u32);
*bucket = unsafe { Bucket::alloc(len, columns) };
}
Vec {
buckets: buckets.map(Bucket::new),
inflight: AtomicU64::new(0),
columns,
}
}
pub fn columns(&self) -> u32 {
self.columns
}
#[inline]
pub fn count(&self) -> u32 {
self.inflight
.load(Ordering::Acquire)
.min(MAX_ENTRIES as u64) as u32
}
#[inline]
pub unsafe fn get_unchecked(&self, index: u32) -> Item<'_, T> {
let location = Location::of(index);
unsafe {
let entries = self
.buckets
.get_unchecked(location.bucket as usize)
.entries
.load(Ordering::Relaxed);
debug_assert!(!entries.is_null());
let entry = Bucket::<T>::get(entries, location.entry, self.columns);
let _ = (*entry).active.load(Ordering::Acquire);
Entry::read(entry, self.columns)
}
}
pub fn get(&self, index: u32) -> Option<Item<'_, T>> {
let location = Location::of(index);
unsafe {
let entries = self
.buckets
.get_unchecked(location.bucket as usize)
.entries
.load(Ordering::Relaxed);
if entries.is_null() {
return None;
}
let entry = Bucket::<T>::get(entries, location.entry, self.columns);
(*entry)
.active
.load(Ordering::Acquire)
.then(|| Entry::read(entry, self.columns))
}
}
pub fn push(&self, value: T, fill_columns: impl FnOnce(&T, &mut [Utf32String])) -> u32 {
let index = self.inflight.fetch_add(1, Ordering::Release);
let index: u32 = index.try_into().expect("overflowed maximum capacity");
let location = Location::of(index);
if index == (location.bucket_len - (location.bucket_len >> 3)) {
if let Some(next_bucket) = self.buckets.get(location.bucket as usize + 1) {
Vec::get_or_alloc(next_bucket, location.bucket_len << 1, self.columns);
}
}
let bucket = unsafe { self.buckets.get_unchecked(location.bucket as usize) };
let mut entries = bucket.entries.load(Ordering::Acquire);
if entries.is_null() {
entries = Vec::get_or_alloc(bucket, location.bucket_len, self.columns);
}
unsafe {
let entry = Bucket::get(entries, location.entry, self.columns);
for col in Entry::matcher_cols_raw(entry, self.columns) {
col.get().write(MaybeUninit::new(Utf32String::default()))
}
fill_columns(&value, Entry::matcher_cols_mut(entry, self.columns));
(*entry).slot.get().write(MaybeUninit::new(value));
(*entry).active.store(true, Ordering::Release);
}
index
}
pub fn extend<I>(&self, values: I, fill_columns: impl Fn(&T, &mut [Utf32String]))
where
I: IntoIterator<Item = T> + ExactSizeIterator,
{
let count: u32 = values
.len()
.try_into()
.expect("overflowed maximum capacity");
if count == 0 {
assert!(
values.into_iter().next().is_none(),
"The `values` variable reported incorrect length."
);
return;
}
let start_index: u32 = self
.inflight
.fetch_add(u64::from(count), Ordering::Release)
.try_into()
.expect("overflowed maximum capacity");
let start_location = Location::of(start_index);
let end_location = Location::of(start_index + count);
let alloc_entry = end_location.alloc_next_bucket_entry();
if end_location.entry >= alloc_entry
&& (start_location.bucket != end_location.bucket || start_location.entry <= alloc_entry)
{
if let Some(next_bucket) = self.buckets.get(end_location.bucket as usize + 1) {
Vec::get_or_alloc(next_bucket, end_location.bucket_len << 1, self.columns);
}
}
let mut bucket = unsafe { self.buckets.get_unchecked(start_location.bucket as usize) };
let mut entries = bucket.entries.load(Ordering::Acquire);
if entries.is_null() {
entries = Vec::get_or_alloc(
bucket,
Location::bucket_len(start_location.bucket),
self.columns,
);
}
let mut location;
let count = count as usize;
for (i, v) in values.into_iter().enumerate() {
assert!(i < count);
location =
Location::of(start_index + u32::try_from(i).expect("overflowed maximum capacity"));
if location.entry == 0 && i != 0 {
bucket = unsafe { self.buckets.get_unchecked(location.bucket as usize) };
entries = bucket.entries.load(Ordering::Acquire);
if entries.is_null() {
entries = Vec::get_or_alloc(
bucket,
Location::bucket_len(location.bucket),
self.columns,
);
}
}
unsafe {
let entry = Bucket::get(entries, location.entry, self.columns);
for col in Entry::matcher_cols_raw(entry, self.columns) {
col.get().write(MaybeUninit::new(Utf32String::default()));
}
fill_columns(&v, Entry::matcher_cols_mut(entry, self.columns));
(*entry).slot.get().write(MaybeUninit::new(v));
(*entry).active.store(true, Ordering::Release);
}
}
}
fn get_or_alloc(bucket: &Bucket<T>, len: u32, cols: u32) -> *mut Entry<T> {
let entries = unsafe { Bucket::alloc(len, cols) };
match bucket.entries.compare_exchange(
ptr::null_mut(),
entries,
Ordering::Release,
Ordering::Acquire,
) {
Ok(_) => entries,
Err(found) => unsafe {
Bucket::dealloc(entries, len, cols);
found
},
}
}
pub unsafe fn snapshot(&self, start: u32) -> Iter<'_, T> {
let end = self
.inflight
.load(Ordering::Acquire)
.min(MAX_ENTRIES as u64) as u32;
assert!(start <= end, "index {start} is out of bounds!");
Iter {
location: Location::of(start),
vec: self,
idx: start,
end,
}
}
pub unsafe fn par_snapshot(&self, start: u32) -> ParIter<'_, T> {
let end = self
.inflight
.load(Ordering::Acquire)
.min(MAX_ENTRIES as u64) as u32;
assert!(start <= end, "index {start} is out of bounds!");
ParIter {
start,
end,
vec: self,
}
}
}
impl<T> Drop for Vec<T> {
fn drop(&mut self) {
for (i, bucket) in self.buckets.iter_mut().enumerate() {
let entries = *bucket.entries.get_mut();
if entries.is_null() {
break;
}
let len = Location::bucket_len(i as u32);
unsafe { Bucket::dealloc(entries, len, self.columns) }
}
}
}
type SnapshotItem<'v, T> = (u32, Option<Item<'v, T>>);
pub struct Iter<'v, T> {
location: Location,
idx: u32,
end: u32,
vec: &'v Vec<T>,
}
impl<T> Iter<'_, T> {
pub fn end(&self) -> u32 {
self.end
}
}
impl<'v, T> Iterator for Iter<'v, T> {
type Item = SnapshotItem<'v, T>;
fn size_hint(&self) -> (usize, Option<usize>) {
(
(self.end - self.idx) as usize,
Some((self.end - self.idx) as usize),
)
}
fn next(&mut self) -> Option<SnapshotItem<'v, T>> {
if self.end == self.idx {
return None;
}
debug_assert!(self.idx < self.end, "huh {} {}", self.idx, self.end);
debug_assert!(self.end as u64 <= self.vec.inflight.load(Ordering::Relaxed));
loop {
let entries = unsafe {
self.vec
.buckets
.get_unchecked(self.location.bucket as usize)
.entries
.load(Ordering::Relaxed)
};
debug_assert!(self.location.bucket < BUCKETS);
if self.location.entry < self.location.bucket_len {
if entries.is_null() {
let index = self.idx;
self.location.entry += 1;
self.idx += 1;
return Some((index, None));
}
let entry = unsafe { Bucket::get(entries, self.location.entry, self.vec.columns) };
let index = self.idx;
self.location.entry += 1;
self.idx += 1;
let entry = unsafe {
(*entry)
.active
.load(Ordering::Acquire)
.then(|| Entry::read(entry, self.vec.columns))
};
return Some((index, entry));
}
self.location.entry = 0;
self.location.bucket += 1;
if self.location.bucket < BUCKETS {
self.location.bucket_len = Location::bucket_len(self.location.bucket);
}
}
}
}
impl<T> ExactSizeIterator for Iter<'_, T> {}
impl<T> DoubleEndedIterator for Iter<'_, T> {
fn next_back(&mut self) -> Option<Self::Item> {
unimplemented!()
}
}
pub struct ParIter<'v, T> {
end: u32,
start: u32,
vec: &'v Vec<T>,
}
impl<T> ParIter<'_, T> {
pub fn end(&self) -> u32 {
self.end
}
}
impl<'v, T: Send + Sync> rayon::iter::ParallelIterator for ParIter<'v, T> {
type Item = SnapshotItem<'v, T>;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: rayon::iter::plumbing::UnindexedConsumer<Self::Item>,
{
rayon::iter::plumbing::bridge(self, consumer)
}
fn opt_len(&self) -> Option<usize> {
Some((self.end - self.start) as usize)
}
}
impl<T: Send + Sync> rayon::iter::IndexedParallelIterator for ParIter<'_, T> {
fn len(&self) -> usize {
(self.end - self.start) as usize
}
fn drive<C: rayon::iter::plumbing::Consumer<Self::Item>>(self, consumer: C) -> C::Result {
rayon::iter::plumbing::bridge(self, consumer)
}
fn with_producer<CB>(self, callback: CB) -> CB::Output
where
CB: rayon::iter::plumbing::ProducerCallback<Self::Item>,
{
callback.callback(ParIterProducer {
start: self.start,
end: self.end,
vec: self.vec,
})
}
}
struct ParIterProducer<'v, T: Send> {
start: u32,
end: u32,
vec: &'v Vec<T>,
}
impl<'v, T: 'v + Send + Sync> rayon::iter::plumbing::Producer for ParIterProducer<'v, T> {
type Item = SnapshotItem<'v, T>;
type IntoIter = Iter<'v, T>;
fn into_iter(self) -> Self::IntoIter {
debug_assert!(self.start <= self.end);
Iter {
location: Location::of(self.start),
idx: self.start,
end: self.end,
vec: self.vec,
}
}
fn split_at(self, index: usize) -> (Self, Self) {
assert!(index <= (self.end - self.start) as usize);
let index = index as u32;
(
ParIterProducer {
start: self.start,
end: self.start + index,
vec: self.vec,
},
ParIterProducer {
start: self.start + index,
end: self.end,
vec: self.vec,
},
)
}
}
struct Bucket<T> {
entries: AtomicPtr<Entry<T>>,
}
impl<T> Bucket<T> {
fn layout(len: u32, layout: Layout) -> Layout {
Layout::from_size_align(layout.size() * len as usize, layout.align())
.expect("exceeded maximum allocation size")
}
unsafe fn alloc(len: u32, cols: u32) -> *mut Entry<T> {
let layout = Entry::<T>::layout(cols);
let arr_layout = Self::layout(len, layout);
let entries = std::alloc::alloc(arr_layout);
if entries.is_null() {
std::alloc::handle_alloc_error(arr_layout)
}
for i in 0..len {
let active = entries.add(i as usize * layout.size()) as *mut AtomicBool;
active.write(AtomicBool::new(false))
}
entries as *mut Entry<T>
}
unsafe fn dealloc(entries: *mut Entry<T>, len: u32, cols: u32) {
let layout = Entry::<T>::layout(cols);
let arr_layout = Self::layout(len, layout);
for i in 0..len {
let entry = Bucket::get(entries, i, cols);
if *(*entry).active.get_mut() {
ptr::drop_in_place((*(*entry).slot.get()).as_mut_ptr());
for matcher_col in Entry::matcher_cols_raw(entry, cols) {
ptr::drop_in_place((*matcher_col.get()).as_mut_ptr());
}
}
}
std::alloc::dealloc(entries as *mut u8, arr_layout)
}
unsafe fn get(entries: *mut Entry<T>, idx: u32, cols: u32) -> *mut Entry<T> {
let layout = Entry::<T>::layout(cols);
let ptr = entries as *mut u8;
ptr.add(layout.size() * idx as usize) as *mut Entry<T>
}
fn new(entries: *mut Entry<T>) -> Bucket<T> {
Bucket {
entries: AtomicPtr::new(entries),
}
}
}
#[repr(C)]
struct Entry<T> {
active: AtomicBool,
slot: UnsafeCell<MaybeUninit<T>>,
tail: [UnsafeCell<MaybeUninit<Utf32String>>; 0],
}
impl<T> Entry<T> {
fn layout(cols: u32) -> Layout {
let head = Layout::new::<Self>();
let tail = Layout::array::<Utf32String>(cols as usize).expect("invalid memory layout");
head.extend(tail)
.expect("invalid memory layout")
.0
.pad_to_align()
}
unsafe fn matcher_cols_raw<'a>(
ptr: *mut Entry<T>,
cols: u32,
) -> &'a [UnsafeCell<MaybeUninit<Utf32String>>] {
let tail = std::ptr::addr_of!((*ptr).tail) as *const u8;
let offset = tail.offset_from(ptr as *mut u8) as usize;
let ptr = (ptr as *mut u8).add(offset) as *mut _;
slice::from_raw_parts(ptr, cols as usize)
}
unsafe fn matcher_cols_mut<'a>(ptr: *mut Entry<T>, cols: u32) -> &'a mut [Utf32String] {
let tail = std::ptr::addr_of!((*ptr).tail) as *const u8;
let offset = tail.offset_from(ptr as *mut u8) as usize;
let ptr = (ptr as *mut u8).add(offset) as *mut _;
slice::from_raw_parts_mut(ptr, cols as usize)
}
unsafe fn read<'a>(ptr: *mut Entry<T>, cols: u32) -> Item<'a, T> {
let data = (*(*ptr).slot.get()).assume_init_ref();
let tail = std::ptr::addr_of!((*ptr).tail) as *const u8;
let offset = tail.offset_from(ptr as *mut u8) as usize;
let ptr = (ptr as *mut u8).add(offset) as *mut _;
let matcher_columns = slice::from_raw_parts(ptr, cols as usize);
Item {
data,
matcher_columns,
}
}
}
#[derive(Debug)]
struct Location {
bucket: u32,
bucket_len: u32,
entry: u32,
}
const SKIP: u32 = 32;
const SKIP_BUCKET: u32 = (u32::BITS - SKIP.leading_zeros()) - 1;
impl Location {
fn of(index: u32) -> Location {
let skipped = index.checked_add(SKIP).expect("exceeded maximum length");
let bucket = u32::BITS - skipped.leading_zeros();
let bucket = bucket - (SKIP_BUCKET + 1);
let bucket_len = Location::bucket_len(bucket);
let entry = skipped ^ bucket_len;
Location {
bucket,
bucket_len,
entry,
}
}
fn bucket_len(bucket: u32) -> u32 {
1 << (bucket + SKIP_BUCKET)
}
fn alloc_next_bucket_entry(&self) -> u32 {
self.bucket_len - (self.bucket_len >> 3)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn location() {
assert_eq!(Location::bucket_len(0), 32);
for i in 0..32 {
let loc = Location::of(i);
assert_eq!(loc.bucket_len, 32);
assert_eq!(loc.bucket, 0);
assert_eq!(loc.entry, i);
}
assert_eq!(Location::bucket_len(1), 64);
for i in 33..96 {
let loc = Location::of(i);
assert_eq!(loc.bucket_len, 64);
assert_eq!(loc.bucket, 1);
assert_eq!(loc.entry, i - 32);
}
assert_eq!(Location::bucket_len(2), 128);
for i in 96..224 {
let loc = Location::of(i);
assert_eq!(loc.bucket_len, 128);
assert_eq!(loc.bucket, 2);
assert_eq!(loc.entry, i - 96);
}
let max = Location::of(MAX_ENTRIES);
assert_eq!(max.bucket, BUCKETS - 1);
assert_eq!(max.bucket_len, 1 << 31);
assert_eq!(max.entry, (1 << 31) - 1);
}
#[test]
fn extend_unique_bucket() {
let vec = Vec::<u32>::with_capacity(1, 1);
vec.extend(0..10, |_, _| {});
assert_eq!(vec.count(), 10);
for i in 0..10 {
assert_eq!(*vec.get(i).unwrap().data, i);
}
assert!(vec.get(10).is_none());
}
#[test]
fn extend_over_two_buckets() {
let vec = Vec::<u32>::with_capacity(1, 1);
vec.extend(0..100, |_, _| {});
assert_eq!(vec.count(), 100);
for i in 0..100 {
assert_eq!(*vec.get(i).unwrap().data, i);
}
assert!(vec.get(100).is_none());
}
#[test]
fn extend_over_more_than_two_buckets() {
let vec = Vec::<u32>::with_capacity(1, 1);
vec.extend(0..1000, |_, _| {});
assert_eq!(vec.count(), 1000);
for i in 0..1000 {
assert_eq!(*vec.get(i).unwrap().data, i);
}
assert!(vec.get(1000).is_none());
}
#[test]
fn extend_with_incorrect_reported_len_is_caught() {
struct IncorrectLenIter {
len: usize,
iter: std::ops::Range<u32>,
}
impl Iterator for IncorrectLenIter {
type Item = u32;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
}
impl ExactSizeIterator for IncorrectLenIter {
fn len(&self) -> usize {
self.len
}
}
let vec = Vec::<u32>::with_capacity(1, 1);
let iter = IncorrectLenIter {
len: 10,
iter: (0..12),
};
assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_err());
let vec = Vec::<u32>::with_capacity(1, 1);
let iter = IncorrectLenIter {
len: 12,
iter: (0..10),
};
assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_ok());
assert_eq!(vec.count(), 12);
for i in 0..10 {
assert_eq!(*vec.get(i).unwrap().data, i);
}
assert!(vec.get(10).is_none());
let vec = Vec::<u32>::with_capacity(1, 1);
let iter = IncorrectLenIter {
len: 0,
iter: (0..2),
};
assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_err());
}
#[test]
fn extend_over_max_capacity() {
let vec = Vec::<u32>::with_capacity(1, 1);
let count = MAX_ENTRIES as usize + 2;
let iter = std::iter::repeat(0).take(count);
assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_err());
}
}