use crate::misc::PhantomOnce;
use std::fmt;
use std::sync::atomic::{AtomicU8, Ordering};
pub struct Once {
state: AtomicU8,
_phantom: PhantomOnce,
}
impl fmt::Debug for Once {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("Once { .. }")
}
}
impl Once {
pub const fn new() -> Self {
Self {
state: AtomicU8::new(OnceState::default().state),
_phantom: PhantomOnce {},
}
}
pub fn call_once<F: FnOnce()>(&self, f: F) {
let (_guard, s) = self.lock();
if s.poisoned() {
panic!("`Once.call_once()' is called while the instance is poisoned.");
}
if s.finished() {
return;
}
f();
let s = s.finish();
self.state.store(s.state, Ordering::Relaxed);
}
pub fn call_once_force<F: FnOnce(&OnceState)>(&self, f: F) {
let (_guard, s) = self.lock();
if s.finished() {
return;
}
f(&s);
let s = s.finish();
let s = s.unpoison();
self.state.store(s.state, Ordering::Relaxed);
}
#[must_use]
pub fn is_completed(&self) -> bool {
let s = OnceState::new(self.state.load(Ordering::Relaxed));
(!s.poisoned()) && (s.finished())
}
fn lock(&self) -> (OnceGuard, OnceState) {
let mut expected = OnceState::default();
loop {
let desired = expected.acquire_lock();
let current = OnceState::new(self.state.compare_and_swap(
expected.state,
desired.state,
Ordering::Acquire,
));
if current.locked() {
expected = current.release_lock();
std::thread::yield_now();
continue;
}
if current.state == expected.state {
return (OnceGuard { once: &self }, desired);
}
expected = current;
}
}
}
struct OnceGuard<'a> {
once: &'a Once,
}
impl Drop for OnceGuard<'_> {
fn drop(&mut self) {
let mut s = OnceState::new(self.once.state.load(Ordering::Relaxed));
debug_assert!(s.locked());
if std::thread::panicking() {
s = s.poison();
}
s = s.release_lock();
self.once.state.store(s.state, Ordering::Release);
}
}
#[derive(Debug)]
pub struct OnceState {
state: u8,
}
impl OnceState {
const INIT: u8 = 0;
const LOCK: u8 = 1;
const FINISHED: u8 = 2;
const POISONED: u8 = 4;
#[must_use]
const fn default() -> Self {
Self { state: Self::INIT }
}
#[must_use]
const fn new(state: u8) -> Self {
Self { state }
}
#[must_use]
const fn locked(&self) -> bool {
(self.state & Self::LOCK) != 0
}
#[must_use]
const fn finished(&self) -> bool {
(self.state & Self::FINISHED) != 0
}
#[must_use]
pub const fn poisoned(&self) -> bool {
(self.state & Self::POISONED) != 0
}
#[must_use]
fn acquire_lock(&self) -> Self {
debug_assert!(!self.locked());
Self::new(self.state | Self::LOCK)
}
#[must_use]
fn release_lock(&self) -> Self {
debug_assert!(self.locked());
Self::new(self.state ^ Self::LOCK)
}
#[must_use]
fn finish(&self) -> Self {
debug_assert!(!self.finished());
Self::new(self.state | Self::FINISHED)
}
#[must_use]
fn poison(&self) -> Self {
Self::new(self.state | Self::POISONED)
}
#[must_use]
fn unpoison(&self) -> Self {
Self::new(self.state ^ Self::POISONED)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn call_once_invoke_task_only_once() {
let mut val = 0;
let once = Once::new();
assert_eq!(0, val);
once.call_once(|| val = 1);
assert_eq!(1, val);
once.call_once(|| val = 2);
assert_eq!(1, val);
}
#[test]
fn call_once_force_do_nothing_after_call_once_succeeded() {
let mut val = 0;
let once = Once::new();
assert_eq!(0, val);
once.call_once(|| val = 1);
assert_eq!(1, val);
once.call_once_force(|_| val = 2);
assert_eq!(1, val);
}
}