Injector

// Bits indicating the state of a slot:
// * If a task has been written into the slot, `WRITE` is set.
// * If a task has been read from the slot, `READ` is set.
// * If the block is being destroyed, `DESTROY` is set.
const WRITE: usize = 1;
const READ: usize = 2;
const DESTROY: usize = 4;

// Each block covers one "lap" of indices.
const LAP: Index = 64;
// The maximum number of values a block can hold.
const BLOCK_CAP: usize = LAP as usize - 1;
// How many lower bits are reserved for metadata.
const SHIFT: usize = 1;
// Indicates that the block is not the last one.
const HAS_NEXT: Index = 1;

struct Slot<T> {
    task: UnsafeCell<MaybeUninit<T>>,
    state: AtomicUsize,
}

struct Block<T> {
    next: AtomicPtr<Block<T>>,
    slots: [Slot<T>; BLOCK_CAP],
}

struct Position<T> {
    index: AtomicIndex,
    block: AtomicPtr<Block<T>>,
}

pub struct Injector<T> {
    head: CachePadded<Position<T>>,
    tail: CachePadded<Position<T>>,
    _marker: PhantomData<T>,
}

pub fn push(&self, task: T) {
    let backoff = Backoff::new();
    let mut tail = self.tail.index.load(Ordering::Acquire);
    let mut block = self.tail.block.load(Ordering::Acquire);
    let mut next_block = None;

    loop {
        // Calculate the offset of the index into the block.
        let offset = ((tail >> SHIFT) % LAP) as usize;

        // If we reached the end of the block, wait until the next one is installed.
        if offset == BLOCK_CAP {
            backoff.snooze();
            tail = self.tail.index.load(Ordering::Acquire);
            block = self.tail.block.load(Ordering::Acquire);
            continue;
        }

        // If we're going to have to install the next block, allocate it in advance in order to
        // make the wait for other threads as short as possible.
        if offset + 1 == BLOCK_CAP && next_block.is_none() {
            next_block = Some(Block::<T>::new());
        }

        let new_tail = tail + (1 << SHIFT);

        // Try advancing the tail forward.
        match self.tail.index.compare_exchange_weak(
            tail,
            new_tail,
            Ordering::SeqCst,
            Ordering::Acquire,
        ) {
            Ok(_) => unsafe {
                // If we've reached the end of the block, install the next one.
                if offset + 1 == BLOCK_CAP {
                    let next_block = Box::into_raw(next_block.unwrap());
                    let next_index = new_tail.wrapping_add(1 << SHIFT);

                    self.tail.block.store(next_block, Ordering::Release);
                    self.tail.index.store(next_index, Ordering::Release);
                    (*block).next.store(next_block, Ordering::Release);
                }

                // Write the task into the slot.
                let slot = (*block).slots.get_unchecked(offset);
                slot.task.get().write(MaybeUninit::new(task));
                slot.state.fetch_or(WRITE, Ordering::Release);

                return;
            },
            Err(t) => {
                tail = t;
                block = self.tail.block.load(Ordering::Acquire);
                backoff.spin();
            }
        }
    }
}

let mut new_head = head + (1 << SHIFT);

if new_head & HAS_NEXT == 0 {
    atomic::fence(Ordering::SeqCst);
    let tail = self.tail.index.load(Ordering::Relaxed);

    // If the tail equals the head, that means the queue is empty.
    if head >> SHIFT == tail >> SHIFT {
        return Steal::Empty;
    }

    // If head and tail are not in the same block, set `HAS_NEXT` in head.
    if (head >> SHIFT) / LAP != (tail >> SHIFT) / LAP {
        new_head |= HAS_NEXT;
    }
}

// Try moving the head index forward.
if self
    .head
    .index
    .compare_exchange_weak(head, new_head, Ordering::SeqCst, Ordering::Acquire)
    .is_err()
{
    return Steal::Retry;
}

unsafe fn destroy(this: *mut Block<T>, count: usize) {
    // It is not necessary to set the `DESTROY` bit in the last slot because that slot has
    // begun destruction of the block.
    for i in (0..count).rev() {
        let slot = (*this).slots.get_unchecked(i);

        // Mark the `DESTROY` bit if a thread is still using the slot.
        if slot.state.load(Ordering::Acquire) & READ == 0
            && slot.state.fetch_or(DESTROY, Ordering::AcqRel) & READ == 0
        {
            // If a thread is still using the slot, it will continue destruction of the block.
            return;
        }
    }

    // No thread is using the block, now it is safe to destroy it.
    drop(Box::from_raw(this));
}

pub fn steal_batch_and_pop(&self, dest: &Worker<T>) -> Steal<T> {
    // TODO: we use `MAX_BATCH + 1` as the hard limit for Injecter as the performance is slightly
    // better, but we may change it in the future to be compatible with the same method in Stealer.
    self.steal_batch_with_limit_and_pop(dest, MAX_BATCH + 1)
}