std: use a more efficient `Once` on platforms without threads

4 files changed, 94 insertions, 0 deletions

diff --git a/library/std/src/sys/unsupported/mod.rs b/library/std/src/sys/unsupported/mod.rs
index 7bf6d40b76d..15b22c620d5 100644
--- a/library/std/src/sys/unsupported/mod.rs
+++ b/library/std/src/sys/unsupported/mod.rs
@@ -9,6 +9,7 @@ pub mod fs;
 pub mod io;
 pub mod locks;
 pub mod net;
+pub mod once;
 pub mod os;
 #[path = "../unix/os_str.rs"]
 pub mod os_str;
diff --git a/library/std/src/sys/unsupported/once.rs b/library/std/src/sys/unsupported/once.rs
new file mode 100644
index 00000000000..b4bb4975f41
--- /dev/null
+++ b/library/std/src/sys/unsupported/once.rs
@@ -0,0 +1,89 @@
+use crate::cell::Cell;
+use crate::sync as public;
+
+pub struct Once {
+    state: Cell<State>,
+}
+
+pub struct OnceState {
+    poisoned: bool,
+    set_state_to: Cell<State>,
+}
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum State {
+    Incomplete,
+    Poisoned,
+    Running,
+    Complete,
+}
+
+struct CompletionGuard<'a> {
+    state: &'a Cell<State>,
+    set_state_on_drop_to: State,
+}
+
+impl<'a> Drop for CompletionGuard<'a> {
+    fn drop(&mut self) {
+        self.state.set(self.set_state_on_drop_to);
+    }
+}
+
+// Safety: threads are not supported on this platform.
+unsafe impl Sync for Once {}
+
+impl Once {
+    #[inline]
+    #[rustc_const_stable(feature = "const_once_new", since = "1.32.0")]
+    pub const fn new() -> Once {
+        Once { state: Cell::new(State::Incomplete) }
+    }
+
+    #[inline]
+    pub fn is_completed(&self) -> bool {
+        self.state.get() == State::Complete
+    }
+
+    #[cold]
+    #[track_caller]
+    pub fn call(&self, ignore_poisoning: bool, f: &mut impl FnMut(&public::OnceState)) {
+        let state = self.state.get();
+        match state {
+            State::Poisoned if !ignore_poisoning => {
+                // Panic to propagate the poison.
+                panic!("Once instance has previously been poisoned");
+            }
+            State::Incomplete | State::Poisoned => {
+                self.state.set(State::Running);
+                // `guard` will set the new state on drop.
+                let mut guard =
+                    CompletionGuard { state: &self.state, set_state_on_drop_to: State::Poisoned };
+                // Run the function, letting it know if we're poisoned or not.
+                let f_state = public::OnceState {
+                    inner: OnceState {
+                        poisoned: state == State::Poisoned,
+                        set_state_to: Cell::new(State::Complete),
+                    },
+                };
+                f(&f_state);
+                guard.set_state_on_drop_to = f_state.inner.set_state_to.get();
+            }
+            State::Running => {
+                panic!("one-time initialization may not be performed recursively");
+            }
+            State::Complete => {}
+        }
+    }
+}
+
+impl OnceState {
+    #[inline]
+    pub fn is_poisoned(&self) -> bool {
+        self.poisoned
+    }
+
+    #[inline]
+    pub fn poison(&self) {
+        self.set_state_to.set(State::Poisoned)
+    }
+}
diff --git a/library/std/src/sys/wasi/mod.rs b/library/std/src/sys/wasi/mod.rs
index c8c47763a34..1dc3f2b2026 100644
--- a/library/std/src/sys/wasi/mod.rs
+++ b/library/std/src/sys/wasi/mod.rs
@@ -32,6 +32,8 @@ pub mod io;
 #[path = "../unsupported/locks/mod.rs"]
 pub mod locks;
 pub mod net;
+#[path = "../unsupported/once.rs"]
+pub mod once;
 pub mod os;
 #[path = "../unix/os_str.rs"]
 pub mod os_str;
diff --git a/library/std/src/sys/wasm/mod.rs b/library/std/src/sys/wasm/mod.rs
index d68c3e5f1df..77ebe3c4ac6 100644
--- a/library/std/src/sys/wasm/mod.rs
+++ b/library/std/src/sys/wasm/mod.rs
@@ -66,6 +66,8 @@ cfg_if::cfg_if! {
     } else {
         #[path = "../unsupported/locks/mod.rs"]
         pub mod locks;
+        #[path = "../unsupported/once.rs"]
+        pub mod once;
         #[path = "../unsupported/thread.rs"]
         pub mod thread;
     }