use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
#[repr(align(64))]
#[derive(Default)]
struct Padded(AtomicUsize);
pub struct SpscRing<T, const CAP: usize> {
mask: usize,
head: Padded, tail: Padded, slots: Box<[Slot<T>]>,
}
struct Slot<T> {
sequence: AtomicUsize,
value: UnsafeCell<Option<T>>,
}
unsafe impl<T: Send, const CAP: usize> Send for SpscRing<T, CAP> {}
unsafe impl<T: Send, const CAP: usize> Sync for SpscRing<T, CAP> {}
impl<T, const CAP: usize> SpscRing<T, CAP> {
#[must_use]
pub fn new() -> Self {
const {
assert!(CAP.is_power_of_two(), "CAP must be a power of two");
}
let slots = (0..CAP)
.map(|i| Slot {
sequence: AtomicUsize::new(i),
value: UnsafeCell::new(None),
})
.collect::<Vec<_>>();
Self {
mask: CAP - 1,
head: Padded::default(),
tail: Padded::default(),
slots: slots.into_boxed_slice(),
}
}
#[inline]
pub fn try_push(&self, value: T) -> Option<T> {
let pos = self.head.0.load(Ordering::Relaxed);
let slot = &self.slots[pos & self.mask];
if slot.sequence.load(Ordering::Acquire) != pos {
return Some(value); }
unsafe {
*slot.value.get() = Some(value);
}
slot.sequence.store(pos.wrapping_add(1), Ordering::Release);
self.head.0.store(pos.wrapping_add(1), Ordering::Relaxed);
None
}
#[inline]
pub fn try_pop(&self) -> Option<T> {
let pos = self.tail.0.load(Ordering::Relaxed);
let slot = &self.slots[pos & self.mask];
if slot.sequence.load(Ordering::Acquire) != pos + 1 {
return None; }
let value = unsafe { (*slot.value.get()).take() };
slot.sequence
.store(pos.wrapping_add(CAP), Ordering::Release);
self.tail.0.store(pos.wrapping_add(1), Ordering::Relaxed);
value
}
#[must_use]
pub fn len(&self) -> usize {
self.head
.0
.load(Ordering::Relaxed)
.wrapping_sub(self.tail.0.load(Ordering::Relaxed))
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<T, const CAP: usize> Default for SpscRing<T, CAP> {
fn default() -> Self {
Self::new()
}
}
#[must_use]
pub fn channel<T, const CAP: usize>() -> (SpscSender<T, CAP>, SpscReceiver<T, CAP>) {
let ring = std::sync::Arc::new(SpscRing::new());
(
SpscSender {
ring: std::sync::Arc::clone(&ring),
},
SpscReceiver { ring },
)
}
pub struct SpscSender<T, const CAP: usize> {
ring: std::sync::Arc<SpscRing<T, CAP>>,
}
impl<T, const CAP: usize> SpscSender<T, CAP> {
#[inline]
pub fn send(&self, value: T) -> Result<(), T> {
match self.ring.try_push(value) {
None => Ok(()),
Some(back) => Err(back),
}
}
}
pub struct SpscReceiver<T, const CAP: usize> {
ring: std::sync::Arc<SpscRing<T, CAP>>,
}
impl<T, const CAP: usize> SpscReceiver<T, CAP> {
#[inline]
pub fn recv(&self) -> Option<T> {
self.ring.try_pop()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn spsc_basic() {
let (tx, rx) = channel::<u64, 4>();
tx.send(1).expect("fits");
tx.send(2).expect("fits");
assert_eq!(rx.recv(), Some(1));
assert_eq!(rx.recv(), Some(2));
assert_eq!(rx.recv(), None);
}
#[test]
fn full_returns_back() {
let (tx, rx) = channel::<u8, 2>();
tx.send(1).unwrap();
tx.send(2).unwrap();
assert_eq!(tx.send(3), Err(3));
assert_eq!(rx.recv(), Some(1));
tx.send(3).expect("space now");
assert_eq!(rx.recv(), Some(2));
assert_eq!(rx.recv(), Some(3));
}
#[test]
fn wraparound() {
let (tx, rx) = channel::<usize, 4>();
for i in 0..10_000usize {
tx.send(i).expect("loop drains");
assert_eq!(rx.recv(), Some(i));
}
}
#[test]
fn cross_thread_sum() {
let (tx, rx) = channel::<u64, 1024>();
std::thread::scope(|s| {
s.spawn(|| {
for i in 0..100_000u64 {
while tx.send(i).is_err() {
std::hint::spin_loop();
}
}
});
let got: u64 = (0..100_000u64)
.map(|_| {
loop {
if let Some(v) = rx.recv() {
break v;
} else {
std::hint::spin_loop();
}
}
})
.sum();
assert_eq!(got, 100_000u64 * 99_999 / 2);
});
}
}