FallibleLazyLock

lib.rs

use std::cell::UnsafeCell;
use std::mem::ManuallyDrop;
use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
use std::sync::atomic::{AtomicU8, Ordering};

const INITIALIZED_BIT: u8 = 0b01;
const LOCKED_BIT: u8 = 0b10;

union Data<T, F> {
    value: ManuallyDrop<T>,
    f: ManuallyDrop<F>,
}

#[derive(Debug)]
pub struct FallibleLazyLock<T, E, F = fn() -> Result<T, E>>
where
    F: Fn() -> Result<T, E>,
{
    state: AtomicU8,
    data: UnsafeCell<Data<T, F>>,
}

impl<T, E, F> FallibleLazyLock<T, E, F>
where
    F: Fn() -> Result<T, E>,
{
    pub const fn new(f: F) -> Self {
        Self {
            state: AtomicU8::new(0),
            data: UnsafeCell::new(Data {
                f: ManuallyDrop::new(f),
            }),
        }
    }

    pub fn try_force(&self) -> Result<&T, E> {
        let mut spinwait = parking_lot_core::SpinWait::new();

        let mut state = self.state.load(Ordering::Acquire);
        let data = self.data.get();

        loop {
            if state & INITIALIZED_BIT != 0 {
                break Ok(unsafe { &(*data).value });
            }

            if state & LOCKED_BIT != 0 && spinwait.spin() {
                unsafe {
                    parking_lot_core::park(
                        &raw const self.state as usize,
                        || self.state.load(Ordering::Relaxed) & LOCKED_BIT != 0,
                        || {},
                        |_, _| unreachable!(),
                        parking_lot_core::DEFAULT_PARK_TOKEN,
                        None,
                    );
                }
            }

            if let Err(new_state) =
                self.state
                    .compare_exchange_weak(0, LOCKED_BIT, Ordering::Relaxed, Ordering::Acquire)
            {
                state = new_state;
                continue;
            }

            match catch_unwind(AssertUnwindSafe(unsafe { &*(*data).f })) {
                Ok(Ok(value)) => unsafe {
                    ManuallyDrop::drop(&mut (*data).f);

                    *data = Data {
                        value: ManuallyDrop::new(value),
                    };

                    self.state.store(INITIALIZED_BIT, Ordering::Release);

                    parking_lot_core::unpark_all(
                        &raw const self.state as usize,
                        parking_lot_core::DEFAULT_UNPARK_TOKEN,
                    );

                    break Ok(&(*data).value);
                },
                Ok(Err(error)) => {
                    self.state.store(0, Ordering::Relaxed);

                    unsafe {
                        parking_lot_core::unpark_one(&raw const self.state as usize, |_| {
                            parking_lot_core::DEFAULT_UNPARK_TOKEN
                        });
                    }

                    break Err(error);
                }
                Err(panic) => {
                    self.state.store(0, Ordering::Relaxed);

                    unsafe {
                        parking_lot_core::unpark_one(&raw const self.state as usize, |_| {
                            parking_lot_core::DEFAULT_UNPARK_TOKEN
                        });
                    }

                    resume_unwind(panic);
                }
            }
        }
    }

    pub fn try_force_mut(&mut self) -> Result<&mut T, E> {
        let state = self.state.get_mut();
        let data = self.data.get_mut();

        if *state & INITIALIZED_BIT != 0 {
            return Ok(unsafe { &mut data.value });
        }

        unsafe {
            let value = (*data.f)()?;

            ManuallyDrop::drop(&mut data.f);

            *data = Data {
                value: ManuallyDrop::new(value),
            };

            *state = INITIALIZED_BIT;

            Ok(&mut data.value)
        }
    }

    pub fn get(&self) -> Option<&T> {
        let state = self.state.load(Ordering::Acquire);
        let data = self.data.get();

        if state & INITIALIZED_BIT != 0 {
            Some(unsafe { &(*data).value })
        } else {
            None
        }
    }

    pub fn get_mut(&mut self) -> Option<&mut T> {
        let state = self.state.get_mut();
        let data = self.data.get_mut();

        if *state & INITIALIZED_BIT != 0 {
            Some(unsafe { &mut data.value })
        } else {
            None
        }
    }

    pub fn into_inner(self) -> Result<T, F> {
        let this = &mut *ManuallyDrop::new(self);
        let state = this.state.get_mut();
        let data = this.data.get_mut();

        if *state & INITIALIZED_BIT != 0 {
            Ok(unsafe { ManuallyDrop::take(&mut data.value) })
        } else {
            Err(unsafe { ManuallyDrop::take(&mut data.f) })
        }
    }
}

impl<T, E, F> Drop for FallibleLazyLock<T, E, F>
where
    F: Fn() -> Result<T, E>,
{
    fn drop(&mut self) {
        let state = self.state.get_mut();
        let data = self.data.get_mut();

        if *state & INITIALIZED_BIT != 0 {
            unsafe { ManuallyDrop::drop(&mut data.value) };
        } else {
            unsafe { ManuallyDrop::drop(&mut data.f) };
        }
    }
}

unsafe impl<T, E, F> Sync for FallibleLazyLock<T, E, F>
where
    T: Sync + Send,
    F: Fn() -> Result<T, E> + Send,
{
}