riscv-rt: Add __post_init Rust function

riscv-rt/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
 
+### Added
+
+- New `post-init` feature to run a Rust `__post_init` function before jumping to `main`.
+- New `#[riscv_rt::post_init]` attribute to aid in the definition of the `__post_init` function. 
+
 ### Changed
 
 - `main` function no longer needs to be close to `_start`. A linker script may copy

riscv-rt/Cargo.toml

@@ -14,7 +14,7 @@ links = "riscv-rt" # Prevent multiple versions of riscv-rt being linked
 
 [package.metadata.docs.rs]
 default-target = "riscv64imac-unknown-none-elf"
-features = ["pre-init"]
+features = ["pre-init", "post-init"]
 targets = [
     "riscv32i-unknown-none-elf", "riscv32imc-unknown-none-elf", "riscv32imac-unknown-none-elf",
     "riscv64imac-unknown-none-elf", "riscv64gc-unknown-none-elf",
@@ -33,6 +33,7 @@ panic-halt = "1.0.0"
 
 [features]
 pre-init = []
+post-init = []
 s-mode = ["riscv-rt-macros/s-mode"]
 single-hart = []
 v-trap = ["riscv-rt-macros/v-trap"]

riscv-rt/macros/src/lib.rs

@@ -63,26 +63,6 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
         .into();
     }
 
-    fn check_correct_type(argument: &PatType, ty: &str) -> Option<TokenStream> {
-        let inv_type_message = format!("argument type must be {ty}");
-
-        if !is_correct_type(&argument.ty, ty) {
-            let error = parse::Error::new(argument.ty.span(), inv_type_message);
-
-            Some(error.to_compile_error().into())
-        } else {
-            None
-        }
-    }
-    fn check_argument_type(argument: &FnArg, ty: &str) -> Option<TokenStream> {
-        let argument_error = parse::Error::new(argument.span(), "invalid argument");
-        let argument_error = argument_error.to_compile_error().into();
-
-        match argument {
-            FnArg::Typed(argument) => check_correct_type(argument, ty),
-            FnArg::Receiver(_) => Some(argument_error),
-        }
-    }
     #[cfg(not(feature = "u-boot"))]
     for argument in f.sig.inputs.iter() {
         if let Some(message) = check_argument_type(argument, "usize") {
@@ -181,6 +161,28 @@ fn is_correct_type(ty: &Type, name: &str) -> bool {
     }
 }
 
+fn check_correct_type(argument: &PatType, ty: &str) -> Option<TokenStream> {
+    let inv_type_message = format!("argument type must be {ty}");
+
+    if !is_correct_type(&argument.ty, ty) {
+        let error = parse::Error::new(argument.ty.span(), inv_type_message);
+
+        Some(error.to_compile_error().into())
+    } else {
+        None
+    }
+}
+
+fn check_argument_type(argument: &FnArg, ty: &str) -> Option<TokenStream> {
+    let argument_error = parse::Error::new(argument.span(), "invalid argument");
+    let argument_error = argument_error.to_compile_error().into();
+
+    match argument {
+        FnArg::Typed(argument) => check_correct_type(argument, ty),
+        FnArg::Receiver(_) => Some(argument_error),
+    }
+}
+
 /// Attribute to mark which function will be called at the beginning of the reset handler.
 /// You must enable the `pre_init` feature in the `riscv-rt` crate to use this macro.
 ///
@@ -263,6 +265,93 @@ pub fn pre_init(args: TokenStream, input: TokenStream) -> TokenStream {
     .into()
 }
 
+/// Attribute to mark which function will be called before jumping to the entry point.
+/// You must enable the `post-init` feature in the `riscv-rt` crate to use this macro.
+///
+/// In contrast with `__pre_init`, this function is called after the static variables
+/// are initialized, so it is safe to access them. It is also safe to run Rust code.
+///
+/// The function must have the signature of `[unsafe] fn([usize])`, where the argument
+/// corresponds to the hart ID of the current hart. This is useful for multi-hart systems
+/// to perform hart-specific initialization.
+///
+/// # IMPORTANT
+///
+/// This attribute can appear at most *once* in the dependency graph.
+///
+/// # Examples
+///
+/// ```
+/// use riscv_rt_macros::post_init;
+/// #[post_init]
+/// unsafe fn before_main(hart_id: usize) {
+///     // do something here
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn post_init(args: TokenStream, input: TokenStream) -> TokenStream {
+    let f = parse_macro_input!(input as ItemFn);
+
+    // check the function arguments
+    if f.sig.inputs.len() > 1 {
+        return parse::Error::new(
+            f.sig.inputs.last().unwrap().span(),
+            "`#[post_init]` function has too many arguments",
+        )
+        .to_compile_error()
+        .into();
+    }
+    for argument in f.sig.inputs.iter() {
+        if let Some(message) = check_argument_type(argument, "usize") {
+            return message;
+        };
+    }
+
+    // check the function signature
+    let valid_signature = f.sig.constness.is_none()
+        && f.sig.asyncness.is_none()
+        && f.vis == Visibility::Inherited
+        && f.sig.abi.is_none()
+        && f.sig.generics.params.is_empty()
+        && f.sig.generics.where_clause.is_none()
+        && f.sig.variadic.is_none()
+        && match f.sig.output {
+            ReturnType::Default => true,
+            ReturnType::Type(_, ref ty) => match **ty {
+                Type::Tuple(ref tuple) => tuple.elems.is_empty(),
+                _ => false,
+            },
+        };
+
+    if !valid_signature {
+        return parse::Error::new(
+            f.span(),
+            "`#[post_init]` function must have signature `[unsafe] fn([usize])`",
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    if !args.is_empty() {
+        return parse::Error::new(Span::call_site(), "This attribute accepts no arguments")
+            .to_compile_error()
+            .into();
+    }
+
+    // XXX should we blacklist other attributes?
+    let attrs = f.attrs;
+    let ident = f.sig.ident;
+    let args = f.sig.inputs;
+    let block = f.block;
+
+    quote!(
+        #[export_name = "__post_init"]
+        #(#attrs)*
+        unsafe fn #ident(#args) #block
+    )
+    .into()
+}
+
 struct AsmLoopArgs {
     asm_template: String,
     count_from: usize,

riscv-rt/src/lib.rs

@@ -553,6 +553,14 @@
 //! );
 //! ```
 //!
+//! ## `post-init`
+//!
+//! When enabled, the runtime will execute the `__post_init` function to be run before jumping to the main function.
+//! If the feature is enabled, the `__post_init` function must be defined in the user code (i.e., no default implementation
+//! is provided by this crate). If the feature is disabled, the `__post_init` function is not required.
+//!
+//! You can use the [`#[post_init]`][attr-post-init] attribute to define a post-init function with Rust.
+//!
 //! ## `single-hart`
 //!
 //! Saves a little code size if there is only one hart on the target.
@@ -595,6 +603,7 @@
 //! [attr-external-interrupt]: attr.external_interrupt.html
 //! [attr-core-interrupt]: attr.core_interrupt.html
 //! [attr-pre-init]: attr.pre_init.html
+//! [attr-post-init]: attr.post_init.html
 
 // NOTE: Adapted from cortex-m/src/lib.rs
 #![no_std]
@@ -624,6 +633,8 @@ use riscv::register::{
 pub use riscv_pac::*;
 pub use riscv_rt_macros::{core_interrupt, entry, exception, external_interrupt};
 
+#[cfg(feature = "post-init")]
+pub use riscv_rt_macros::post_init;
 #[cfg(feature = "pre-init")]
 pub use riscv_rt_macros::pre_init;
 
@@ -650,10 +661,14 @@ pub static __ONCE__: () = ();
 #[export_name = "_start_rust"]
 pub unsafe extern "C" fn start_rust(a0: usize, a1: usize, a2: usize) -> ! {
     extern "Rust" {
+        #[cfg(feature = "post-init")]
+        fn __post_init(a0: usize);
         fn _setup_interrupts();
         fn hal_main(a0: usize, a1: usize, a2: usize) -> !;
     }
 
+    #[cfg(feature = "post-init")]
+    __post_init(a0);
     _setup_interrupts();
     hal_main(a0, a1, a2);
 }

tests-trybuild/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2021"
 
 [dependencies]
 riscv = { path = "../riscv" }
-riscv-rt = { path = "../riscv-rt", features = ["no-exceptions", "no-interrupts"] }
+riscv-rt = { path = "../riscv-rt", features = ["no-exceptions", "no-interrupts", "post-init"] }
 trybuild = "1.0"
 
 [features]

tests-trybuild/tests/riscv-rt/post_init/fail_arg_count.rs

@@ -0,0 +1,4 @@
+#[riscv_rt::post_init]
+fn before_main(_hart_id: usize, _dtb: usize) {}
+
+fn main() {}

tests-trybuild/tests/riscv-rt/post_init/fail_arg_count.stderr

@@ -0,0 +1,5 @@
+error: `#[post_init]` function has too many arguments
+ --> tests/riscv-rt/post_init/fail_arg_count.rs:2:33
+  |
+2 | fn before_main(_hart_id: usize, _dtb: usize) {}
+  |                                 ^^^^

tests-trybuild/tests/riscv-rt/post_init/fail_arg_type.rs

@@ -0,0 +1,4 @@
+#[riscv_rt::post_init]
+fn before_main(_hart_id: String) {}
+
+fn main() {}

tests-trybuild/tests/riscv-rt/post_init/fail_arg_type.stderr

@@ -0,0 +1,5 @@
+error: argument type must be usize
+ --> tests/riscv-rt/post_init/fail_arg_type.rs:2:26
+  |
+2 | fn before_main(_hart_id: String) {}
+  |                          ^^^^^^

tests-trybuild/tests/riscv-rt/post_init/fail_async.rs

@@ -0,0 +1,4 @@
+#[riscv_rt::post_init]
+async fn before_main() {}
+
+fn main() {}

tests-trybuild/tests/riscv-rt/post_init/fail_async.stderr

@@ -0,0 +1,5 @@
+error: `#[post_init]` function must have signature `[unsafe] fn([usize])`
+ --> tests/riscv-rt/post_init/fail_async.rs:2:1
+  |
+2 | async fn before_main() {}
+  | ^^^^^

tests-trybuild/tests/riscv-rt/post_init/pass_empty.rs

@@ -0,0 +1,4 @@
+#[riscv_rt::post_init]
+fn before_main() {}
+
+fn main() {}

tests-trybuild/tests/riscv-rt/post_init/pass_safe.rs

@@ -0,0 +1,4 @@
+#[riscv_rt::post_init]
+fn before_main(_hart_id: usize) {}
+
+fn main() {}

tests-trybuild/tests/riscv-rt/post_init/pass_unsafe.rs

@@ -0,0 +1,4 @@
+#[riscv_rt::post_init]
+unsafe fn before_main(_hart_id: usize) {}
+
+fn main() {}