lib: rust_support: Add init hook support

Bug: 304851081
Change-Id: Iebc3de4c684264e2f13d8f05ff9b1aea1327d7c3

4 files changed, 95 insertions, 0 deletions

diff --git a/lib/rust_support/bindings.h b/lib/rust_support/bindings.h
index e628e4ea..dc7c72a0 100644
--- a/lib/rust_support/bindings.h
+++ b/lib/rust_support/bindings.h
@@ -2,4 +2,5 @@
 #include <err.h>
 #include <kernel/mutex.h>
 #include <kernel/vm.h>
+#include <lk/init.h>
 #include <panic.h>
diff --git a/lib/rust_support/init.rs b/lib/rust_support/init.rs
new file mode 100644
index 00000000..6bcae416
--- /dev/null
+++ b/lib/rust_support/init.rs
@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2024 Google Inc. All rights reserved
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files
+ * (the "Software"), to deal in the Software without restriction,
+ * including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+use crate::sys::uint;
+use core::ffi::c_char;
+
+pub use crate::sys::lk_init_flags;
+pub use crate::sys::lk_init_level;
+pub use crate::sys::lk_init_struct;
+
+// SAFETY: lk_init_struct does not use interior mutability.
+unsafe impl Sync for lk_init_struct {}
+
+// lk_init_level constants have large gaps between them and some modules
+// add or subtract from these constants to indicate that it wants to run
+// right before or after other init hooks at a given level. Add add and sub
+// functions to lk_init_level to allow this for rust init hooks as well.
+impl lk_init_level {
+    pub const fn add(mut self, rhs: uint) -> Self {
+        self.0 += rhs;
+        self
+    }
+    pub const fn sub(mut self, rhs: uint) -> Self {
+        self.0 -= rhs;
+        self
+    }
+}
+
+impl lk_init_struct {
+    pub const fn new(
+        level: lk_init_level,
+        flags: lk_init_flags,
+        hook: extern "C" fn(uint),
+        name: *const c_char,
+    ) -> Self {
+        lk_init_struct { level: level.0, flags: flags.0, hook: Option::Some(hook), name: name }
+    }
+}
+
+#[macro_export]
+macro_rules! LK_INIT_HOOK_FLAGS {
+    ($name:ident, $hook:expr, $level:expr, $flags:expr) => {
+        #[link_section = ".lk_init"]
+        #[no_mangle]
+        #[used]
+        static $name: $crate::init::lk_init_struct = $crate::init::lk_init_struct::new(
+            $level,
+            $flags,
+            $hook,
+            (concat!(stringify!($name), "\0").as_bytes()).as_ptr().cast(),
+        );
+    };
+}
+
+#[macro_export]
+macro_rules! LK_INIT_HOOK {
+    ($name:ident, $hook:expr, $level:expr) => {
+        $crate::LK_INIT_HOOK_FLAGS!(
+            $name,
+            $hook,
+            $level,
+            $crate::init::lk_init_flags::LK_INIT_FLAG_PRIMARY_CPU
+        );
+    };
+}
diff --git a/lib/rust_support/lib.rs b/lib/rust_support/lib.rs
index 3e9a8793..0143dba6 100644
--- a/lib/rust_support/lib.rs
+++ b/lib/rust_support/lib.rs
@@ -39,6 +39,7 @@ mod sys {
 }
 
 pub mod err;
+pub mod init;
 pub mod mmu;
 pub mod sync;
 pub mod vmm;
diff --git a/lib/rust_support/rules.mk b/lib/rust_support/rules.mk
index acadcd24..3840b782 100644
--- a/lib/rust_support/rules.mk
+++ b/lib/rust_support/rules.mk
@@ -46,6 +46,9 @@ MODULE_BINDGEN_ALLOW_FUNCTIONS := \
 	vmm_alloc_contiguous \
 	vmm_free_region \
 
+MODULE_BINDGEN_ALLOW_TYPES := \
+	lk_init_.* \
+
 MODULE_BINDGEN_ALLOW_VARS := \
 	_kernel_aspace \
 	ARCH_MMU_FLAG_.* \
@@ -53,6 +56,11 @@ MODULE_BINDGEN_ALLOW_VARS := \
 	PAGE_SIZE \
 	PAGE_SIZE_SHIFT \
 
+MODULE_BINDGEN_FLAGS := \
+	--newtype-enum lk_init_level \
+	--bitfield-enum lk_init_flags \
+	--no-prepend-enum-name \
+
 MODULE_BINDGEN_SRC_HEADER := $(LOCAL_DIR)/bindings.h
 
 include make/module.mk