use super::less;
use core::cell::Cell;
use core::mem::MaybeUninit;
use core::ptr;
use core::sync::atomic::{
fence, AtomicU32, AtomicUsize,
Ordering::{Acquire, Relaxed, Release},
};
pub(crate) struct TokenLink(AtomicUsize);
impl TokenLink {
pub(crate) const fn new() -> Self {
Self(AtomicUsize::new(0))
}
pub(crate) fn lock(&self, token: &Token, last: u16, idx: u16) -> *const Token {
debug_assert!(last != idx);
let new_status = ((last as u32) << 16) | idx as u32;
loop {
match token
.status
.compare_exchange_weak(0, new_status, Relaxed, Relaxed)
{
Ok(_) => return self.lock_with(token),
Err(old) => {
if (old as u16) != last && (old >> 16) as u16 != idx {
continue;
}
let current = self.try_lock_update();
if current.is_null() {
continue;
}
let status = token.status.load(Relaxed);
if status == 0 {
token.status.store(new_status, Relaxed);
token.set_next(current);
return token;
} else if (status as u16) == last {
token.set_last_and_idx((status >> 16) as u16, idx);
return current;
} else {
token.set_last_and_idx(last, status as u16);
return current;
}
}
}
}
}
fn try_lock_update(&self) -> *const Token {
let old = self.0.load(Relaxed);
if old > 0
&& (old & 1) == 0
&& self
.0
.compare_exchange_weak(old, 1, Release, Relaxed)
.is_ok()
{
fence(Acquire);
return old as *const Token;
}
ptr::null()
}
fn lock_with(&self, token: &Token) -> *const Token {
let addr = token as *const Token as usize;
token.set_next(ptr::null());
let mut old = self.0.load(Relaxed);
loop {
let new = if (old & 1) == 1 {
token.set_next((old & !1) as *const Token);
addr | 1
} else {
1
};
match self.0.compare_exchange_weak(old, new, Release, Relaxed) {
Ok(_) => {
return if (old & 1) == 0 {
fence(Acquire);
token.set_next((old & !1) as *const Token);
token
} else {
ptr::null()
};
}
Err(val) => old = val,
}
}
}
#[inline(always)]
fn next_range(old: usize, last: u16) -> bool {
debug_assert!(old > 0);
if old != 1 {
let task = unsafe { &*((old & !1) as *const Token) };
return task.find(last);
}
false
}
#[inline(always)]
pub(crate) fn unlock(&self, token: *const Token, last: u16) -> *const Token {
if token.is_null() {
self.unlock_with_null(last)
} else {
self.unlock_with_next(unsafe { &*token }, last)
}
}
fn unlock_with_null(&self, last: u16) -> *const Token {
let mut old = self.0.load(Relaxed);
loop {
let next_range = Self::next_range(old, last);
let new = if !next_range { old & !1 } else { 1 };
match self.0.compare_exchange_weak(old, new, Release, Relaxed) {
Ok(_) => {
return if new == 1 {
fence(Acquire);
(old & !1) as *const Token
} else {
ptr::null()
};
}
Err(val) => old = val,
}
}
}
fn unlock_with_next(&self, token: &Token, last: u16) -> *const Token {
let end = token.last();
let mut old = self.0.load(Relaxed);
loop {
let next_range = Self::next_range(old, last);
let new = if !next_range {
end.set_next((old & !1) as *const Token);
token as *const Token as usize
} else {
1
};
match self.0.compare_exchange_weak(old, new, Release, Relaxed) {
Ok(_) => {
return if new == 1 {
fence(Acquire);
end.set_next((old & !1) as *const Token);
token as *const Token
} else {
ptr::null()
};
}
Err(val) => old = val,
}
}
}
}
pub trait TokenFactory {
fn new() -> Self;
fn get(&self) -> &Token;
}
pub struct NullTokenFactory;
static NULL_TOKEN: Token = Token::new();
impl TokenFactory for NullTokenFactory {
fn new() -> Self {
NullTokenFactory
}
fn get(&self) -> &Token {
&NULL_TOKEN
}
}
pub type TokenArray = Tokens<2>;
#[repr(C)]
pub struct Tokens<const CNT: usize> {
idx: Cell<usize>,
tokens: MaybeUninit<[Token; CNT]>,
}
impl<const CNT: usize> TokenFactory for Tokens<CNT> {
fn new() -> Self {
debug_assert!((CNT & (CNT - 1)) == 0);
Self {
tokens: MaybeUninit::zeroed(),
idx: Cell::new(0),
}
}
fn get(&self) -> &Token {
let idx = self.idx.get();
self.idx.set((idx + 1) & (CNT - 1));
unsafe { self.tokens.assume_init_ref().get_unchecked(idx) }
}
}
#[repr(C)]
pub struct Token {
status: AtomicU32,
next: Cell<*const Token>,
}
unsafe impl Send for Token {}
unsafe impl Sync for Token {}
impl Token {
pub(crate) const fn new() -> Self {
Self {
status: AtomicU32::new(0),
next: Cell::new(ptr::null()),
}
}
pub(crate) fn find(&self, last: u16) -> bool {
let mut token = self;
loop {
let (expected, _) = token.last_and_idx();
if expected == last {
return true;
}
let next = token.next();
if next.is_null() {
return false;
}
token = unsafe { &*next };
}
}
pub(crate) fn unlock(&self) -> *const Token {
let next = self.next.replace(ptr::null());
self.status.store(0, Release);
next
}
pub(crate) fn last_and_idx(&self) -> (u16, u16) {
let status = self.status.load(Relaxed);
((status >> 16) as u16, status as u16)
}
pub(crate) fn set_last_and_idx(&self, last: u16, idx: u16) {
let new_status = ((last as u32) << 16) | idx as u32;
self.status.store(new_status, Relaxed);
}
fn try_merge_next(&self, idx: u16) -> u16 {
let next = self.next();
if next.is_null() {
return idx;
}
let next = unsafe { &*next };
let (next_last, next_idx) = next.last_and_idx();
if idx == next_last {
self.set_next(next.unlock());
next_idx
} else {
idx
}
}
pub(crate) fn next(&self) -> *const Token {
self.next.get()
}
pub(crate) fn set_next(&self, token: *const Token) {
self.next.set(token);
}
pub(crate) fn sort(&self) -> &Token {
let mut sorted = SortedToken::new();
let mut token = self as *const Token;
while !token.is_null() {
let token_ref = unsafe { &*token };
let next = token_ref.next();
sorted.insert(token_ref);
token = next;
}
unsafe { &*sorted.token }
}
pub(crate) fn last(&self) -> &Self {
let mut last = self;
while !last.next().is_null() {
last = unsafe { &*last.next() };
}
last
}
}
struct SortedToken {
token: *const Token,
}
impl SortedToken {
fn new() -> Self {
Self { token: ptr::null() }
}
fn insert(&mut self, token: &Token) {
token.set_next(ptr::null());
if !self.token.is_null() {
let (last, idx) = token.last_and_idx();
let mut prev = ptr::null::<Token>();
let mut pos = self.token;
loop {
let cur = unsafe { &*pos };
let (cur_last, cur_idx) = cur.last_and_idx();
if idx == cur_last {
cur.set_last_and_idx(last, cur_idx);
token.unlock();
return;
} else if last == cur_idx {
let idx = cur.try_merge_next(idx);
cur.set_last_and_idx(cur_last, idx);
token.unlock();
return;
} else if less(idx, cur_last) {
token.set_next(cur);
if prev.is_null() {
self.token = token;
} else {
unsafe { &*prev }.set_next(token);
}
return;
} else if cur.next().is_null() {
cur.set_next(token);
return;
}
prev = pos;
pos = cur.next();
}
} else {
self.token = token;
}
}
}