Remove critical-section requirement for no_std with atomics (#4322)

Author: daxpedda

.github/workflows/main.yml

@@ -121,8 +121,8 @@ jobs:
     - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p wasm-bindgen -Zbuild-std=core,alloc -- -D warnings
     - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p js-sys -Zbuild-std=core,alloc -- -D warnings
     - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p web-sys -Zbuild-std=core,alloc -- -D warnings
-    - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p wasm-bindgen-futures --features once_cell/critical-section -Zbuild-std=core,alloc -- -D warnings
-    - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p wasm-bindgen-test --features once_cell/critical-section -Zbuild-std=core,alloc -- -D warnings
+    - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p wasm-bindgen-futures -Zbuild-std=core,alloc -- -D warnings
+    - run: cargo clippy --no-deps --no-default-features --target wasm32-unknown-unknown -p wasm-bindgen-test -Zbuild-std=core,alloc -- -D warnings
 
   # Run `cargo clippy` over the project
   clippy_project:

CHANGELOG.md

@@ -13,6 +13,9 @@
 * Add clear error message to communicate new feature resolver version requirements.
   [#4312](https://github.com/rustwasm/wasm-bindgen/pull/4312)
 
+* Remove `once_cell/critical-section` requirement for `no_std` with atomics.
+  [#4322](https://github.com/rustwasm/wasm-bindgen/pull/4322)
+
 ### Fixed
 
 * Fix macro-hygiene for calls to `std::thread_local!`.

crates/backend/src/codegen.rs

@@ -201,8 +201,8 @@ impl TryToTokens for ast::LinkToModule {
                 #program
                 #extern_fn
 
-                static __VAL: #wasm_bindgen::__rt::once_cell::sync::Lazy<#wasm_bindgen::__rt::alloc::string::String> =
-                    #wasm_bindgen::__rt::once_cell::sync::Lazy::new(|| unsafe {
+                static __VAL: #wasm_bindgen::__rt::LazyLock<#wasm_bindgen::__rt::alloc::string::String> =
+                    #wasm_bindgen::__rt::LazyLock::new(|| unsafe {
                         <#wasm_bindgen::__rt::alloc::string::String as #wasm_bindgen::convert::FromWasmAbi>::from_abi(#name().join())
                     });
 

crates/test/Cargo.toml

@@ -16,7 +16,6 @@ std = ["wasm-bindgen/std", "js-sys/std", "wasm-bindgen-futures/std", "scoped-tls
 [dependencies]
 gg-alloc = { version = "1.0", optional = true }
 js-sys = { path = '../js-sys', version = '=0.3.74', default-features = false }
-once_cell = { version = "1.12", default-features = false }
 scoped-tls = { version = "1.0", optional = true }
 wasm-bindgen = { path = '../..', version = '=0.2.97', default-features = false }
 wasm-bindgen-futures = { path = '../futures', version = '=0.4.47', default-features = false }

src/lib.rs

@@ -1572,6 +1572,8 @@ pub mod __rt {
     use core::convert::Infallible;
     use core::mem;
     use core::ops::{Deref, DerefMut};
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    use core::sync::atomic::{AtomicU8, Ordering};
 
     pub extern crate alloc;
     pub extern crate core;
@@ -1582,16 +1584,6 @@ pub mod __rt {
     use alloc::boxed::Box;
     use alloc::rc::Rc;
 
-    pub mod once_cell {
-        #[cfg(any(target_feature = "atomics", feature = "std"))]
-        pub use once_cell::*;
-
-        #[cfg(all(not(target_feature = "atomics"), not(feature = "std")))]
-        pub mod sync {
-            pub use super::super::LazyCell as Lazy;
-        }
-    }
-
     /// Wrapper around [`::once_cell::unsync::Lazy`] adding some compatibility methods with
     /// [`std::thread::LocalKey`] and adding `Send + Sync` when `atomics` is not enabled.
     #[cfg(not(feature = "std"))]
@@ -1633,6 +1625,89 @@ pub mod __rt {
         }
     }
 
+    #[cfg(feature = "std")]
+    pub use once_cell::sync::Lazy as LazyLock;
+
+    #[cfg(all(not(target_feature = "atomics"), not(feature = "std")))]
+    pub use LazyCell as LazyLock;
+
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    pub struct LazyLock<T, F = fn() -> T> {
+        state: AtomicU8,
+        data: UnsafeCell<Data<T, F>>,
+    }
+
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    enum Data<T, F> {
+        Value(T),
+        Init(F),
+    }
+
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    unsafe impl<T, F> Sync for LazyLock<T, F> {}
+
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    unsafe impl<T, F> Send for LazyLock<T, F> {}
+
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    impl<T, F> LazyLock<T, F> {
+        const STATE_UNINIT: u8 = 0;
+        const STATE_INITIALIZING: u8 = 1;
+        const STATE_INIT: u8 = 2;
+
+        pub const fn new(init: F) -> LazyLock<T, F> {
+            Self {
+                state: AtomicU8::new(Self::STATE_UNINIT),
+                data: UnsafeCell::new(Data::Init(init)),
+            }
+        }
+    }
+
+    #[cfg(all(target_feature = "atomics", not(feature = "std")))]
+    impl<T> Deref for LazyLock<T> {
+        type Target = T;
+
+        fn deref(&self) -> &T {
+            let mut state = self.state.load(Ordering::Acquire);
+
+            loop {
+                match state {
+                    Self::STATE_INIT => {
+                        let Data::Value(value) = (unsafe { &*self.data.get() }) else {
+                            unreachable!()
+                        };
+                        return value;
+                    }
+                    Self::STATE_UNINIT => {
+                        if let Err(new_state) = self.state.compare_exchange_weak(
+                            Self::STATE_UNINIT,
+                            Self::STATE_INITIALIZING,
+                            Ordering::Acquire,
+                            Ordering::Relaxed,
+                        ) {
+                            state = new_state;
+                            continue;
+                        }
+
+                        let data = unsafe { &mut *self.data.get() };
+                        let Data::Init(init) = data else {
+                            unreachable!()
+                        };
+                        *data = Data::Value(init());
+                        self.state.store(Self::STATE_INIT, Ordering::Release);
+                        state = Self::STATE_INIT;
+                    }
+                    Self::STATE_INITIALIZING => {
+                        // TODO: Block here if possible. This would require
+                        // detecting if we can in the first place.
+                        state = self.state.load(Ordering::Acquire);
+                    }
+                    _ => unreachable!(),
+                }
+            }
+        }
+    }
+
     #[inline]
     pub fn assert_not_null<T>(s: *mut T) {
         if s.is_null() {