use core::task::Waker;
use crate::error::CallError;
use crate::port::{Correlation, Interest};
const SLOT_BITS: u32 = 16;
const SLOT_MASK: u64 = (1 << SLOT_BITS) - 1;
const GENERATION_MASK: u64 = (1 << (64 - SLOT_BITS)) - 1;
#[must_use = "a waker or a reclaim handed back and ignored leaves a task waiting"]
#[derive(Debug)]
pub enum Settled {
Recorded(Option<Waker>),
Reclaimed,
Unknown,
}
#[must_use = "a waker or a reclaim handed back and ignored leaves a task waiting"]
#[derive(Debug)]
pub enum Forgotten {
Reclaimed,
Marked(Option<Waker>),
Unknown,
}
#[derive(Clone, Copy, Debug)]
enum State {
Free,
InFlight { forgotten: bool },
Settled(Result<(), CallError>),
}
#[derive(Debug)]
struct Entry {
generation: u64,
state: State,
reservation: u64,
waker: Option<Waker>,
}
impl Entry {
const FREE: Entry = Entry {
generation: 0,
state: State::Free,
reservation: 0,
waker: None,
};
}
#[derive(Debug)]
pub struct Table<const N: usize> {
slots: [Entry; N],
budget: Option<u64>,
}
impl<const N: usize> Table<N> {
const SLOT_INDEX_FITS: () = assert!(
N <= 1 << SLOT_BITS,
"a correlate::Table holds at most 65536 slots"
);
pub const fn new(budget: Option<u64>) -> Self {
let () = Self::SLOT_INDEX_FITS;
Table {
slots: [Entry::FREE; N],
budget,
}
}
pub fn insert(&mut self, reservation: u64) -> Option<Correlation> {
let index = self
.slots
.iter()
.position(|entry| matches!(entry.state, State::Free))?;
if let Some(free) = self.budget {
self.budget = Some(free.checked_sub(reservation)?);
}
let entry = &mut self.slots[index];
entry.state = State::InFlight { forgotten: false };
entry.reservation = reservation;
Some(Correlation((entry.generation << SLOT_BITS) | index as u64))
}
pub fn settle(&mut self, c: Correlation, outcome: Result<(), CallError>) -> Settled {
let Some(index) = self.held(c) else {
return Settled::Unknown;
};
match self.slots[index].state {
State::InFlight { forgotten: false } => {
let entry = &mut self.slots[index];
entry.state = State::Settled(outcome);
Settled::Recorded(entry.waker.take())
}
State::InFlight { forgotten: true } => {
self.reclaim(index);
Settled::Reclaimed
}
State::Settled(_) | State::Free => Settled::Unknown,
}
}
#[must_use]
pub fn outcome(&self, c: Correlation) -> Option<Result<(), CallError>> {
match self.slots[self.held(c)?].state {
State::Settled(outcome) => Some(outcome),
State::InFlight { .. } | State::Free => None,
}
}
pub fn forget(&mut self, c: Correlation) -> Forgotten {
let Some(index) = self.held(c) else {
return Forgotten::Unknown;
};
match self.slots[index].state {
State::Settled(_) => {
self.reclaim(index);
Forgotten::Reclaimed
}
State::InFlight { forgotten: false } => {
let entry = &mut self.slots[index];
entry.state = State::InFlight { forgotten: true };
Forgotten::Marked(entry.waker.take())
}
State::InFlight { forgotten: true } | State::Free => Forgotten::Unknown,
}
}
pub fn wake_on(&mut self, c: Correlation, waker: &Waker) -> Option<Waker> {
let Some(index) = self.held(c) else {
return Some(waker.clone());
};
let entry = &mut self.slots[index];
match entry.state {
State::InFlight { forgotten: false } => store(&mut entry.waker, waker),
State::InFlight { forgotten: true } | State::Settled(_) | State::Free => {
Some(waker.clone())
}
}
}
#[must_use]
pub fn slot(c: Correlation) -> usize {
(c.0 & SLOT_MASK) as usize
}
fn held(&self, c: Correlation) -> Option<usize> {
let index = Self::slot(c);
let entry = self.slots.get(index)?;
let generation = c.0 >> SLOT_BITS;
(entry.generation == generation && !matches!(entry.state, State::Free)).then_some(index)
}
fn reclaim(&mut self, index: usize) {
let entry = &mut self.slots[index];
if let Some(free) = self.budget {
self.budget = Some(free.saturating_add(entry.reservation));
}
entry.state = State::Free;
entry.reservation = 0;
entry.waker = None;
entry.generation = entry.generation.wrapping_add(1) & GENERATION_MASK;
}
}
#[derive(Debug, Default)]
pub struct Waiters {
slot: Option<Waker>,
event: Option<Waker>,
claim: Option<Waker>,
}
impl Waiters {
#[must_use]
pub const fn new() -> Self {
Waiters {
slot: None,
event: None,
claim: None,
}
}
pub fn register(&mut self, what: Interest, waker: &Waker) -> Option<Waker> {
match self.kind(what) {
Some(stored) => store(stored, waker),
None => Some(waker.clone()),
}
}
pub fn take(&mut self, what: Interest) -> Option<Waker> {
self.kind(what)?.take()
}
pub fn take_all(&mut self) -> impl Iterator<Item = Waker> + use<> {
[self.slot.take(), self.event.take(), self.claim.take()]
.into_iter()
.flatten()
}
fn kind(&mut self, what: Interest) -> Option<&mut Option<Waker>> {
match what {
Interest::Slot => Some(&mut self.slot),
Interest::Event(_) => Some(&mut self.event),
Interest::Claim(_) => Some(&mut self.claim),
Interest::Outcome(_) => None,
}
}
}
fn store(stored: &mut Option<Waker>, waker: &Waker) -> Option<Waker> {
stored
.replace(waker.clone())
.filter(|displaced| !displaced.will_wake(waker))
}