[naga]: Switch off of LazyLock to once_cell::racy::OnceBox (#7587)

* Switch off of `LazyLock` to a custom `RacyLock`

* Formatting

* Switch to `OnceBox` internally

* Clippy

* Simplify `trunc` usage

Co-Authored-By: Connor Fitzgerald <[email protected]>

* Switch to `expect`

Co-Authored-By: Connor Fitzgerald <[email protected]>

* Fix expectation conditions

Co-Authored-By: Connor Fitzgerald <[email protected]>

* Update documentation

---------

Co-authored-by: Connor Fitzgerald <[email protected]>

Author: bushrat011899

naga/build.rs

@@ -6,5 +6,7 @@ fn main() {
         msl_out: { any(feature = "msl-out", all(target_vendor = "apple", feature = "msl-out-if-target-apple")) },
         spv_out: { feature = "spv-out" },
         wgsl_out: { feature = "wgsl-out" },
+        std: { any(test, spv_out, feature = "spv-in", feature = "wgsl-in", feature = "stderr") },
+        no_std: { not(std) },
     }
 }

naga/src/back/glsl/keywords.rs

@@ -1,4 +1,4 @@
-use std::sync::LazyLock;
+use crate::racy_lock::RacyLock;
 
 use hashbrown::HashSet;
 
@@ -499,7 +499,7 @@ pub const RESERVED_KEYWORDS: &[&str] = &[
 /// significant time during [`Namer::reset`](crate::proc::Namer::reset).
 ///
 /// See <https://github.com/gfx-rs/wgpu/pull/7338> for benchmarks.
-pub static RESERVED_KEYWORD_SET: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
+pub static RESERVED_KEYWORD_SET: RacyLock<HashSet<&'static str>> = RacyLock::new(|| {
     let mut set = HashSet::default();
     set.reserve(RESERVED_KEYWORDS.len());
     for &word in RESERVED_KEYWORDS {

naga/src/back/hlsl/keywords.rs

@@ -1,4 +1,4 @@
-use std::sync::LazyLock;
+use crate::racy_lock::RacyLock;
 
 use hashbrown::HashSet;
 
@@ -924,7 +924,7 @@ pub const TYPES: &[&str] = &{
 /// significant time during [`Namer::reset`](crate::proc::Namer::reset).
 ///
 /// See <https://github.com/gfx-rs/wgpu/pull/7338> for benchmarks.
-pub static RESERVED_SET: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
+pub static RESERVED_SET: RacyLock<HashSet<&'static str>> = RacyLock::new(|| {
     let mut set = HashSet::default();
     set.reserve(RESERVED.len() + TYPES.len());
     for &word in RESERVED {

naga/src/back/msl/keywords.rs

@@ -1,4 +1,4 @@
-use std::sync::LazyLock;
+use crate::racy_lock::RacyLock;
 
 use hashbrown::HashSet;
 
@@ -360,7 +360,7 @@ pub const RESERVED: &[&str] = &[
 /// significant time during [`Namer::reset`](crate::proc::Namer::reset).
 ///
 /// See <https://github.com/gfx-rs/wgpu/pull/7338> for benchmarks.
-pub static RESERVED_SET: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
+pub static RESERVED_SET: RacyLock<HashSet<&'static str>> = RacyLock::new(|| {
     let mut set = HashSet::default();
     set.reserve(RESERVED.len());
     for &word in RESERVED {

naga/src/back/pipeline_constants.rs

@@ -16,6 +16,9 @@ use crate::{
     Span, Statement, TypeInner, WithSpan,
 };
 
+#[cfg(no_std)]
+use num_traits::float::FloatCore as _;
+
 #[derive(Error, Debug, Clone)]
 #[cfg_attr(test, derive(PartialEq))]
 pub enum PipelineConstantError {

naga/src/common/wgsl/to_wgsl.rs

@@ -14,7 +14,7 @@ use alloc::string::{String, ToString};
 ///
 /// - If a type's WGSL form requires dynamic formatting, so that
 ///   returning a `&'static str` isn't feasible, consider implementing
-///   [`std::fmt::Display`] on some wrapper type instead.
+///   [`core::fmt::Display`] on some wrapper type instead.
 pub trait ToWgsl: Sized {
     /// Return WGSL source code representation of `self`.
     fn to_wgsl(self) -> &'static str;
@@ -32,7 +32,7 @@ pub trait ToWgsl: Sized {
 ///
 /// - If a type's WGSL form requires dynamic formatting, so that
 ///   returning a `&'static str` isn't feasible, consider implementing
-///   [`std::fmt::Display`] on some wrapper type instead.
+///   [`core::fmt::Display`] on some wrapper type instead.
 pub trait TryToWgsl: Sized {
     /// Return the WGSL form of `self` as a `'static` string.
     ///

naga/src/keywords/wgsl.rs

@@ -4,7 +4,7 @@ Keywords for [WGSL][wgsl] (WebGPU Shading Language).
 [wgsl]: https://gpuweb.github.io/gpuweb/wgsl.html
 */
 
-use std::sync::LazyLock;
+use crate::racy_lock::RacyLock;
 
 use hashbrown::HashSet;
 
@@ -238,7 +238,7 @@ pub const RESERVED: &[&str] = &[
 /// significant time during [`Namer::reset`](crate::proc::Namer::reset).
 ///
 /// See <https://github.com/gfx-rs/wgpu/pull/7338> for benchmarks.
-pub static RESERVED_SET: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
+pub static RESERVED_SET: RacyLock<HashSet<&'static str>> = RacyLock::new(|| {
     let mut set = HashSet::default();
     set.reserve(RESERVED.len());
     for &word in RESERVED {

naga/src/lib.rs

@@ -100,20 +100,7 @@ void main() {
 )]
 #![no_std]
 
-#[cfg(any(
-    test,
-    spv_out,
-
-    // Need OnceLock
-    hlsl_out,
-    msl_out,
-    wgsl_out,
-
-    feature = "spv-in",
-    feature = "wgsl-in",
-
-    feature = "stderr",
-))]
+#[cfg(std)]
 extern crate std;
 
 extern crate alloc;
@@ -130,6 +117,7 @@ pub mod ir;
 pub mod keywords;
 mod non_max_u32;
 pub mod proc;
+mod racy_lock;
 mod span;
 pub mod valid;
 

naga/src/racy_lock.rs

@@ -0,0 +1,41 @@
+#![cfg_attr(
+    not(any(glsl_out, hlsl_out, msl_out, feature = "wgsl-in", wgsl_out)),
+    expect(
+        dead_code,
+        reason = "RacyLock is only required for the above configurations"
+    )
+)]
+
+use alloc::boxed::Box;
+use once_cell::race::OnceBox;
+
+/// An alternative to [`LazyLock`] based on [`OnceBox`].
+///
+/// [`LazyLock`]: https://doc.rust-lang.org/stable/std/sync/struct.LazyLock.html
+pub struct RacyLock<T: 'static> {
+    inner: OnceBox<T>,
+    init: fn() -> T,
+}
+
+impl<T: 'static> RacyLock<T> {
+    /// Creates a new [`RacyLock`], which will initialize using the provided `init` function.
+    pub const fn new(init: fn() -> T) -> Self {
+        Self {
+            inner: OnceBox::new(),
+            init,
+        }
+    }
+
+    /// Loads the internal value, initializing it if required.
+    pub fn get(&self) -> &T {
+        self.inner.get_or_init(|| Box::new((self.init)()))
+    }
+}
+
+impl<T: 'static> core::ops::Deref for RacyLock<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        self.get()
+    }
+}