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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
//! A collection of lazy initialized values that are created by `Future`s.
//!
//! [OnceCell]'s API should be familiar to anyone who has used the
//! [`once_cell`](https://crates.io/crates/once_cell) crate or the proposed `std::cell::OnceCell`.
//! It provides an async version of a cell that can only be initialized once, permitting tasks to
//! wait on the initialization if it is already running instead of racing multiple initialization
//! tasks.
//!
//! Unlike threads, tasks can be cancelled at any point where they block.  [OnceCell] deals with
//! this by allowing another initializer to run if the task currently initializing the cell is
//! dropped.  This also allows for fallible initialization using [OnceCell::get_or_try_init], and
//! for the initializing `Future` to contain borrows or use references to thread-local data.
//!
//! [Lazy] takes the opposite approach: it wraps a single `Future` which is cooperatively run to
//! completion by any polling task.  This requires that the initialization function be independent
//! of the calling context, but will never restart an initializing function just because the
//! surrounding task was cancelled.
//!
//! # Overhead
//!
//! Both cells use two `usize`s to store state and do not retain any allocations after
//! initialization is complete.  [OnceCell] and [Lazy] only allocate if there is contention.
//!
//! # Features
//!
//! ## The `critical-section` feature
//!
//! If this feature is enabled, the [`critical-section`](https://crates.io/crates/critical-section)
//! crate is used instead of an `std` mutex.  You must depend on that crate and select a locking
//! implementation; see [its documentation](https://docs.rs/critical-section/) for details.
//!
//! ## The `unpin` feature
//!
//! This feature enables the `unpin` module which contains an alternative API for [Lazy] that does
//! not rely on pinning the object during initialization, even for futures that are not [Unpin].
//! In general, prefer the types in the crate root and, if needed, box futures to make them unpin.
//!
//! ## The `std` feature
//!
//! This is currently a no-op, but might in the future be used to expose APIs that depends on
//! types only in `std`.  It does *not* control the locking implementation.

#![cfg_attr(feature = "critical-section", no_std)]
extern crate alloc;

#[cfg(any(not(feature = "critical-section"), feature = "std"))]
extern crate std;

use alloc::{boxed::Box, vec, vec::Vec};

use core::{
    cell::UnsafeCell,
    convert::Infallible,
    future::Future,
    panic::{RefUnwindSafe, UnwindSafe},
    pin::Pin,
    ptr,
    sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
    task,
};

#[cfg(feature = "critical-section")]
struct Mutex<T> {
    data: UnsafeCell<T>,
    locked: core::sync::atomic::AtomicBool,
}

#[cfg(feature = "critical-section")]
impl<T> Mutex<T> {
    const fn new(data: T) -> Self {
        Mutex { data: UnsafeCell::new(data), locked: core::sync::atomic::AtomicBool::new(false) }
    }
}

#[cfg(not(feature = "critical-section"))]
use std::sync::Mutex;

#[cfg(feature = "critical-section")]
fn with_lock<T, R>(mutex: &Mutex<T>, f: impl FnOnce(&mut T) -> R) -> R {
    struct Guard<'a, T>(&'a Mutex<T>);
    impl<'a, T> Drop for Guard<'a, T> {
        fn drop(&mut self) {
            self.0.locked.store(false, Ordering::Relaxed);
        }
    }
    critical_section::with(|_| {
        if mutex.locked.swap(true, Ordering::Relaxed) {
            // Note: this can in theory happen if the delegated Clone impl on a Waker provided in
            // an initialization context turns around and tries to initialize the same cell.  This
            // is an absurd thing to do, but it's safe so we can't assume nobody will ever do it.
            panic!("Attempted reentrant locking");
        }
        let guard = Guard(mutex);
        // Safety: we just checked that we were the one to set `locked` to true, and the data in
        // this Mutex will only be accessed while the lock is true.  We use Relaxed memory ordering
        // instead of Acquire/Release because critical_section::with itself must provide an
        // Acquire/Release barrier around its closure, and also guarantees that there will not be
        // more than one such closure executing at a time.
        let rv = unsafe { f(&mut *mutex.data.get()) };
        drop(guard);
        rv
    })
}

#[cfg(not(feature = "critical-section"))]
fn with_lock<T, R>(mutex: &Mutex<T>, f: impl FnOnce(&mut T) -> R) -> R {
    f(&mut *mutex.lock().unwrap())
}

/// Types that do not rely on pinning during initialization.
///
/// This module is only built if the `unpin` crate feature is enabled.
///
/// This module contains [OnceFuture](unpin::OnceFuture) and its wrappers [Lazy](unpin::Lazy) and
/// [ConstLazy](unpin::ConstLazy), which provide lazy initialization without requiring the
/// resulting structure be pinned.
///
/// This is the API exposed by the 0.3 version of this crate for `Lazy`.
#[cfg(feature = "unpin")]
pub mod unpin;

/// A cell which can be written to only once.
///
/// This allows initialization using an async closure that borrows from its environment.
///
/// ```
/// # async fn run() {
/// use std::rc::Rc;
/// use std::sync::Arc;
/// use async_once_cell::OnceCell;
///
/// let non_send_value = Rc::new(4);
/// let shared = Arc::new(OnceCell::new());
///
/// let value : &i32 = shared.get_or_init(async {
///     *non_send_value
/// }).await;
/// assert_eq!(value, &4);
///
/// // A second init is not called
/// let second = shared.get_or_init(async {
///     unreachable!()
/// }).await;
/// assert_eq!(second, &4);
///
/// # }
/// ```
#[derive(Debug)]
pub struct OnceCell<T> {
    value: UnsafeCell<Option<T>>,
    inner: Inner,
}

// Safety: our UnsafeCell should be treated like an RwLock<T>
unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
unsafe impl<T: Send> Send for OnceCell<T> {}
impl<T> Unpin for OnceCell<T> {}
impl<T: RefUnwindSafe + UnwindSafe> RefUnwindSafe for OnceCell<T> {}
impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {}

/// Monomorphic portion of the state
#[derive(Debug)]
struct Inner {
    state: AtomicUsize,
    queue: AtomicPtr<Queue>,
}

/// Transient state during initialization
///
/// Unlike the sync OnceCell, this cannot be a linked list through stack frames, because Futures
/// can be freed at any point by any thread.  Instead, this structure is allocated on the heap
/// during the first initialization call and freed after the value is set (or when the OnceCell is
/// dropped, if the value never gets set).
struct Queue {
    wakers: Mutex<Option<Vec<task::Waker>>>,
}

/// This is somewhat like Arc<Queue>, but holds the refcount in Inner instead of Queue so it can be
/// freed once the cell's initialization is complete.
struct QueueRef<'a> {
    inner: &'a Inner,
    queue: *const Queue,
}
// Safety: the queue is a reference (only the lack of a valid lifetime requires it to be a pointer)
unsafe impl<'a> Sync for QueueRef<'a> {}
unsafe impl<'a> Send for QueueRef<'a> {}

#[derive(Debug)]
struct QuickInitGuard<'a>(&'a Inner);

/// A Future that waits for acquisition of a QueueHead
struct QueueWaiter<'a> {
    guard: Option<QueueRef<'a>>,
}

/// A guard for the actual initialization of the OnceCell
struct QueueHead<'a> {
    guard: QueueRef<'a>,
}

const NEW: usize = 0x0;
const QINIT_BIT: usize = 1 + (usize::MAX >> 2);
const READY_BIT: usize = 1 + (usize::MAX >> 1);

impl Inner {
    const fn new() -> Self {
        Inner { state: AtomicUsize::new(NEW), queue: AtomicPtr::new(ptr::null_mut()) }
    }

    const fn new_ready() -> Self {
        Inner { state: AtomicUsize::new(READY_BIT), queue: AtomicPtr::new(ptr::null_mut()) }
    }

    /// Initialize the queue (if needed) and return a waiter that can be polled to get a QueueHead
    /// that gives permission to initialize the OnceCell.
    ///
    /// The Queue referenced in the returned QueueRef will not be freed until the cell is populated
    /// and all references have been dropped.  If any references remain, further calls to
    /// initialize will return the existing queue.
    #[cold]
    fn initialize(&self, try_quick: bool) -> Result<QueueWaiter, QuickInitGuard> {
        if try_quick {
            if self
                .state
                .compare_exchange(NEW, QINIT_BIT, Ordering::Acquire, Ordering::Relaxed)
                .is_ok()
            {
                // On success, we know that there were no other QueueRef objects active, and we
                // just set QINIT_BIT which makes us the only party allowed to create a QueueHead.
                // This remains true even if the queue is created later.
                return Err(QuickInitGuard(self));
            }
        }

        // Increment the queue's reference count.  This ensures that queue won't be freed until we exit.
        let prev_state = self.state.fetch_add(1, Ordering::Acquire);

        // Note: unlike Arc, refcount overflow is impossible.  The only way to increment the
        // refcount is by calling poll on the Future returned by get_or_try_init, which is !Unpin.
        // The poll call requires a Pinned pointer to this Future, and the contract of Pin requires
        // Drop to be called on any !Unpin value that was pinned before the memory is reused.
        // Because the Drop impl of QueueRef decrements the refcount, an overflow would require
        // more than (usize::MAX / 4) QueueRef objects in memory, which is impossible as these
        // objects take up more than 4 bytes.

        let mut guard = QueueRef { inner: self, queue: self.queue.load(Ordering::Acquire) };

        if guard.queue.is_null() && prev_state & READY_BIT == 0 {
            let wakers = Mutex::new(None);

            // Race with other callers of initialize to create the queue
            let new_queue = Box::into_raw(Box::new(Queue { wakers }));

            match self.queue.compare_exchange(
                ptr::null_mut(),
                new_queue,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_null) => {
                    // Normal case: it was actually set.  The Release part of AcqRel orders this
                    // with all Acquires on the queue.
                    guard.queue = new_queue;
                }
                Err(actual) => {
                    // we lost the race, but we have the (non-null) value now.
                    guard.queue = actual;
                    // Safety: we just allocated it, and nobody else has seen it
                    unsafe {
                        drop(Box::from_raw(new_queue));
                    }
                }
            }
        }
        Ok(QueueWaiter { guard: Some(guard) })
    }

    fn set_ready(&self) {
        // This Release pairs with the Acquire any time we check READY_BIT, and ensures that the
        // writes to the cell's value are visible to the cell's readers.
        let prev_state = self.state.fetch_or(READY_BIT, Ordering::Release);

        debug_assert_eq!(prev_state & READY_BIT, 0, "Invalid state: someone else set READY_BIT");
    }
}

impl<'a> Drop for QueueRef<'a> {
    fn drop(&mut self) {
        // Release the reference to queue
        let prev_state = self.inner.state.fetch_sub(1, Ordering::Release);
        // Note: as of now, self.queue may be invalid

        let curr_state = prev_state - 1;
        if curr_state == READY_BIT || curr_state == READY_BIT | QINIT_BIT {
            // We just removed the only waiter on an initialized cell.  This means the
            // queue is no longer needed.  Acquire the queue again so we can free it.
            let queue = self.inner.queue.swap(ptr::null_mut(), Ordering::Acquire);
            if !queue.is_null() {
                // Safety: the last guard is being freed, and queue is only used by guard-holders.
                // Due to the swap, we are the only one who is freeing this particular queue.
                unsafe {
                    drop(Box::from_raw(queue));
                }
            }
        }
    }
}

impl<'a> Drop for QuickInitGuard<'a> {
    fn drop(&mut self) {
        let prev_state = self.0.state.load(Ordering::Relaxed);
        if prev_state == QINIT_BIT | READY_BIT || prev_state == QINIT_BIT {
            let target = prev_state & !QINIT_BIT;
            // Try to finish the fast path of initialization if possible.
            if self
                .0
                .state
                .compare_exchange(prev_state, target, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                // If init succeeded, the Release in set_ready already ordered the value.  If init
                // failed, we made no writes that need to be ordered and there are no waiters to
                // wake, so we can leave the state at NEW.

                if target == READY_BIT {
                    // It's possible (though unlikely) that someone created the queue but abandoned
                    // their QueueRef before we finished our poll, resulting in us not observing
                    // them.  No wakes are needed in this case because there are no waiting tasks,
                    // but we should still clean up the allocation.
                    let queue = self.0.queue.swap(ptr::null_mut(), Ordering::Relaxed);
                    if !queue.is_null() {
                        // Synchronize with both the fetch_sub that lowered the refcount and the
                        // queue initialization.
                        core::sync::atomic::fence(Ordering::Acquire);
                        // Safety: we observed no active QueueRefs, and queue is only used by
                        // guard-holders.  Due to the swap, we are the only one who is freeing this
                        // particular queue.
                        unsafe {
                            drop(Box::from_raw(queue));
                        }
                    }
                }
                return;
            }
        }

        // Slow path: get a guard, create the QueueHead we should have been holding, then drop it
        // so that the tasks are woken as intended.  This is needed regardless of if we succeeded
        // or not - either waiters need to run init themselves, or they need to read the value we
        // set.
        //
        // The guard is guaranteed to have been created with no QueueHead available because
        // QINIT_BIT is still set.
        let waiter = self.0.initialize(false).expect("Got a QuickInitGuard in slow init");
        let guard = waiter.guard.expect("No guard available even without polling");
        if guard.queue.is_null() {
            // The queue was already freed by someone else before we got our QueueRef (this must
            // have happened between the load of prev_state and initialize, because otherwise we
            // would have taken the fast path).  This implies that all other tasks have noticed
            // READY_BIT and do not need waking, so there is nothing left for us to do except
            // release our reference.
            drop(guard);
        } else {
            // Safety: the guard holds a place on the waiter list and we just checked that the
            // queue is non-null.  It will remain valid until guard is dropped.
            let queue = unsafe { &*guard.queue };

            with_lock(&queue.wakers, |lock| {
                // Ensure that nobody else can grab the QueueHead between when we release QINIT_BIT
                // and when our QueueHead is dropped.
                lock.get_or_insert_with(Vec::new);
                // Allow someone else to take the head position once we drop it.  Ordering is
                // handled by the Mutex.
                self.0.state.fetch_and(!QINIT_BIT, Ordering::Relaxed);
            });

            // Safety: we just took the head position, and we were the QuickInitGuard
            drop(QueueHead { guard })
        }
    }
}

impl Drop for Inner {
    fn drop(&mut self) {
        let queue = *self.queue.get_mut();
        if !queue.is_null() {
            // Safety: nobody else could have a reference
            unsafe {
                drop(Box::from_raw(queue));
            }
        }
    }
}

impl<'a> Future for QueueWaiter<'a> {
    type Output = Option<QueueHead<'a>>;
    fn poll(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
    ) -> task::Poll<Option<QueueHead<'a>>> {
        let guard = self.guard.as_ref().expect("Polled future after finished");

        // Fast path for waiters that get notified after the value is set
        let state = guard.inner.state.load(Ordering::Acquire);
        if state & READY_BIT != 0 {
            return task::Poll::Ready(None);
        }

        // Safety: the guard holds a place on the waiter list and we just checked that the state is
        // not ready, so the queue is non-null and will remain valid until guard is dropped.
        let queue = unsafe { &*guard.queue };
        let rv = with_lock(&queue.wakers, |lock| {
            // Another task might have called set_ready() and dropped its QueueHead between our
            // optimistic lock-free check and our lock acquisition.  Don't return a QueueHead unless we
            // know for sure that we are allowed to initialize.
            let state = guard.inner.state.load(Ordering::Acquire);
            if state & READY_BIT != 0 {
                return task::Poll::Ready(None);
            }

            match lock.as_mut() {
                None if state & QINIT_BIT == 0 => {
                    // take the head position and start a waker queue
                    *lock = Some(Vec::new());

                    task::Poll::Ready(Some(()))
                }
                None => {
                    // Someone else has a QuickInitGuard; they will wake us when they finish.
                    let waker = cx.waker().clone();
                    *lock = Some(vec![waker]);
                    task::Poll::Pending
                }
                Some(wakers) => {
                    // Wait for the QueueHead to be dropped
                    let my_waker = cx.waker();
                    for waker in wakers.iter() {
                        if waker.will_wake(my_waker) {
                            return task::Poll::Pending;
                        }
                    }
                    wakers.push(my_waker.clone());
                    task::Poll::Pending
                }
            }
        });

        // Safety: If rv is Ready/Some, we know:
        //  - we are holding a QueueRef (in guard) that prevents state from being 0
        //  - creating a new QuickInitGuard requires the state to be 0
        //  - we just checked QINIT_BIT and saw there isn't a QuickInitGuard active
        //  - the queue was None, meaning there are no current QueueHeads
        //  - we just set the queue to Some, claiming the head
        //
        // If rv is Ready/None, this is due to READY_BIT being set.
        // If rv is Pending, we have a waker in the queue.
        rv.map(|o| o.map(|()| QueueHead { guard: self.guard.take().unwrap() }))
    }
}

impl<'a> Drop for QueueHead<'a> {
    fn drop(&mut self) {
        // Safety: if queue is not null, then it is valid as long as the guard is alive
        if let Some(queue) = unsafe { self.guard.queue.as_ref() } {
            // Take the waker queue so the next QueueWaiter can make a new one
            let wakers = with_lock(&queue.wakers, |w| w.take())
                .expect("QueueHead dropped without a waker list");
            for waker in wakers {
                waker.wake();
            }
        }
    }
}

impl<T> OnceCell<T> {
    /// Creates a new empty cell.
    pub const fn new() -> Self {
        Self { value: UnsafeCell::new(None), inner: Inner::new() }
    }

    /// Creates a new cell with the given contents.
    ///
    /// If value is `None`, this is equivalent to `new()`.
    pub const fn new_with(value: Option<T>) -> Self {
        let inner = match value {
            Some(_) => Inner::new_ready(),
            None => Inner::new(),
        };
        Self { value: UnsafeCell::new(value), inner }
    }

    /// Gets the contents of the cell, initializing it with `init` if the cell was empty.
    ///
    /// Many tasks may call `get_or_init` concurrently with different initializing futures, but
    /// it is guaranteed that only one future will be executed as long as the resulting future is
    /// polled to completion.
    ///
    /// If `init` panics, the panic is propagated to the caller, and the cell remains uninitialized.
    ///
    /// If the Future returned by this function is dropped prior to completion, the cell remains
    /// uninitialized, and another `init` function will be started (if any are available).
    ///
    /// It is an error to reentrantly initialize the cell from `init`.  The current implementation
    /// deadlocks, but will recover if the offending task is dropped or if the future is actually
    /// able to proceed despite the reentrant call never returning.
    pub async fn get_or_init(&self, init: impl Future<Output = T>) -> &T {
        match self.get_or_try_init(async move { Ok::<T, Infallible>(init.await) }).await {
            Ok(t) => t,
            Err(e) => match e {},
        }
    }

    /// Gets the contents of the cell, initializing it with `init` if the cell was empty.   If the
    /// cell was empty and `init` failed, an error is returned.
    ///
    /// If `init` panics or returns an error, the panic or error is propagated to the caller, and
    /// the cell remains uninitialized.  In this case, another `init` function from a concurrent
    /// caller will be selected to execute, if one is available.
    ///
    /// If the Future returned by this function is dropped prior to completion, the cell remains
    /// uninitialized, and another `init` function will be started.
    ///
    /// It is an error to reentrantly initialize the cell from `init`.  The current implementation
    /// deadlocks, but will recover if the offending task is dropped or if the future is actually
    /// able to proceed despite the reentrant call never returning.
    pub async fn get_or_try_init<E>(
        &self,
        init: impl Future<Output = Result<T, E>>,
    ) -> Result<&T, E> {
        let state = self.inner.state.load(Ordering::Acquire);

        if state & READY_BIT == 0 {
            self.init_slow(state == NEW, init).await?;
        }

        // Safety: initialized on all paths
        Ok(unsafe { (&*self.value.get()).as_ref().unwrap() })
    }

    #[cold]
    async fn init_slow<E>(
        &self,
        try_quick: bool,
        init: impl Future<Output = Result<T, E>>,
    ) -> Result<(), E> {
        match self.inner.initialize(try_quick) {
            Err(guard) => {
                // Try to proceed assuming no contention.
                let value = init.await?;
                // Safety: the guard acts like QueueHead even if there is contention.
                unsafe {
                    *self.value.get() = Some(value);
                }
                self.inner.set_ready();
                drop(guard);
            }
            Ok(guard) => {
                if let Some(init_lock) = guard.await {
                    // We hold the QueueHead, so we know that nobody else has successfully run an init
                    // poll and that nobody else can start until it is dropped.  On error, panic, or
                    // drop of this Future, the head will be passed to another waiter.
                    let value = init.await?;

                    // Safety: We still hold the head, so nobody else can write to value
                    unsafe {
                        *self.value.get() = Some(value);
                    }
                    // mark the cell ready before giving up the head
                    init_lock.guard.inner.set_ready();
                    // drop of QueueHead notifies other Futures
                    // drop of QueueRef (might) free the Queue
                } else {
                    // someone initialized it while waiting on the queue
                }
            }
        }
        Ok(())
    }

    /// Gets the reference to the underlying value.
    ///
    /// Returns `None` if the cell is empty or being initialized. This method never blocks.
    pub fn get(&self) -> Option<&T> {
        let state = self.inner.state.load(Ordering::Acquire);

        if state & READY_BIT == 0 {
            None
        } else {
            unsafe { (&*self.value.get()).as_ref() }
        }
    }

    /// Gets a mutable reference to the underlying value.
    pub fn get_mut(&mut self) -> Option<&mut T> {
        self.value.get_mut().as_mut()
    }

    /// Takes the value out of this `OnceCell`, moving it back to an uninitialized state.
    pub fn take(&mut self) -> Option<T> {
        self.inner = Inner::new();
        self.value.get_mut().take()
    }

    /// Consumes the OnceCell, returning the wrapped value. Returns None if the cell was empty.
    pub fn into_inner(self) -> Option<T> {
        self.value.into_inner()
    }
}

impl<T> Default for OnceCell<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> From<T> for OnceCell<T> {
    fn from(value: T) -> Self {
        Self::new_with(Some(value))
    }
}

#[derive(Debug)]
enum LazyState<T, F> {
    Running(F),
    Ready(T),
}

/// A value which is computed on demand by running a future.
///
/// Unlike [OnceCell], if a task is cancelled, the initializing future's execution will be
/// continued by other (concurrent or future) callers of [Lazy::get].
///
/// ```
/// # async fn run() {
/// use std::sync::Arc;
/// use async_once_cell::Lazy;
///
/// struct Data {
///     id: u32,
/// }
///
/// let shared = Arc::pin(Lazy::new(async move {
///     Data { id: 4 }
/// }));
///
/// assert_eq!(shared.as_ref().get().await.id, 4);
/// # }
/// ```
#[derive(Debug)]
pub struct Lazy<T, F> {
    value: UnsafeCell<LazyState<T, F>>,
    inner: Inner,
}

// Safety: our UnsafeCell should be treated like an RwLock<(T, F)>
unsafe impl<T: Sync + Send, F: Sync + Send> Sync for Lazy<T, F> {}
unsafe impl<T: Send, F: Send> Send for Lazy<T, F> {}
impl<T: Unpin, F: Unpin> Unpin for Lazy<T, F> {}
impl<T: RefUnwindSafe + UnwindSafe, F: RefUnwindSafe + UnwindSafe> RefUnwindSafe for Lazy<T, F> {}
impl<T: UnwindSafe, F: UnwindSafe> UnwindSafe for Lazy<T, F> {}

impl<T, F> Lazy<T, F>
where
    F: Future<Output = T>,
{
    /// Creates a new lazy value with the given initializing future.
    pub const fn new(future: F) -> Self {
        Self::from_future(future)
    }

    /// Forces the evaluation of this lazy value and returns a reference to the result.
    ///
    /// The [Pin::static_ref] function may be useful if this is a static value.
    pub async fn get(self: Pin<&Self>) -> Pin<&T> {
        let state = self.inner.state.load(Ordering::Acquire);

        if state & READY_BIT == 0 {
            self.init_slow(state == NEW).await;
        }

        // Safety: initialized on all paths, and pinned like self
        unsafe {
            match &*self.value.get() {
                LazyState::Ready(v) => Pin::new_unchecked(v),
                _ => unreachable!(),
            }
        }
    }

    #[cold]
    async fn init_slow(self: Pin<&Self>, try_quick: bool) {
        match self.inner.initialize(try_quick) {
            Err(guard) => {
                let init = unsafe {
                    match &mut *self.value.get() {
                        LazyState::Running(f) => Pin::new_unchecked(f),
                        _ => unreachable!(),
                    }
                };
                let value = init.await;
                // Safety: the guard acts like QueueHead even if there is contention.
                // This overwrites the pinned future, dropping it in place
                unsafe {
                    *self.value.get() = LazyState::Ready(value);
                }
                self.inner.set_ready();
                drop(guard);
            }
            Ok(guard) => {
                if let Some(init_lock) = guard.await {
                    let init = unsafe {
                        match &mut *self.value.get() {
                            LazyState::Running(f) => Pin::new_unchecked(f),
                            _ => unreachable!(),
                        }
                    };
                    // We hold the QueueHead, so we know that nobody else has successfully run an init
                    // poll and that nobody else can start until it is dropped.  On error, panic, or
                    // drop of this Future, the head will be passed to another waiter.
                    let value = init.await;

                    // Safety: We still hold the head, so nobody else can write to value
                    // This overwrites the pinned future, dropping it in place
                    unsafe {
                        *self.value.get() = LazyState::Ready(value);
                    }
                    // mark the cell ready before giving up the head
                    init_lock.guard.inner.set_ready();
                    // drop of QueueHead notifies other Futures
                    // drop of QueueRef (might) free the Queue
                } else {
                    // someone initialized it while waiting on the queue
                }
            }
        }
    }
}

impl<T, F> Lazy<T, F>
where
    F: Future<Output = T> + Unpin,
{
    /// Forces the evaluation of this lazy value and returns a reference to the result.
    ///
    /// Unlike [Self::get], this does not require pinning the object.
    pub async fn get_unpin(&self) -> &T {
        // The get() function itself does not use the fact that T is pinned, and Pin::deref already
        // exposes a &T from Pin<&T> (although not with the right lifetime).
        unsafe { Pin::into_inner_unchecked(Pin::new_unchecked(self).get().await) }
    }
}

impl<T, F> Lazy<T, F> {
    /// Creates a new lazy value with the given initializing future.
    ///
    /// This is equivalent to [Self::new] but with no type bound.
    pub const fn from_future(future: F) -> Self {
        Self { value: UnsafeCell::new(LazyState::Running(future)), inner: Inner::new() }
    }

    /// Creates an already-initialized lazy value.
    pub const fn with_value(value: T) -> Self {
        Self { value: UnsafeCell::new(LazyState::Ready(value)), inner: Inner::new_ready() }
    }

    /// Gets the value without blocking or starting the initialization.
    pub fn try_get(&self) -> Option<&T> {
        let state = self.inner.state.load(Ordering::Acquire);

        if state & READY_BIT == 0 {
            None
        } else {
            match unsafe { &*self.value.get() } {
                LazyState::Ready(v) => Some(v),
                _ => unreachable!(),
            }
        }
    }

    /// Gets the value without blocking or starting the initialization.
    ///
    /// This requires mutable access to self, so rust's aliasing rules prevent any concurrent
    /// access and allow violating the usual rules for accessing this cell.
    pub fn try_get_mut(self: Pin<&mut Self>) -> Option<Pin<&mut T>> {
        unsafe {
            match self.get_unchecked_mut().value.get_mut() {
                LazyState::Ready(v) => Some(Pin::new_unchecked(v)),
                _ => None,
            }
        }
    }

    /// Gets the value without blocking or starting the initialization.
    ///
    /// This requires mutable access to self, so rust's aliasing rules prevent any concurrent
    /// access and allow violating the usual rules for accessing this cell.
    pub fn try_get_mut_unpin(&mut self) -> Option<&mut T> {
        match self.value.get_mut() {
            LazyState::Ready(v) => Some(v),
            _ => None,
        }
    }

    /// Gets the value if it was set.
    pub fn into_inner(self) -> Option<T> {
        match self.value.into_inner() {
            LazyState::Ready(v) => Some(v),
            _ => None,
        }
    }
}