use std::{cell::Cell, rc::Rc, task};
use crate::task::LocalWaker;
pub(super) struct Counter(Rc<CounterInner>);
struct CounterInner {
count: Cell<usize>,
capacity: usize,
task: LocalWaker,
}
impl Counter {
pub(super) fn new(capacity: usize) -> Self {
Counter(Rc::new(CounterInner {
capacity,
count: Cell::new(0),
task: LocalWaker::new(),
}))
}
pub(super) fn get(&self) -> CounterGuard {
CounterGuard::new(self.0.clone())
}
pub(super) fn available(&self, cx: &mut task::Context<'_>) -> bool {
self.0.available(cx)
}
pub(super) fn total(&self) -> usize {
self.0.count.get()
}
pub(super) fn priv_clone(&self) -> Self {
Counter(self.0.clone())
}
}
pub(super) struct CounterGuard(Rc<CounterInner>);
impl CounterGuard {
fn new(inner: Rc<CounterInner>) -> Self {
inner.inc();
CounterGuard(inner)
}
}
impl Unpin for CounterGuard {}
impl Drop for CounterGuard {
fn drop(&mut self) {
self.0.dec();
}
}
impl CounterInner {
fn inc(&self) {
self.count.set(self.count.get() + 1);
}
fn dec(&self) {
let num = self.count.get();
self.count.set(num - 1);
if num == self.capacity {
self.task.wake();
}
}
fn available(&self, cx: &mut task::Context<'_>) -> bool {
if self.count.get() < self.capacity {
true
} else {
self.task.register(cx.waker());
false
}
}
}