use std::array;
use crate::marker::ThreadBound;
pub struct FairCredits<const N: usize = 1> {
capacity: [usize; N],
available: [usize; N],
protected: [usize; N],
held: Box<[[usize; N]]>,
reserve: [usize; N],
_thread: ThreadBound,
}
impl<const N: usize> FairCredits<N> {
pub fn from_capacities(capacity: [usize; N], lane_count: usize) -> Self {
assert!(N > 0, "credit dimension count must be positive");
assert!(lane_count > 0, "credit lane count must be positive");
let reserve = capacity.map(|amount| {
if lane_count == 1 {
amount
} else {
amount / lane_count / 2
}
});
Self::with_reserve_per_lane(capacity, lane_count, reserve)
}
fn with_reserve_per_lane(
capacity: [usize; N],
lane_count: usize,
reserve_per_lane: [usize; N],
) -> Self {
assert!(N > 0, "credit dimension count must be positive");
assert!(lane_count > 0, "credit lane count must be positive");
let protected = array::from_fn(|dimension| {
assert!(
reserve_per_lane[dimension] <= capacity[dimension] / lane_count,
"credit reserve exceeds capacity"
);
reserve_per_lane[dimension] * lane_count
});
Self {
capacity,
available: capacity,
protected,
held: vec![[0; N]; lane_count].into_boxed_slice(),
reserve: reserve_per_lane,
_thread: ThreadBound::NEW,
}
}
fn can_acquire_all(&self, lane: usize, amount: [usize; N]) -> bool {
let Some(held) = self.held.get(lane) else {
return false;
};
for dimension in 0..N {
if amount[dimension] > self.available[dimension] {
return false;
}
let own = self
.unclaimed(dimension, held[dimension])
.min(amount[dimension]);
if amount[dimension] - own > self.available[dimension] - self.protected[dimension] {
return false;
}
}
true
}
pub fn try_acquire_all(&mut self, lane: usize, amount: [usize; N]) -> bool {
let Some(held) = self.held.get(lane).copied() else {
return false;
};
let mut own = [0; N];
for dimension in 0..N {
if amount[dimension] > self.available[dimension] {
return false;
}
own[dimension] = self
.unclaimed(dimension, held[dimension])
.min(amount[dimension]);
if amount[dimension] - own[dimension]
> self.available[dimension] - self.protected[dimension]
{
return false;
}
}
for dimension in 0..N {
self.held[lane][dimension] = held[dimension] + amount[dimension];
self.available[dimension] -= amount[dimension];
self.protected[dimension] -= own[dimension];
}
true
}
pub fn release_all(&mut self, lane: usize, amount: [usize; N]) {
let held = self.held[lane];
for dimension in 0..N {
assert!(
held[dimension] >= amount[dimension],
"cannot release credits that are not held"
);
}
for dimension in 0..N {
let next = held[dimension] - amount[dimension];
let before = self.unclaimed(dimension, held[dimension]);
let after = self.unclaimed(dimension, next);
self.held[lane][dimension] = next;
self.available[dimension] += amount[dimension];
self.protected[dimension] += after - before;
}
}
fn unclaimed(&self, dimension: usize, held: usize) -> usize {
self.reserve[dimension].saturating_sub(held)
}
}
impl FairCredits {
pub fn with_reserve(capacity: usize, lane_count: usize, reserve_per_lane: usize) -> Self {
Self::with_reserve_per_lane([capacity], lane_count, [reserve_per_lane])
}
pub fn used(&self) -> usize {
self.capacity[0] - self.available[0]
}
pub fn held_by(&self, lane: usize) -> Option<usize> {
self.held.get(lane).map(|held| held[0])
}
pub fn reserved_for(&self, lane: usize) -> Option<usize> {
self.held.get(lane).map(|_| self.reserve[0])
}
pub fn shared_available(&self) -> usize {
self.available[0] - self.protected[0]
}
#[inline]
pub fn can_acquire(&self, lane: usize, amount: usize) -> bool {
self.can_acquire_all(lane, [amount])
}
#[inline]
pub fn try_acquire(&mut self, lane: usize, amount: usize) -> bool {
self.try_acquire_all(lane, [amount])
}
#[inline]
pub fn release(&mut self, lane: usize, amount: usize) {
self.release_all(lane, [amount]);
}
}