rust: introduce `InPlaceModule`

This allows modules to be initialised in-place in pinned memory, which
enables the usage of pinned types (e.g., mutexes, spinlocks, driver
registrations, etc.) in modules without any extra allocations.

Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
Signed-off-by: Danilo Krummrich <dakr@kernel.org>

2 files changed, 35 insertions, 16 deletions

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index b62451f64f6e..4fdb0d91f2ad 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -83,6 +83,29 @@ pub trait Module: Sized + Sync + Send {
     fn init(module: &'static ThisModule) -> error::Result<Self>;
 }
 
+/// A module that is pinned and initialised in-place.
+pub trait InPlaceModule: Sync + Send {
+    /// Creates an initialiser for the module.
+    ///
+    /// It is called when the module is loaded.
+    fn init(module: &'static ThisModule) -> impl init::PinInit<Self, error::Error>;
+}
+
+impl<T: Module> InPlaceModule for T {
+    fn init(module: &'static ThisModule) -> impl init::PinInit<Self, error::Error> {
+        let initer = move |slot: *mut Self| {
+            let m = <Self as Module>::init(module)?;
+
+            // SAFETY: `slot` is valid for write per the contract with `pin_init_from_closure`.
+            unsafe { slot.write(m) };
+            Ok(())
+        };
+
+        // SAFETY: On success, `initer` always fully initialises an instance of `Self`.
+        unsafe { init::pin_init_from_closure(initer) }
+    }
+}
+
 /// Equivalent to `THIS_MODULE` in the C API.
 ///
 /// C header: [`include/linux/init.h`](srctree/include/linux/init.h)
diff --git a/rust/macros/module.rs b/rust/macros/module.rs
index aef3b132f32b..a03266a78cfb 100644
--- a/rust/macros/module.rs
+++ b/rust/macros/module.rs
@@ -232,6 +232,7 @@ pub(crate) fn module(ts: TokenStream) -> TokenStream {
             mod __module_init {{
                 mod __module_init {{
                     use super::super::{type_};
+                    use kernel::init::PinInit;
 
                     /// The \"Rust loadable module\" mark.
                     //
@@ -242,7 +243,8 @@ pub(crate) fn module(ts: TokenStream) -> TokenStream {
                     #[used]
                     static __IS_RUST_MODULE: () = ();
 
-                    static mut __MOD: Option<{type_}> = None;
+                    static mut __MOD: core::mem::MaybeUninit<{type_}> =
+                        core::mem::MaybeUninit::uninit();
 
                     // Loadable modules need to export the `{{init,cleanup}}_module` identifiers.
                     /// # Safety
@@ -331,20 +333,14 @@ pub(crate) fn module(ts: TokenStream) -> TokenStream {
                     ///
                     /// This function must only be called once.
                     unsafe fn __init() -> core::ffi::c_int {{
-                        match <{type_} as kernel::Module>::init(&super::super::THIS_MODULE) {{
-                            Ok(m) => {{
-                                // SAFETY: No data race, since `__MOD` can only be accessed by this
-                                // module and there only `__init` and `__exit` access it. These
-                                // functions are only called once and `__exit` cannot be called
-                                // before or during `__init`.
-                                unsafe {{
-                                    __MOD = Some(m);
-                                }}
-                                return 0;
-                            }}
-                            Err(e) => {{
-                                return e.to_errno();
-                            }}
+                        let initer =
+                            <{type_} as kernel::InPlaceModule>::init(&super::super::THIS_MODULE);
+                        // SAFETY: No data race, since `__MOD` can only be accessed by this module
+                        // and there only `__init` and `__exit` access it. These functions are only
+                        // called once and `__exit` cannot be called before or during `__init`.
+                        match unsafe {{ initer.__pinned_init(__MOD.as_mut_ptr()) }} {{
+                            Ok(m) => 0,
+                            Err(e) => e.to_errno(),
                         }}
                     }}
 
@@ -359,7 +355,7 @@ pub(crate) fn module(ts: TokenStream) -> TokenStream {
                         // called once and `__init` was already called.
                         unsafe {{
                             // Invokes `drop()` on `__MOD`, which should be used for cleanup.
-                            __MOD = None;
+                            __MOD.assume_init_drop();
                         }}
                     }}