Lazily-initialized Box<T>

lazy_box.rs

//! A crate for things that are
//! 1) Lazily initialized
//! 2) Expensive to create
//! 3) Immutable after creation
//! 4) Used on multiple threads
//! 5) Stored in a Box
//!
//! `LazyBox<T>` is better than `Mutex<Option<Box<T>>>` because after creation accessing
//! `T` does not require any locking, just a single boolean load with
//! `Ordering::Acquire` (which on x86 is just a compiler barrier, not an actual
//! memory barrier).
//!
//! Inspired by the lazy_init crate, but with more efficient storage.
//!
//! It could be made more generic to any NonNull, so as to have LazyArc too.

use std::mem::drop;
use std::ptr;
use std::sync::Mutex;
use std::sync::atomic::{AtomicPtr, Ordering};


/// `LazyBox<T>` is a synchronized holder type, that holds a lazily-initialized value of
/// type T.
pub struct LazyBox<T>
    where T: Sync + Send
{
    lock: Mutex<()>,
    value: AtomicPtr<T>
}

// Public API.
impl<T> LazyBox<T>
    where T: Sync + Send
{
    /// Construct a new, empty `LazyBox<T>` with an argument of type T.
    pub fn new() -> LazyBox<T> {
        LazyBox {
            lock: Mutex::new(()),
            value: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Get a reference to the stored value, invoking `f` to create it
    /// if the `LazyBox<T>` has yet to be initialized.  It is
    /// guaranteed that if multiple calls to `get_or_create` race, only one
    /// will invoke its closure, and every call will receive a reference to the
    /// newly created value.
    pub fn get_or_create<F, U>(&self, f: F) -> Result<&T, U>
        where F: FnOnce() -> Result<Box<T>, U>
    {
        let value = self.value.load(Ordering::Acquire);
        if !value.is_null() {
            // Safe because the reference we return has the same lifetime as self.
            unsafe {
                return Ok(value.as_ref().unwrap())
            }
        }

        // We *may* not be initialized. We have to block to be certain.
        let _lock = self.lock.lock().unwrap();
        let value = self.value.load(Ordering::Acquire);
        if !value.is_null() {
            // We raced, and someone else initialized us. We can fall
            // through now.  Safe as above.
            unsafe {
                return Ok(value.as_ref().unwrap())
            }
        }

        let result = f()?;
        let ptr = Box::into_raw(result);
        self.value.store(ptr as *mut T, Ordering::Release);

        // Safe to unwrap, because the pointer cannot be null.
        unsafe { Ok(ptr.as_ref().unwrap()) }
    }

    /// Get a reference to the stored value, returning `Some(&T)` if the
    /// `LazyBox<T>` has been initialized or `None` if it has not.  It
    /// is guaranteed that if a reference is returned it is to the created
    /// value inside the `LazyBox<T>`.
    pub fn get(&self) -> Option<&T> {
        let ptr = self.value.load(Ordering::Acquire);

        // Safe because the reference we return has the same lifetime as self,
        // and the box exists until it is freed by drop.
        unsafe { ptr.as_ref() }
    }

    fn take(&mut self) -> Option<Box<T>> {
        let ptr = self.value.load(Ordering::Acquire);

        self.value.store(ptr::null_mut(), Ordering::Relaxed);
        // Safe to move out of the pointer because there are no other references to self,
        // and the pointer won't be extracted again because we just cleared it.
        unsafe {
            ptr.as_mut().map(|ptr| Box::from_raw(ptr))
        }
    }

    /// Unwrap the contained value, returning `Some` if the `LazyBox<T>` has
    /// been initialized or None if it has not.
    pub fn into_inner(mut self) -> Option<Box<T>> {
        // We don't want a LazyBox to be initialized more than once, so the
        // take() method (which takes &mut self) is not public.  into_inner()
        // however consumes self so it can be made public.
        self.take()
    }
}

impl<T> Drop for LazyBox<T>
    where T: Sync + Send
{
    fn drop(&mut self) {
        drop(self.take())
    }
}

// `#[derive(Default)]` automatically adds `T: Default` trait bound, but that
// is too restrictive, because `LazyBox<T>` always has a default value for any `T`.
impl<T> Default for LazyBox<T>
    where T: Sync + Send
{
    fn default() -> Self {
        LazyBox::new()
    }
}

#[cfg(test)]
mod tests {

    use std::{thread, time};
    use std::sync::atomic::{AtomicUsize, Ordering};
    use super::LazyBox;

    #[test]
    fn test_lazy() {
        let lazy_value: LazyBox<u8> = LazyBox::new();

        assert_eq!(lazy_value.get(), None);
        assert_eq!(lazy_value.get_or_create(|| Err("bad")), Err("bad"));
        assert_eq!(lazy_value.get(), None);

        let n = AtomicUsize::new(0);

        let value = lazy_value.get_or_create::<_, ()>(|| {
            n.fetch_add(1, Ordering::Relaxed);
            Ok(Box::new(42))
        });
        assert!(value.is_ok());
        assert_eq!(*value.unwrap(), 42);

        let value = lazy_value.get();
        assert!(value.is_some());
        assert_eq!(*value.unwrap(), 42);

        let value = lazy_value.get_or_create::<_, ()>(|| {
            n.fetch_add(1, Ordering::Relaxed);
            Ok(Box::new(43))
        });
        assert!(value.is_ok());
        assert_eq!(*value.unwrap(), 42);

        assert_eq!(*lazy_value.into_inner().unwrap(), 42);
        assert_eq!(n.load(Ordering::SeqCst), 1);
    }
}