hirun/channel/

token.rs

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;
        }
    }
}