1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! The `init_once` crate provides a mechanic to attempt to read a value without
//! blocking the caller, in case it is being initialized concurrently. Such an
//! abstraction might be useful in cache implementations whose consumers might
//! not want to block on the cache to fill up with data.

#![no_std]
#![deny(missing_docs)]

use core::cell::UnsafeCell;
use core::mem::{needs_drop, MaybeUninit};
use core::sync::atomic::{self, AtomicUsize};
use core::task::Poll;

mod init_once_state {
    pub const EMPTY: usize = 0;
    pub const INITIALIZING: usize = 1;
    pub const INITIALIZED: usize = 2;
}

/// Initialization state of an [`InitOnce`] instance.
pub enum InitState<'a, T> {
    /// The inner value is currently being initialized by another caller.
    Initializing,
    /// The inner value is initialized.
    Initialized(&'a T),
    /// The inner value requires initialization via [`PollInit`].
    Polling(PollInit<'a, T>),
}

/// Lazily initialize a value of some arbitrary type `T`.
/// Reading the value doesn't block the caller, if it is
/// being initialized concurrently.
pub struct InitOnce<T> {
    cell: UnsafeCell<MaybeUninit<T>>,
    state: AtomicUsize,
}

/// Polling mechanism to initialize a value contained in some
/// [`InitOnce`] instance.
pub struct PollInit<'a, T> {
    init_once: &'a InitOnce<T>,
}

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

impl<T> Drop for InitOnce<T> {
    // NB: it is guaranteed that at least one thread calls `Drop`, since
    // we must block to initialize from at least one thread
    fn drop(&mut self) {
        if needs_drop::<T>()
            && self.state.load(atomic::Ordering::SeqCst) == init_once_state::INITIALIZED
        {
            // SAFETY: the value is initialized, so we can drop it
            unsafe {
                self.cell.get_mut().assume_init_drop();
            }
        }
    }
}

impl<T> Default for InitOnce<T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl<T> InitOnce<T> {
    /// Create an uninitialized [`InitOnce`].
    pub const fn new() -> Self {
        Self {
            cell: UnsafeCell::new(MaybeUninit::uninit()),
            state: AtomicUsize::new(init_once_state::EMPTY),
        }
    }

    #[must_use]
    fn poll_init_begin(&self) -> PollInit<'_, T> {
        PollInit { init_once: self }
    }

    /// Query the state of an [`InitOnce`] instance.
    ///
    /// If the current state is [`InitState::Polling`], the caller is
    /// responsible for polling the init function to completion.
    #[must_use = "The state of an InitOnce (i.e. InitState) must always be consumed. If you do \
             not poll the value initializer to completion, the value will never be initialized."]
    pub fn state(&self) -> InitState<'_, T> {
        self.state
            .compare_exchange(
                init_once_state::EMPTY,
                init_once_state::INITIALIZING,
                atomic::Ordering::SeqCst,
                atomic::Ordering::SeqCst,
            )
            .map_or_else(
                |current_value| match current_value {
                    init_once_state::INITIALIZING => InitState::Initializing,
                    init_once_state::INITIALIZED => {
                        InitState::Initialized(unsafe { (*self.cell.get()).assume_init_ref() })
                    }
                    _ => unreachable!(),
                },
                |_| InitState::Polling(self.poll_init_begin()),
            )
    }

    /// Attempt to initialize this [`InitOnce`] with the value returned by `init`.
    pub fn try_init<F>(&self, mut init: F) -> Option<&T>
    where
        F: FnMut() -> T,
    {
        match self.state() {
            InitState::Initialized(value) => Some(value),
            InitState::Initializing => None,
            InitState::Polling(poller) => match poller.poll_init(|| Poll::Ready(init())) {
                Poll::Ready(value) => Some(value),
                Poll::Pending => unreachable!(),
            },
        }
    }
}

impl<'init_once, T> PollInit<'init_once, T> {
    /// Check if the value returned by `init` is ready.
    pub fn poll_init<F>(&self, mut init: F) -> Poll<&'init_once T>
    where
        F: FnMut() -> Poll<T>,
    {
        let value = core::task::ready!(init());

        let slot = unsafe { (*self.init_once.cell.get()).as_mut_ptr() };
        unsafe {
            core::ptr::write(slot, value);
        }

        self.init_once
            .state
            .store(init_once_state::INITIALIZED, atomic::Ordering::SeqCst);

        Poll::Ready(unsafe { (*self.init_once.cell.get()).assume_init_ref() })
    }
}