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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

use core::cell::UnsafeCell;
use core::fmt;
use core::hint::spin_loop;
use core::sync::atomic::{AtomicUsize, Ordering};

/// A synchronization primitive which can be used to run a one-time global
/// initialization. Unlike its std equivalent, this is generalized so that the
/// closure returns a value and it is stored. Once therefore acts something like
/// a future, too.
pub struct Once<T> {
    state: AtomicUsize,
    data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
}

impl<T: fmt::Debug> fmt::Debug for Once<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.r#try() {
            Some(s) => write!(f, "Once {{ data: ")
                .and_then(|()| s.fmt(f))
                .and_then(|()| write!(f, "}}")),
            None => write!(f, "Once {{ <uninitialized> }}"),
        }
    }
}

// Same unsafe impls as `std::sync::RwLock`, because this also allows for
// concurrent reads.
unsafe impl<T: Send + Sync> Sync for Once<T> {}
unsafe impl<T: Send> Send for Once<T> {}

// Four states that a Once can be in, encoded into the lower bits of `state` in
// the Once structure.
const INCOMPLETE: usize = 0x0;
const RUNNING: usize = 0x1;
const COMPLETE: usize = 0x2;
const PANICKED: usize = 0x3;

use core::hint::unreachable_unchecked as unreachable;

impl<T> Once<T> {
    /// Initialization constant of `Once`.
    pub const INIT: Self = Once {
        state: AtomicUsize::new(INCOMPLETE),
        data: UnsafeCell::new(None),
    };

    /// Creates a new `Once` value.
    pub const fn new() -> Once<T> {
        Self::INIT
    }

    fn force_get<'a>(&'a self) -> &'a T {
        match unsafe { &*self.data.get() }.as_ref() {
            None => unsafe { unreachable() },
            Some(p) => p,
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). The
    /// returned pointer will point to the result from the closure that was
    /// run.
    pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
    where
        F: FnOnce() -> T,
    {
        let mut status = self.state.load(Ordering::SeqCst);

        if status == INCOMPLETE {
            status = match self.state.compare_exchange(
                INCOMPLETE,
                RUNNING,
                Ordering::SeqCst,
                Ordering::SeqCst,
            ) {
                Ok(status) => {
                    debug_assert_eq!(
                        status, INCOMPLETE,
                        "if compare_exchange succeeded, previous status must be incomplete",
                    );
                    // We init
                    // We use a guard (Finish) to catch panics caused by builder
                    let mut finish = Finish {
                        state: &self.state,
                        panicked: true,
                    };
                    unsafe { *self.data.get() = Some(builder()) };
                    finish.panicked = false;

                    self.state.store(COMPLETE, Ordering::SeqCst);

                    // This next line is strictly an optimization
                    return self.force_get();
                }
                Err(status) => status,
            }
        }

        loop {
            match status {
                INCOMPLETE => unreachable!(),
                RUNNING => {
                    // We spin
                    spin_loop();
                    status = self.state.load(Ordering::SeqCst)
                }
                PANICKED => panic!("Once has panicked"),
                COMPLETE => return self.force_get(),
                _ => unsafe { unreachable() },
            }
        }
    }

    /// Returns a pointer iff the `Once` was previously initialized
    pub fn r#try<'a>(&'a self) -> Option<&'a T> {
        match self.state.load(Ordering::SeqCst) {
            COMPLETE => Some(self.force_get()),
            _ => None,
        }
    }

    /// Like try, but will spin if the `Once` is in the process of being
    /// initialized
    pub fn wait<'a>(&'a self) -> Option<&'a T> {
        loop {
            match self.state.load(Ordering::SeqCst) {
                INCOMPLETE => return None,

                RUNNING => {
                    spin_loop() // We spin
                }
                COMPLETE => return Some(self.force_get()),
                PANICKED => panic!("Once has panicked"),
                _ => unsafe { unreachable() },
            }
        }
    }
}

struct Finish<'a> {
    state: &'a AtomicUsize,
    panicked: bool,
}

impl<'a> Drop for Finish<'a> {
    fn drop(&mut self) {
        if self.panicked {
            self.state.store(PANICKED, Ordering::SeqCst);
        }
    }
}