#![no_std]
#![allow(clippy::bool_comparison)]
use core::cell::Cell;
use core::marker::PhantomData;
use core::mem::{self, MaybeUninit};
use core::ptr;
use core::sync::atomic::{AtomicBool, Ordering};
use core::fmt;
type PhantomUnsync = PhantomData<Cell<u8>>;
pub struct AtomicTake<T> {
taken: AtomicBool,
value: MaybeUninit<T>,
_unsync: PhantomUnsync,
}
impl<T> AtomicTake<T> {
pub const fn new(value: T) -> Self {
AtomicTake {
taken: AtomicBool::new(false),
value: MaybeUninit::new(value),
_unsync: PhantomData,
}
}
pub const fn empty() -> Self {
AtomicTake {
taken: AtomicBool::new(true),
value: MaybeUninit::uninit(),
_unsync: PhantomData,
}
}
pub fn take(&self) -> Option<T> {
if self.taken.swap(true, Ordering::Relaxed) == false {
unsafe { Some(ptr::read(self.value.as_ptr())) }
} else {
None
}
}
pub fn take_mut(&mut self) -> Option<T> {
if mem::replace(self.taken.get_mut(), true) == false {
unsafe { Some(ptr::read(self.value.as_ptr())) }
} else {
None
}
}
pub fn is_taken(&self) -> bool {
self.taken.load(Ordering::Relaxed)
}
pub fn insert(&mut self, value: T) -> Option<T> {
let previous = self.take_mut();
unsafe {
ptr::write(self.value.as_mut_ptr(), value);
*self.taken.get_mut() = false;
}
previous
}
}
impl<T> Drop for AtomicTake<T> {
fn drop(&mut self) {
if !*self.taken.get_mut() {
unsafe {
ptr::drop_in_place(self.value.as_mut_ptr());
}
}
}
}
unsafe impl<T: Send> Sync for AtomicTake<T> {}
impl<T> From<T> for AtomicTake<T> {
fn from(t: T) -> AtomicTake<T> {
AtomicTake::new(t)
}
}
impl<T> fmt::Debug for AtomicTake<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.is_taken() {
write!(f, "Empty AtomicTake")
} else {
write!(f, "Non-Empty AtomicTake")
}
}
}
#[cfg(test)]
mod tests {
use crate::AtomicTake;
struct CountDrops {
counter: *mut u32,
}
impl Drop for CountDrops {
fn drop(&mut self) {
unsafe {
*self.counter += 1;
}
}
}
struct PanicOnDrop;
impl Drop for PanicOnDrop {
fn drop(&mut self) {
panic!("Panic on drop called.");
}
}
#[test]
fn drop_calls_drop() {
let mut counter = 0;
let take = AtomicTake::new(CountDrops {
counter: &mut counter,
});
drop(take);
assert_eq!(counter, 1);
}
#[test]
fn taken_not_dropped_twice() {
let mut counter = 0;
let take = AtomicTake::new(CountDrops {
counter: &mut counter,
});
take.take();
assert_eq!(counter, 1);
drop(take);
assert_eq!(counter, 1);
}
#[test]
fn taken_mut_not_dropped_twice() {
let mut counter = 0;
let mut take = AtomicTake::new(CountDrops {
counter: &mut counter,
});
take.take_mut();
assert_eq!(counter, 1);
drop(take);
assert_eq!(counter, 1);
}
#[test]
fn insert_dropped_once() {
let mut counter1 = 0;
let mut counter2 = 0;
let mut take = AtomicTake::new(CountDrops {
counter: &mut counter1,
});
assert!(!take.is_taken());
take.insert(CountDrops {
counter: &mut counter2,
});
assert!(!take.is_taken());
drop(take);
assert_eq!(counter1, 1);
assert_eq!(counter2, 1);
}
#[test]
fn insert_take() {
let mut counter1 = 0;
let mut counter2 = 0;
let mut take = AtomicTake::new(CountDrops {
counter: &mut counter1,
});
take.insert(CountDrops {
counter: &mut counter2,
});
assert!(!take.is_taken());
assert_eq!(counter1, 1);
assert_eq!(counter2, 0);
drop(take);
assert_eq!(counter1, 1);
assert_eq!(counter2, 1);
}
#[test]
fn empty_no_drop() {
let take: AtomicTake<PanicOnDrop> = AtomicTake::empty();
assert!(take.is_taken());
drop(take);
}
#[test]
fn empty_insert() {
let mut take = AtomicTake::empty();
assert!(take.is_taken());
let mut counter = 0;
take.insert(CountDrops {
counter: &mut counter,
});
assert!(!take.is_taken());
drop(take);
assert_eq!(counter, 1);
}
}