rust: print: Add support for calling a function exactly once

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>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Gary Guo <gary@garyguo.net>
Link: https://patch.msgid.link/20251117002452.4068692-2-fujita.tomonori@gmail.com
[ Added prefix to title. - Miguel ]
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

1 files changed, 83 insertions, 0 deletions

diff --git a/rust/kernel/print.rs b/rust/kernel/print.rs
index 2d743d78d220..af32054be5a2 100644
--- a/rust/kernel/print.rs
+++ b/rust/kernel/print.rs
@@ -11,6 +11,11 @@ use crate::{
     fmt,
     prelude::*,
     str::RawFormatter,
+    sync::atomic::{
+        Atomic,
+        AtomicType,
+        Relaxed, //
+    },
 };
 
 // Called from `vsprintf` with format specifier `%pA`.
@@ -423,3 +428,81 @@ macro_rules! pr_cont (
         $crate::print_macro!($crate::print::format_strings::CONT, true, $($arg)*)
     )
 );
+
+/// 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(),
+    {
+        // Avoid expensive cmpxchg if already completed.
+        // ORDERING: `Relaxed` is used here since no synchronization is required.
+        let old = self.0.load(Relaxed);
+        if old == State::Complete {
+            return false;
+        }
+
+        // ORDERING: `Relaxed` is used here since no synchronization is required.
+        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::print::OnceLite = $crate::print::OnceLite::new();
+        ONCE.call_once(|| { $($e)* });
+    }};
+}