[PATCH v2 4/5] rust: introduce `InPlaceModule`

[Wedson Almeida Filho]

From: Wedson Almeida Filho <walmeida@xxxxxxxxxxxxx>

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.

Drivers that don't need this may continue to implement `Module` without
any changes.

Reviewed-by: Benno Lossin <benno.lossin@xxxxxxxxx>
Tested-by: Valentin Obst <kernel@xxxxxxxxxxxxxxx>
Signed-off-by: Wedson Almeida Filho <walmeida@xxxxxxxxxxxxx>
---
 rust/kernel/lib.rs    | 23 +++++++++++++++++++++++
 rust/macros/module.rs | 18 ++++++------------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 9141a95efb25..64aee4fbc53b 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -72,6 +72,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/export.h`](srctree/include/linux/export.h)
diff --git a/rust/macros/module.rs b/rust/macros/module.rs
index 6da1246742a5..4e5b5a68c3af 100644
--- a/rust/macros/module.rs
+++ b/rust/macros/module.rs
@@ -208,7 +208,7 @@ 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();
 
             // SAFETY: `__this_module` is constructed by the kernel at load time and will not be
             // freed until the module is unloaded.
@@ -275,23 +275,17 @@ pub(crate) fn module(ts: TokenStream) -> TokenStream {
             }}
 
             fn __init() -> ::core::ffi::c_int {{
-                match <{type_} as ::kernel::Module>::init(&THIS_MODULE) {{
-                    Ok(m) => {{
-                        unsafe {{
-                            __MOD = Some(m);
-                        }}
-                        return 0;
-                    }}
-                    Err(e) => {{
-                        return e.to_errno();
-                    }}
+                let initer = <{type_} as ::kernel::InPlaceModule>::init(&THIS_MODULE);
+                match unsafe {{ ::kernel::init::PinInit::__pinned_init(initer, __MOD.as_mut_ptr()) }} {{
+                    Ok(m) => 0,
+                    Err(e) => e.to_errno(),
                 }}
             }}
 
             fn __exit() {{
                 unsafe {{
                     // Invokes `drop()` on `__MOD`, which should be used for cleanup.
-                    __MOD = None;
+                    __MOD.assume_init_drop();
                 }}
             }}
 
-- 
2.34.1