[PATCH 1/2] rust: introduce `InPlaceModule`

[Wedson Almeida Filho <wedsonaf@gmail.com>]

From: Wedson Almeida Filho <walmeida@microsoft.com>

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.

Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
---
 rust/kernel/lib.rs    | 25 ++++++++++++++++++++++++-
 rust/macros/module.rs | 18 ++++++------------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 5c641233e26d..64aee4fbc53b 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -62,7 +62,7 @@
 /// The top level entrypoint to implementing a kernel module.
 ///
 /// For any teardown or cleanup operations, your type may implement [`Drop`].
-pub trait Module: Sized + Sync {
+pub trait Module: Sized + Sync + Send {
     /// Called at module initialization time.
     ///
     /// Use this method to perform whatever setup or registration your module
@@ -72,6 +72,29 @@ pub trait Module: Sized + Sync {
     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 27979e582e4b..0b2bb4ec2fba 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 {{ initer.__pinned_init(__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