diff --git a/embassy-sync/src/lazy_lock.rs b/embassy-sync/src/lazy_lock.rs index 2b5742491..cf88bfdf8 100644 --- a/embassy-sync/src/lazy_lock.rs +++ b/embassy-sync/src/lazy_lock.rs @@ -1,7 +1,7 @@ //! Synchronization primitive for initializing a value once, allowing others to get a reference to the value. -use core::cell::Cell; -use core::mem::MaybeUninit; +use core::cell::UnsafeCell; +use core::mem::ManuallyDrop; use core::sync::atomic::{AtomicBool, Ordering}; /// The `LazyLock` is a synchronization primitive that allows for @@ -23,8 +23,12 @@ use core::sync::atomic::{AtomicBool, Ordering}; /// ``` pub struct LazyLock T> { init: AtomicBool, - init_fn: Cell>, - data: Cell>, + data: UnsafeCell>, +} + +union Data { + value: ManuallyDrop, + f: ManuallyDrop, } unsafe impl Sync for LazyLock {} @@ -34,8 +38,9 @@ impl T> LazyLock { pub const fn new(init_fn: F) -> Self { Self { init: AtomicBool::new(false), - init_fn: Cell::new(Some(init_fn)), - data: Cell::new(MaybeUninit::zeroed()), + data: UnsafeCell::new(Data { + f: ManuallyDrop::new(init_fn), + }), } } @@ -44,7 +49,7 @@ impl T> LazyLock { #[inline] pub fn get(&self) -> &T { self.ensure_init_fast(); - unsafe { (*self.data.as_ptr()).assume_init_ref() } + unsafe { &(*self.data.get()).value } } /// Consume the `LazyLock`, returning the underlying value. The @@ -53,7 +58,10 @@ impl T> LazyLock { #[inline] pub fn into_inner(self) -> T { self.ensure_init_fast(); - unsafe { self.data.into_inner().assume_init() } + let this = ManuallyDrop::new(self); + let data = unsafe { core::ptr::read(&this.data) }.into_inner(); + + ManuallyDrop::into_inner(unsafe { data.value }) } /// Initialize the `LazyLock` if it has not been initialized yet. @@ -75,10 +83,23 @@ impl T> LazyLock { fn ensure_init(&self) { critical_section::with(|_| { if !self.init.load(Ordering::Acquire) { - let init_fn = self.init_fn.take().unwrap(); - self.data.set(MaybeUninit::new(init_fn())); + let data = unsafe { &mut *self.data.get() }; + let f = unsafe { ManuallyDrop::take(&mut data.f) }; + let value = f(); + data.value = ManuallyDrop::new(value); + self.init.store(true, Ordering::Release); } }); } } + +impl Drop for LazyLock { + fn drop(&mut self) { + if self.init.load(Ordering::Acquire) { + unsafe { ManuallyDrop::drop(&mut self.data.get_mut().value) }; + } else { + unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) }; + } + } +}