use core::{num::NonZeroUsize, ptr::NonNull};
use crate::{
alloc,
atomic::{AtomicUsize, Ordering},
padded::Padded,
spsc,
};
#[repr(C)]
pub(crate) struct Shards<T> {
rc: Padded<AtomicUsize>,
queue_ptrs: spsc::QueuePtr<T>,
}
impl<T> Shards<T> {
pub fn new(max_shards: NonZeroUsize) -> NonNull<Self> {
let layout = Self::layout(max_shards.get());
let Some(ptr) = NonNull::new(unsafe { alloc::alloc(layout) }) else {
alloc::handle_alloc_error(layout)
};
unsafe { ptr.cast::<AtomicUsize>().write(AtomicUsize::new(1)) };
ptr.cast()
}
fn layout(max_shards: usize) -> alloc::Layout {
let (layout, _offset) = alloc::Layout::new::<Padded<AtomicUsize>>()
.extend(alloc::Layout::array::<spsc::QueuePtr<T>>(max_shards).unwrap())
.unwrap();
layout.pad_to_align()
}
#[inline(always)]
pub(crate) fn at(ptr: NonNull<Self>, shard: usize) -> NonNull<spsc::QueuePtr<T>> {
unsafe { _field!(Shards<T>, ptr, queue_ptrs).cast().add(shard) }
}
}
pub(crate) struct ShardsPtr<T> {
ptr: NonNull<Shards<T>>,
max_shards: usize,
}
impl<T> Clone for ShardsPtr<T> {
fn clone(&self) -> Self {
self.rc().fetch_add(1, Ordering::AcqRel);
Self {
ptr: self.ptr,
max_shards: self.max_shards,
}
}
}
impl<T> ShardsPtr<T> {
pub fn new(max_shards: NonZeroUsize, capacity_per_shard: NonZeroUsize) -> Self {
let ptr = Shards::new(max_shards);
for i in 0..max_shards.get() {
let ptr = Shards::at(ptr, i);
unsafe { ptr.write(spsc::QueuePtr::with_size(capacity_per_shard)) };
}
Self {
ptr,
max_shards: max_shards.get(),
}
}
pub(crate) fn clone_queue_ptr(&self, shard: usize) -> spsc::QueuePtr<T> {
unsafe { Shards::at(self.ptr, shard).as_ref() }.clone()
}
fn rc(&self) -> &AtomicUsize {
unsafe { _field!(Shards<T>, self.ptr, rc, AtomicUsize).as_ref() }
}
}
impl<T> Drop for ShardsPtr<T> {
fn drop(&mut self) {
if self.rc().fetch_sub(1, Ordering::AcqRel) == 1 {
unsafe {
_field!(Shards<T>, self.ptr, rc, AtomicUsize).drop_in_place();
for i in 0..self.max_shards {
Shards::at(self.ptr, i).drop_in_place();
}
alloc::dealloc(
self.ptr.as_ptr().cast(),
Shards::<T>::layout(self.max_shards),
);
}
}
}
}