From 05e0f128462ebdacd3dffc27f74502913d782589 Mon Sep 17 00:00:00 2001
From: Samuel Tardieu <sam@rfc1149.net>
Date: Sat, 27 Jul 2024 11:49:02 +0200
Subject: [PATCH] embassy-sync: add LazyLock

`LazyLock` is inspired by Rust 1.80.0's `std::sync::LazyLock` type.
---
 embassy-sync/CHANGELOG.md     |  2 +
 embassy-sync/README.md        |  1 +
 embassy-sync/src/lazy_lock.rs | 84 +++++++++++++++++++++++++++++++++++
 embassy-sync/src/lib.rs       |  1 +
 4 files changed, 88 insertions(+)
 create mode 100644 embassy-sync/src/lazy_lock.rs

diff --git a/embassy-sync/CHANGELOG.md b/embassy-sync/CHANGELOG.md
index a283adc0c..8f2d26fe0 100644
--- a/embassy-sync/CHANGELOG.md
+++ b/embassy-sync/CHANGELOG.md
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Add LazyLock sync primitive.
+
 ## 0.6.0 - 2024-05-29
 
 - Add `capacity`, `free_capacity`, `clear`, `len`, `is_empty` and `is_full` functions to `Channel`.
diff --git a/embassy-sync/README.md b/embassy-sync/README.md
index 2c1c0cf68..dec1fbc32 100644
--- a/embassy-sync/README.md
+++ b/embassy-sync/README.md
@@ -13,6 +13,7 @@ Synchronization primitives and data structures with async support:
 - [`WakerRegistration`](waitqueue::WakerRegistration) - Utility to register and wake a `Waker`.
 - [`AtomicWaker`](waitqueue::AtomicWaker) - A variant of `WakerRegistration` accessible using a non-mut API.
 - [`MultiWakerRegistration`](waitqueue::MultiWakerRegistration) - Utility registering and waking multiple `Waker`'s.
+- [`LazyLock`](lazy_lock::LazyLock) - A value which is initialized on the first access
 
 ## Interoperability
 
diff --git a/embassy-sync/src/lazy_lock.rs b/embassy-sync/src/lazy_lock.rs
new file mode 100644
index 000000000..2b5742491
--- /dev/null
+++ b/embassy-sync/src/lazy_lock.rs
@@ -0,0 +1,84 @@
+//! 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::sync::atomic::{AtomicBool, Ordering};
+
+/// The `LazyLock` is a synchronization primitive that allows for
+/// initializing a value once, and allowing others to obtain a
+/// reference to the value. This is useful for lazy initialization of
+/// a static value.
+///
+/// # Example
+/// ```
+/// use futures_executor::block_on;
+/// use embassy_sync::lazy_lock::LazyLock;
+///
+/// // Define a static value that will be lazily initialized
+/// // at runtime at the first access.
+/// static VALUE: LazyLock<u32> = LazyLock::new(|| 20);
+///
+/// let reference = VALUE.get();
+/// assert_eq!(reference, &20);
+/// ```
+pub struct LazyLock<T, F = fn() -> T> {
+    init: AtomicBool,
+    init_fn: Cell<Option<F>>,
+    data: Cell<MaybeUninit<T>>,
+}
+
+unsafe impl<T, F> Sync for LazyLock<T, F> {}
+
+impl<T, F: FnOnce() -> T> LazyLock<T, F> {
+    /// Create a new uninitialized `StaticLock`.
+    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()),
+        }
+    }
+
+    /// Get a reference to the underlying value, initializing it if it
+    /// has not been done already.
+    #[inline]
+    pub fn get(&self) -> &T {
+        self.ensure_init_fast();
+        unsafe { (*self.data.as_ptr()).assume_init_ref() }
+    }
+
+    /// Consume the `LazyLock`, returning the underlying value. The
+    /// initialization function will be called if it has not been
+    /// already.
+    #[inline]
+    pub fn into_inner(self) -> T {
+        self.ensure_init_fast();
+        unsafe { self.data.into_inner().assume_init() }
+    }
+
+    /// Initialize the `LazyLock` if it has not been initialized yet.
+    /// This function is a fast track to [`Self::ensure_init`]
+    /// which does not require a critical section in most cases when
+    /// the value has been initialized already.
+    /// When this function returns, `self.data` is guaranteed to be
+    /// initialized and visible on the current core.
+    #[inline]
+    fn ensure_init_fast(&self) {
+        if !self.init.load(Ordering::Acquire) {
+            self.ensure_init();
+        }
+    }
+
+    /// Initialize the `LazyLock` if it has not been initialized yet.
+    /// When this function returns, `self.data` is guaranteed to be
+    /// initialized and visible on the current core.
+    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()));
+                self.init.store(true, Ordering::Release);
+            }
+        });
+    }
+}
diff --git a/embassy-sync/src/lib.rs b/embassy-sync/src/lib.rs
index a5eee8d02..014bf1d06 100644
--- a/embassy-sync/src/lib.rs
+++ b/embassy-sync/src/lib.rs
@@ -12,6 +12,7 @@ mod ring_buffer;
 
 pub mod blocking_mutex;
 pub mod channel;
+pub mod lazy_lock;
 pub mod mutex;
 pub mod once_lock;
 pub mod pipe;