Skip to main content

miden_utils_sync/
once_lock.rs

1#[cfg(not(feature = "std"))]
2use core::fmt;
3#[cfg(feature = "std")]
4use std::sync::OnceLock;
5
6#[cfg(not(feature = "std"))]
7use once_cell::race::OnceBox;
8
9/// A wrapper around `once_cell::race::OnceBox` that adds a `take()` method for cache invalidation.
10///
11/// `OnceBox` is designed to be write-once, but we need to be able to invalidate the cache
12/// when the MAST forest is mutated. Since `take()` is only called with `&mut self`, we can
13/// safely replace the entire `OnceBox` with a new empty one.
14#[cfg(not(feature = "std"))]
15pub struct OnceLockCompat<T> {
16    inner: OnceBox<T>,
17}
18
19#[cfg(not(feature = "std"))]
20impl<T> OnceLockCompat<T> {
21    /// Creates a new empty `OnceLockCompat`.
22    pub const fn new() -> Self {
23        Self { inner: OnceBox::new() }
24    }
25
26    /// Gets the value if initialized, or initializes it with the provided closure.
27    ///
28    /// If multiple threads call this simultaneously, they may both execute the closure,
29    /// but only one value will be stored. The losing thread's value will be immediately dropped.
30    pub fn get_or_init<F>(&self, f: F) -> &T
31    where
32        F: FnOnce() -> T,
33    {
34        self.inner.get_or_init(|| alloc::boxed::Box::new(f()))
35    }
36
37    /// Takes the value out of the `OnceLockCompat`, leaving it empty.
38    ///
39    /// Returns `Some(T)` if the value was present, or `None` if it was not initialized.
40    ///
41    /// Note: For the no-std implementation, we can't extract the value from `OnceBox`,
42    /// so we just replace it with a new empty one and return `None`.
43    pub fn take(&mut self) -> Option<T> {
44        // Replace the inner OnceBox with a new empty one.
45        // This invalidates the cache by making the next `get_or_init` recompute.
46        // We can't extract the value from OnceBox, so we just discard it.
47        self.inner = OnceBox::new();
48        None
49    }
50}
51
52#[cfg(not(feature = "std"))]
53impl<T> Default for OnceLockCompat<T> {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59#[cfg(not(feature = "std"))]
60impl<T> Clone for OnceLockCompat<T> {
61    /// Cloning an `OnceLockCompat` creates a new empty `OnceLockCompat`.
62    ///
63    /// The cached value is not cloned because it's derived data that can be
64    /// recomputed. This matches the semantics expected for cached/memoized values
65    /// where the cache is an optimization detail, not part of the logical state.
66    fn clone(&self) -> Self {
67        Self::new()
68    }
69}
70
71#[cfg(not(feature = "std"))]
72impl<T: fmt::Debug> fmt::Debug for OnceLockCompat<T> {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        // Delegate to OnceBox's Debug implementation
75        self.inner.fmt(f)
76    }
77}
78
79/// Type alias for std builds - uses `std::sync::OnceLock` directly
80#[cfg(feature = "std")]
81pub type OnceLockCompat<T> = OnceLock<T>;