[PATCH v2 1/2] rust: Add support for calling a function exactly once

[FUJITA Tomonori <fujita.tomonori@gmail.com>]

Add the Rust equivalent of the kernel's `DO_ONCE_LITE` macro. While it
would be possible to implement the feature entirely as a Rust macro,
the functionality that can be implemented as regular functions has
been extracted and implemented as the `OnceLite` struct for better
code maintainability.

Signed-off-by: FUJITA Tomonori <fujita.tomonori@gmail.com>
---
 rust/kernel/sync.rs           |  1 +
 rust/kernel/sync/once_lite.rs | 83 +++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 rust/kernel/sync/once_lite.rs

diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
index cf5b638a097d..03789283e2ec 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -17,6 +17,7 @@
 mod condvar;
 pub mod lock;
 mod locked_by;
+pub mod once_lite;
 pub mod poll;
 pub mod rcu;
 mod refcount;
diff --git a/rust/kernel/sync/once_lite.rs b/rust/kernel/sync/once_lite.rs
new file mode 100644
index 000000000000..04c48d18ca39
--- /dev/null
+++ b/rust/kernel/sync/once_lite.rs
@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Support for calling a function exactly once.
+//!
+//! C header: [`include/linux/once_lite.h`](srctree/include/linux/once_lite.h)
+
+use super::atomic::{
+    Atomic,
+    AtomicType,
+    Relaxed, //
+};
+
+/// A lightweight `call_once` primitive.
+///
+/// This structure provides the Rust equivalent of the kernel's `DO_ONCE_LITE` macro.
+/// While it would be possible to implement the feature entirely as a Rust macro,
+/// the functionality that can be implemented as regular functions has been
+/// extracted and implemented as the `OnceLite` struct for better code maintainability.
+pub struct OnceLite(Atomic<State>);
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+#[repr(i32)]
+enum State {
+    Incomplete = 0,
+    Complete = 1,
+}
+
+// SAFETY: `State` and `i32` has the same size and alignment, and it's round-trip
+// transmutable to `i32`.
+unsafe impl AtomicType for State {
+    type Repr = i32;
+}
+
+impl OnceLite {
+    /// Creates a new [`OnceLite`] in the incomplete state.
+    #[inline(always)]
+    #[allow(clippy::new_without_default)]
+    pub const fn new() -> Self {
+        OnceLite(Atomic::new(State::Incomplete))
+    }
+
+    /// Calls the provided function exactly once.
+    ///
+    /// There is no other synchronization between two `call_once()`s
+    /// except that only one will execute `f`, in other words, callers
+    /// should not use a failed `call_once()` as a proof that another
+    /// `call_once()` has already finished and the effect is observable
+    /// to this thread.
+    pub fn call_once<F>(&self, f: F) -> bool
+    where
+        F: FnOnce(),
+    {
+        // ORDERING: `Relaxed` is used here since no synchronization is required
+        // for `call_once()`.
+        let old = self.0.xchg(State::Complete, Relaxed);
+        if old == State::Complete {
+            return false;
+        }
+
+        f();
+        true
+    }
+}
+
+/// Run the given function exactly once.
+///
+/// This is equivalent to the kernel's `DO_ONCE_LITE` macro.
+///
+/// # Examples
+///
+/// ```
+/// kernel::do_once_lite! {
+///     kernel::pr_info!("This will be printed only once\n");
+/// };
+/// ```
+#[macro_export]
+macro_rules! do_once_lite {
+    { $($e:tt)* } => {{
+        #[link_section = ".data..once"]
+        static ONCE: $crate::sync::once_lite::OnceLite = $crate::sync::once_lite::OnceLite::new();
+        ONCE.call_once(|| { $($e)* });
+    }};
+}
-- 
2.43.0