Reimplementing lazy-static, and then once_cell's Lazy

Jan 13, 2023

lazy_static is a crate which allows you to lazily initialize static items. For example:

lazy_static! {
  static ref LIST: HashSet<&'static str> = {
    let mut set = HashSet::new();
    set.insert("One");
    set.insert("Two");
    set.insert("Three");
    set
  }
}

I was curious how it all worked under the hood. However I didn't spend enough time actually reading all the code. Instead, I got as far as the no_std implementation which uses the spin crate's Once implementation.

lazy_static!

lazy-static relies on two things to work: A type to hold the data that you'd want to store in a static item; And a Deref implementation on an arbitrary wrapper type. The latter is implemented per static item.

The idea is that the final static item is actually an instance of the arbitrary type which can then be dererenced to fetch the data held by a Lazy instance it knows about. So if you wanted to have a static item LIST: Vec<&'static str>, you would implement the wrapper type like this:

struct StaticSet;

impl Deref for StaticSet {
    type Target = HashSet<&'static str>;

    fn deref(&self) -> &Self::Target {
        static LAZY: Lazy<HashSet<&'static str>> = Lazy::new();
        LAZY.get(|| {
            let mut set = HashSet::new();
            set.insert("One");
            set.insert("Two");
            set.insert("Three");
            set
        })
    }
}

static SET: StaticSet = StaticSet;

StaticList is unique for every lazy static item and it is what you initialize the static item with since it is const. When it is accessed (and automatically dereferenced), for example SET.get("One"), the Deref impl will run the lazy evaluation logic.

For the Lazy type, we want to be able to pass in a closure which would initialize and return the static item's value and which should only ever be called once. Subsequent calls should always return the result of that closures evaluation.

use std::{
    mem::MaybeUninit,
    sync::{Mutex, Once},
};

pub struct Lazy<T> {
    init: Once,
    data: Mutex<MaybeUninit<T>>,
}

impl<T> Lazy<T> {
    pub const fn new() -> Self {
        Self {
            init: Once::new(),
            data: Mutex::new(MaybeUninit::uninit()),
        }
    }

    pub fn get(&'static self, init: impl FnOnce() -> T) -> &T {
        self.init.call_once(|| {
            self.data
                .lock()
                .expect("Poisoned")
                .write(init());
        });

        // SAFETY: When call_once has completed, we know that the inner data has been written 
        // and we can get a reference to it.
        unsafe {
            &*self
                .data
                .lock()
                .expect("Poisoned")
                .as_ptr()
        }
    }
}

const fn new() so that we can initialize the LAZY static item in the deref implementation. get will call the init closure exactly once, so we can guarantee that after the closure has returned, the data will be written in data.

It turned out lazy_static has a no_std and a std implementation which I completely neglected to notice until it was too late. Having a look at the std implementation, they're using std::cell::Cell instead of Mutex which I was using. Cell is Send, but not Sync however, so they've implemented Sync for the Lazy type. In this instance we wouldn't need to worry about Cell's interior mutability not being synchronized across threads because we write to it exactly once when the init closure returns.

use std::{
    cell::Cell,
    mem::MaybeUninit,
    ops::Deref,
    sync::{Mutex, Once},
};

pub struct Lazy<T> {
    init: Once,
    data: Cell<MaybeUninit<T>>,
}

unsafe impl<T: Sync> Sync for Lazy<T> {}

impl<T> Lazy<T> {
    pub const fn new() -> Self {
        Self {
            init: Once::new(),
            data: Cell::new(MaybeUninit::uninit()),
        }
    }

    pub fn get(&'static self, init: impl FnOnce() -> T) -> &T {
        self.init.call_once(|| unsafe {
            (*self.data.as_ptr()).write(init());
        });

        // SAFETY: Once call_once has completed, we know that the inner data has been written
        // and won't be written to again. This is also why we can impl Sync for Lazy<T>.
        unsafe { &*(*self.data.as_ptr()).as_ptr() }
    }
}

Lazy::new(|| {})

And then I noticed this error:

44 | static SET: HashSet<&'static str> = HashSet::new();
   |                                     ^^^^^^^^^^^^^^
   |
   = note: calls in statics are limited to constant functions, tuple structs and tuple variants
   = note: consider wrapping this expression in `Lazy::new(|| ...)` from the `once_cell` crate: https://crates.io/crates/once_cell

The compiler recommends to use once_cell for lazy static items...

I had a look at once_cell::Lazy and it's a very similar data structure to what I had before.

The difference between once_cell's implementation and lazy_static is that with once_cell::Lazy, you explicitly use the Lazy type, whereas lazy_static generates individual types for each static item:

use once_cell::Lazy;

static SET: Lazy<HashSet<&'static str>> = Lazy::new(|| {
    let mut set = HashSet::new();
    set.insert("One");
    set.insert("Two");
    set.insert("Three");
    set
});

This required some modifications to the original code, but not too much. Instead of having an arbitrary wrapper type which implements Deref, I could implement Deref on the Lazy type. new would take the init closure, and that closure would need to be stored as a field on the type. The only change in get is to retrieve the init closure out of the struct, and make the function itself private:

pub struct Lazy<T, F = fn() -> T> {
    init: Once,
    init_fn: Cell<Option<F>>,
    data: Cell<MaybeUninit<T>>,
}

unsafe impl<T: Sync> Sync for Lazy<T> {}

impl<T, F: FnOnce() -> T> Lazy<T, F> {
    pub const fn new(init_fn: F) -> Self {
        Self {
            init: Once::new(),
            init_fn: Cell::new(Some(init_fn)),
            data: Cell::new(MaybeUninit::uninit()),
        }
    }

    fn get(&self) -> &T {
        self.init.call_once(move || unsafe {
            let init_fn = (*self.init_fn.as_ptr()).take().unwrap();
            (*self.data.as_ptr()).write(init_fn());
        });

        // SAFETY: Once call_once has completed, we know that the inner data has been written
        // and we can get a reference to it.
        unsafe { &*(*self.data.as_ptr()).as_ptr() }
    }
}

impl<T, F: FnOnce() -> T> Deref for Lazy<T, F> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        self.get()
    }
}