#![no_std]
use core::{
cell::UnsafeCell,
mem::ManuallyDrop,
sync::atomic::{AtomicBool, Ordering},
};
#[derive(Default)]
pub struct TakeCell<T: ?Sized> {
taken: AtomicBool,
value: UnsafeCell<T>,
}
impl<T> TakeCell<T> {
pub const fn new(v: T) -> Self {
Self {
taken: AtomicBool::new(false),
value: UnsafeCell::new(v),
}
}
pub fn into_inner(self) -> T {
self.value.into_inner()
}
}
impl<T: ?Sized> TakeCell<T> {
pub fn take(&self) -> Option<&mut T> {
match self.taken.swap(true, Ordering::Relaxed) {
true => None,
false => Some(unsafe { &mut *self.value.get() }),
}
}
pub fn is_taken(&self) -> bool {
self.taken.load(Ordering::Relaxed)
}
pub fn get(&mut self) -> &mut T {
self.value.get_mut()
}
pub fn heal(&mut self) {
self.taken = AtomicBool::new(false);
}
pub fn is_taken_unsync(&mut self) -> bool {
*self.taken.get_mut()
}
pub fn take_unsync(&mut self) -> Option<&mut T> {
match self.is_taken_unsync() {
false => {
*self.taken.get_mut() = true;
Some(self.get())
}
true => None,
}
}
#[allow(clippy::mut_from_ref)]
pub unsafe fn steal(&self) -> &mut T {
self.taken.store(true, Ordering::Relaxed);
&mut *self.value.get()
}
}
impl<T> From<T> for TakeCell<T> {
fn from(v: T) -> Self {
Self::new(v)
}
}
unsafe impl<T: ?Sized + Send> Sync for TakeCell<T> {}
#[derive(Default)]
pub struct TakeOwnCell<T>(
TakeCell<ManuallyDrop<T>>,
);
impl<T> TakeOwnCell<T> {
pub const fn new(v: T) -> Self {
Self(TakeCell::new(ManuallyDrop::new(v)))
}
pub fn take(&self) -> Option<T> {
self.0
.take()
.map(|value| unsafe { ManuallyDrop::take(value) })
}
pub fn is_taken(&self) -> bool {
self.0.is_taken()
}
pub fn get(&mut self) -> Option<&mut T> {
match self.is_taken() {
false => {
Some(&mut *self.0.get())
}
true => None,
}
}
pub fn into_inner(mut self) -> Option<T> {
self.take_unsync()
}
pub fn heal(&mut self, v: T) -> (&mut T, Result<(), T>) {
let res = match self.0.is_taken() {
true => {
*self.0.get() = ManuallyDrop::new(v);
Ok(())
}
false => Err(v),
};
self.0.heal();
let uref = &mut *self.0.get();
(uref, res)
}
pub fn is_taken_unsync(&mut self) -> bool {
self.0.is_taken_unsync()
}
pub fn take_unsync(&mut self) -> Option<T> {
self.0
.take_unsync()
.map(|value| unsafe { ManuallyDrop::take(value) })
}
pub unsafe fn steal(&self) -> T {
ManuallyDrop::take(self.0.steal())
}
}
impl<T> From<T> for TakeOwnCell<T> {
fn from(v: T) -> Self {
Self::new(v)
}
}
unsafe impl<T: Send> Sync for TakeOwnCell<T> {}
impl<T> Drop for TakeOwnCell<T> {
fn drop(&mut self) {
let _ = self.take_unsync();
}
}
#[cfg(test)]
mod tests {
use crate::TakeCell;
#[test]
fn it_works() {
let cell = TakeCell::new(0);
{
let uref = cell.take().unwrap();
*uref += 1;
assert_eq!(*uref, 1);
assert!(cell.take().is_none());
*uref += 1;
assert_eq!(*uref, 2);
}
assert!(cell.take().is_none());
assert_eq!(cell.into_inner(), 2);
}
#[test]
fn unsize() {
let cell: TakeCell<[i32; 10]> = TakeCell::new([0; 10]);
let _: &TakeCell<[i32]> = &cell;
let _: &TakeCell<dyn Send> = &cell;
}
#[test]
fn r#static() {
static CELL: TakeCell<i32> = TakeCell::new(0);
{
let uref: &'static mut i32 = CELL.take().unwrap();
*uref += 1;
assert_eq!(*uref, 1);
assert!(CELL.take().is_none());
*uref += 1;
assert_eq!(*uref, 2);
}
assert!(CELL.take().is_none());
}
#[test]
fn steal_takes() {
let cell = TakeCell::new(0);
let uref = unsafe { cell.steal() };
*uref += 1;
assert_eq!(*uref, 1);
assert!(cell.is_taken());
assert!(cell.take().is_none());
*uref += 1;
assert_eq!(*uref, 2);
assert_eq!(cell.into_inner(), 2);
}
}
#[cfg(test)]
mod own_tests {
use crate::TakeOwnCell;
#[test]
fn it_works() {
let cell = TakeOwnCell::new(17);
assert_eq!(cell.take(), Some(17));
assert!(cell.take().is_none());
assert!(cell.into_inner().is_none());
}
#[test]
fn heal() {
let mut cell = TakeOwnCell::new(17);
let (uref, res) = cell.heal(12);
assert_eq!(res, Err(12));
assert_eq!(*uref, 17);
*uref = 0xAA;
assert_eq!(cell.take(), Some(0xAA));
let (uref, res) = cell.heal(12);
assert!(res.is_ok());
assert_eq!(*uref, 12);
*uref = 0xBB;
assert_eq!(cell.into_inner(), Some(0xBB));
}
#[test]
fn r#static() {
static CELL: TakeOwnCell<i32> = TakeOwnCell::new(42);
assert!(!CELL.is_taken());
assert_eq!(CELL.take(), Some(42));
assert!(CELL.is_taken());
assert!(CELL.take().is_none());
}
#[test]
fn steal_takes() {
let cell = TakeOwnCell::new(1);
assert!(!cell.is_taken());
assert_eq!(unsafe { cell.steal() }, 1);
assert!(cell.is_taken());
assert!(cell.take().is_none());
assert!(cell.into_inner().is_none());
}
}