use alloc::boxed::Box;
use alloc::sync::Arc;
use core::iter::FusedIterator;
use core::mem::{drop, MaybeUninit};
use core::panic::{RefUnwindSafe, UnwindSafe};
use core::sync::atomic::Ordering::{Acquire, Relaxed, Release};
use cache_padded::CachePadded;
use crate::config::{AtomicUnsignedLong, AtomicUnsignedShort, UnsignedShort};
use crate::loom_exports::cell::UnsafeCell;
use crate::loom_exports::debug_or_loom_assert;
use crate::{allocate_buffer, pack, unpack, StealError};
#[derive(Debug)]
struct Queue<T> {
push_count: CachePadded<AtomicUnsignedShort>,
pop_count_and_head: CachePadded<AtomicUnsignedLong>,
head: CachePadded<AtomicUnsignedShort>,
buffer: Box<[UnsafeCell<MaybeUninit<T>>]>,
mask: UnsignedShort,
}
impl<T> Queue<T> {
#[inline]
unsafe fn read_at(&self, position: UnsignedShort) -> T {
let index = (position & self.mask) as usize;
(*self.buffer).as_ref()[index].with(|slot| slot.read().assume_init())
}
#[inline]
unsafe fn write_at(&self, position: UnsignedShort, item: T) {
let index = (position & self.mask) as usize;
(*self.buffer).as_ref()[index].with_mut(|slot| slot.write(MaybeUninit::new(item)));
}
fn book_items<C>(
&self,
mut count_fn: C,
max_count: UnsignedShort,
) -> Result<(UnsignedShort, UnsignedShort, UnsignedShort), StealError>
where
C: FnMut(usize) -> usize,
{
let mut pop_count_and_head = self.pop_count_and_head.load(Acquire);
let old_head = self.head.load(Acquire);
loop {
let (pop_count, head) = unpack(pop_count_and_head);
if old_head != head {
return Err(StealError::Busy);
}
let push_count = self.push_count.load(Acquire);
let tail = push_count.wrapping_sub(pop_count);
let item_count = tail.wrapping_sub(head);
if item_count == 0 {
return Err(StealError::Empty);
}
let count = (count_fn(item_count as usize).min(max_count as usize) as UnsignedShort)
.min(item_count);
if count == 0 {
return Err(StealError::Empty);
}
let new_head = head.wrapping_add(count);
let new_pop_count_and_head = pack(pop_count, new_head);
match self.pop_count_and_head.compare_exchange_weak(
pop_count_and_head,
new_pop_count_and_head,
Acquire,
Acquire,
) {
Ok(_) => return Ok((head, new_head, count)),
Err(current) => pop_count_and_head = current,
}
}
}
#[inline]
fn capacity(&self) -> UnsignedShort {
self.mask.wrapping_add(1)
}
}
impl<T> Drop for Queue<T> {
fn drop(&mut self) {
let head = self.head.load(Relaxed);
let push_count = self.push_count.load(Relaxed);
let pop_count = unpack(self.pop_count_and_head.load(Relaxed)).0;
let tail = push_count.wrapping_sub(pop_count);
let count = tail.wrapping_sub(head);
for offset in 0..count {
drop(unsafe { self.read_at(head.wrapping_add(offset)) });
}
}
}
#[derive(Debug)]
pub struct Worker<T> {
queue: Arc<Queue<T>>,
}
impl<T> Worker<T> {
pub fn new(min_capacity: usize) -> Self {
const MAX_CAPACITY: usize = 1 << (UnsignedShort::BITS - 1);
assert!(
min_capacity <= MAX_CAPACITY,
"the capacity of the queue cannot exceed {}",
MAX_CAPACITY
);
let capacity = min_capacity.next_power_of_two();
let buffer = allocate_buffer(capacity);
let mask = capacity as UnsignedShort - 1;
let queue = Arc::new(Queue {
push_count: CachePadded::new(AtomicUnsignedShort::new(0)),
pop_count_and_head: CachePadded::new(AtomicUnsignedLong::new(0)),
head: CachePadded::new(AtomicUnsignedShort::new(0)),
buffer,
mask,
});
Worker { queue }
}
pub fn stealer(&self) -> Stealer<T> {
Stealer {
queue: self.queue.clone(),
}
}
pub fn capacity(&self) -> usize {
self.queue.capacity() as usize
}
pub fn spare_capacity(&self) -> usize {
let push_count = self.queue.push_count.load(Relaxed);
let pop_count = unpack(self.queue.pop_count_and_head.load(Relaxed)).0;
let tail = push_count.wrapping_sub(pop_count);
let head = self.queue.head.load(Relaxed);
let len = tail.wrapping_sub(head);
(self.queue.capacity() - len) as usize
}
pub fn is_empty(&self) -> bool {
let push_count = self.queue.push_count.load(Relaxed);
let (pop_count, head) = unpack(self.queue.pop_count_and_head.load(Relaxed));
let tail = push_count.wrapping_sub(pop_count);
tail == head
}
pub fn push(&self, item: T) -> Result<(), T> {
let push_count = self.queue.push_count.load(Relaxed);
let pop_count = unpack(self.queue.pop_count_and_head.load(Relaxed)).0;
let tail = push_count.wrapping_sub(pop_count);
let head = self.queue.head.load(Acquire);
if tail.wrapping_sub(head) > self.queue.mask {
return Err(item);
}
unsafe { self.queue.write_at(tail, item) };
self.queue
.push_count
.store(push_count.wrapping_add(1), Release);
Ok(())
}
pub fn extend<I: IntoIterator<Item = T>>(&self, iter: I) {
let push_count = self.queue.push_count.load(Relaxed);
let pop_count = unpack(self.queue.pop_count_and_head.load(Relaxed)).0;
let mut tail = push_count.wrapping_sub(pop_count);
let head = self.queue.head.load(Acquire);
let max_tail = head.wrapping_add(self.queue.capacity());
for item in iter {
if tail == max_tail {
break;
}
unsafe { self.queue.write_at(tail, item) };
tail = tail.wrapping_add(1);
}
self.queue
.push_count
.store(tail.wrapping_add(pop_count), Release);
}
pub fn pop(&self) -> Option<T> {
let mut pop_count_and_head = self.queue.pop_count_and_head.load(Relaxed);
let push_count = self.queue.push_count.load(Relaxed);
let (pop_count, mut head) = unpack(pop_count_and_head);
let tail = push_count.wrapping_sub(pop_count);
let new_pop_count = pop_count.wrapping_add(1);
loop {
if tail == head {
return None;
}
let new_pop_count_and_head = pack(new_pop_count, head);
match self.queue.pop_count_and_head.compare_exchange_weak(
pop_count_and_head,
new_pop_count_and_head,
Release,
Relaxed,
) {
Ok(_) => break,
Err(current) => {
pop_count_and_head = current;
head = unpack(current).1;
}
}
}
unsafe { Some(self.queue.read_at(tail.wrapping_sub(1))) }
}
pub fn drain<C>(&self, count_fn: C) -> Result<Drain<'_, T>, StealError>
where
C: FnMut(usize) -> usize,
{
let (old_head, new_head, _) = self.queue.book_items(count_fn, UnsignedShort::MAX)?;
Ok(Drain {
queue: &self.queue,
current: old_head,
end: new_head,
})
}
}
impl<T> UnwindSafe for Worker<T> {}
impl<T> RefUnwindSafe for Worker<T> {}
unsafe impl<T: Send> Send for Worker<T> {}
#[derive(Debug)]
pub struct Drain<'a, T> {
queue: &'a Queue<T>,
current: UnsignedShort,
end: UnsignedShort,
}
impl<'a, T> Iterator for Drain<'a, T> {
type Item = T;
fn next(&mut self) -> Option<T> {
if self.current == self.end {
return None;
}
let item = Some(unsafe { self.queue.read_at(self.current) });
self.current = self.current.wrapping_add(1);
if self.current == self.end {
self.queue.head.store(self.end, Release);
}
item
}
fn size_hint(&self) -> (usize, Option<usize>) {
let sz = self.end.wrapping_sub(self.current) as usize;
(sz, Some(sz))
}
}
impl<'a, T> ExactSizeIterator for Drain<'a, T> {}
impl<'a, T> FusedIterator for Drain<'a, T> {}
impl<'a, T> Drop for Drain<'a, T> {
fn drop(&mut self) {
for _item in self {}
}
}
impl<'a, T> UnwindSafe for Drain<'a, T> {}
impl<'a, T> RefUnwindSafe for Drain<'a, T> {}
unsafe impl<'a, T: Send> Send for Drain<'a, T> {}
unsafe impl<'a, T: Send> Sync for Drain<'a, T> {}
#[derive(Debug)]
pub struct Stealer<T> {
queue: Arc<Queue<T>>,
}
impl<T> Stealer<T> {
pub fn steal<C>(&self, dest: &Worker<T>, count_fn: C) -> Result<usize, StealError>
where
C: FnMut(usize) -> usize,
{
let dest_push_count = dest.queue.push_count.load(Relaxed);
let dest_pop_count = unpack(dest.queue.pop_count_and_head.load(Relaxed)).0;
let dest_tail = dest_push_count.wrapping_sub(dest_pop_count);
let dest_head = dest.queue.head.load(Acquire);
let dest_free_capacity = dest.queue.capacity() - dest_tail.wrapping_sub(dest_head);
debug_or_loom_assert!(dest_free_capacity <= dest.queue.capacity());
let (old_head, new_head, transfer_count) =
self.queue.book_items(count_fn, dest_free_capacity)?;
debug_or_loom_assert!(transfer_count <= dest_free_capacity);
for offset in 0..transfer_count {
unsafe {
let item = self.queue.read_at(old_head.wrapping_add(offset));
dest.queue.write_at(dest_tail.wrapping_add(offset), item);
}
}
dest.queue
.push_count
.store(dest_push_count.wrapping_add(transfer_count), Release);
self.queue.head.store(new_head, Release);
Ok(transfer_count as usize)
}
pub fn steal_and_pop<C>(&self, dest: &Worker<T>, count_fn: C) -> Result<(T, usize), StealError>
where
C: FnMut(usize) -> usize,
{
let dest_push_count = dest.queue.push_count.load(Relaxed);
let dest_pop_count = unpack(dest.queue.pop_count_and_head.load(Relaxed)).0;
let dest_tail = dest_push_count.wrapping_sub(dest_pop_count);
let dest_head = dest.queue.head.load(Acquire);
let dest_free_capacity = dest.queue.capacity() - dest_tail.wrapping_sub(dest_head);
debug_or_loom_assert!(dest_free_capacity <= dest.queue.capacity());
let (old_head, new_head, count) =
self.queue.book_items(count_fn, dest_free_capacity + 1)?;
let transfer_count = count - 1;
debug_or_loom_assert!(transfer_count <= dest_free_capacity);
for offset in 0..transfer_count {
unsafe {
let item = self.queue.read_at(old_head.wrapping_add(offset));
dest.queue.write_at(dest_tail.wrapping_add(offset), item);
}
}
let last_item = unsafe { self.queue.read_at(old_head.wrapping_add(transfer_count)) };
dest.queue
.push_count
.store(dest_push_count.wrapping_add(transfer_count), Release);
self.queue.head.store(new_head, Release);
Ok((last_item, transfer_count as usize))
}
}
impl<T> Clone for Stealer<T> {
fn clone(&self) -> Self {
Stealer {
queue: self.queue.clone(),
}
}
}
impl<T> UnwindSafe for Stealer<T> {}
impl<T> RefUnwindSafe for Stealer<T> {}
unsafe impl<T: Send> Send for Stealer<T> {}
unsafe impl<T: Send> Sync for Stealer<T> {}