riscv-macros: post_init macro

Author: romancardenas

riscv-macros/CHANGELOG.md

@@ -5,7 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
-## [Unreleased]
+## v0.4.1 - Unreleased
+
+### Added
+
+- `post_init` macro for the `riscv-rt` crate (migrated from `riscv-rt-macros`).
 
 ## v0.4.0 - 2025-12-19
 

riscv-macros/Cargo.toml

@@ -9,7 +9,7 @@ keywords = ["riscv", "register", "peripheral"]
 license = "MIT OR Apache-2.0"
 name = "riscv-macros"
 repository = "https://github.com/rust-embedded/riscv"
-version = "0.4.0"
+version = "0.4.1"
 edition = "2021"
 
 [lib]
@@ -18,6 +18,7 @@ proc-macro = true
 [features]
 rt = []
 rt-v-trap = ["rt"]
+riscv-rt = ["rt", "syn/extra-traits", "syn/full"]
 
 [dependencies]
 proc-macro2 = "1.0"

riscv-macros/src/lib.rs

@@ -4,6 +4,9 @@ use syn::{parse_macro_input, DeriveInput};
 
 mod riscv;
 
+#[cfg(feature = "riscv-rt")]
+mod riscv_rt;
+
 /// Attribute-like macro that implements the traits of the `riscv-types` crate for a given enum.
 ///
 /// As these traits are unsafe, the macro must be called with the `unsafe` keyword followed by the trait name.
@@ -59,3 +62,31 @@ pub fn pac_enum(attr: TokenStream, item: 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
+///
+/// ```
+/// #[riscv_macros::post_init]
+/// unsafe fn before_main(hart_id: usize) {
+///     // do something here
+/// }
+/// ```
+#[cfg(feature = "riscv-rt")]
+#[proc_macro_attribute]
+pub fn post_init(args: TokenStream, input: TokenStream) -> TokenStream {
+    riscv_rt::Fn::post_init(args, input)
+}

riscv-macros/src/riscv_rt.rs

@@ -0,0 +1,230 @@
+use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::quote;
+use syn::{
+    parse_macro_input, punctuated::Punctuated, spanned::Spanned, token::Comma, Error, FnArg, Ident,
+    ItemFn, Result, ReturnType, Type, Visibility,
+};
+
+/// Enum representing the supported runtime function attributes
+pub enum Fn {
+    PostInit,
+}
+
+impl Fn {
+    /// Convenience method to generate the token stream for the `post_init` attribute
+    pub fn post_init(args: TokenStream, input: TokenStream) -> TokenStream {
+        match Self::PostInit.check_args_empty(args) {
+            Ok(_) => Self::PostInit.quote_fn(input),
+            Err(e) => e.to_compile_error().into(),
+        }
+    }
+
+    /// Generate the token stream for the function with the given attribute
+    fn quote_fn(&self, item: TokenStream) -> TokenStream {
+        let mut func = parse_macro_input!(item as ItemFn);
+
+        if let Err(e) = self.check_fn(&func) {
+            return e.to_compile_error().into();
+        }
+
+        let export_name = self.export_name(&func);
+        let link_section = self.link_section(&func);
+
+        // Append to function name the prefix __riscv_rt_ (to prevent users from calling it directly)
+        // Note that we do not change the export name, only the internal function name in the Rust code.
+        func.sig.ident = Ident::new(
+            &format!("__riscv_rt_{}", func.sig.ident),
+            func.sig.ident.span(),
+        );
+
+        quote! {
+            #export_name
+            #link_section
+            #func
+        }
+        .into()
+    }
+
+    /// Check if the function signature is valid for the given attribute
+    fn check_fn(&self, f: &ItemFn) -> Result<()> {
+        // First, check that the function is private
+        if f.vis != Visibility::Inherited {
+            let attr = self.attr_name();
+            return Err(Error::new(
+                f.vis.span(),
+                format!("`#[{attr}]` function must be private"),
+            ));
+        }
+        let sig = &f.sig;
+
+        // Next, check common aspects of the signature (constness, asyncness, generics, etc.)
+        let valid_signature = sig.constness.is_none()
+            && sig.asyncness.is_none()
+            && sig.abi.is_none()
+            && sig.generics.params.is_empty()
+            && sig.generics.where_clause.is_none()
+            && sig.variadic.is_none();
+        if !valid_signature {
+            let attr = self.attr_name();
+            let expected = self.expected_signature();
+            return Err(Error::new(
+                sig.span(),
+                format!("`#[{attr}]` function signature must be `{expected}`"),
+            ));
+        }
+
+        // Finally, check that input arguments and output type are valid
+        self.check_inputs(&sig.inputs)?;
+        self.check_output(&sig.output)
+    }
+
+    /// Utility method for printing attribute name in error messages
+    const fn attr_name(&self) -> &'static str {
+        // Use this match to specify attribute names for different functions in the future
+        match self {
+            Self::PostInit => "post_init",
+        }
+    }
+
+    /// Utility method for printing expected function signature in error messages
+    const fn expected_signature(&self) -> &'static str {
+        // Use this match to specify expected signatures for different functions in the future
+        match self {
+            Self::PostInit => "[unsafe] fn([usize])",
+        }
+    }
+
+    /// Check if the function has valid input arguments for the given attribute
+    fn check_inputs(&self, inputs: &Punctuated<FnArg, Comma>) -> Result<()> {
+        // Use this match to specify expected input arguments for different functions in the future
+        match self {
+            Self::PostInit => self.check_fn_args(inputs, &["usize"]),
+        }
+    }
+
+    /// Check if the function has a valid output type for the given attribute
+    fn check_output(&self, output: &ReturnType) -> Result<()> {
+        // Use this match to specify expected output types for different functions in the future
+        match self {
+            Self::PostInit => match output {
+                // post_init return type is ()
+                ReturnType::Default => Ok(()),
+                ReturnType::Type(_, ty) => match **ty {
+                    Type::Tuple(ref tuple) => {
+                        if tuple.elems.is_empty() {
+                            Ok(())
+                        } else {
+                            Err(Error::new(tuple.span(), "return type must be ()"))
+                        }
+                    }
+                    _ => Err(Error::new(ty.span(), "return type must be ()")),
+                },
+            },
+        }
+    }
+
+    /// The export name for the given attribute
+    fn export_name(&self, _f: &ItemFn) -> Option<TokenStream2> {
+        // Use this match to specify export names for different functions in the future
+        let export_name = match self {
+            Self::PostInit => Some("__post_init".to_string()),
+        };
+
+        export_name.map(|name| quote! { #[export_name = #name] })
+    }
+
+    /// The link section attribute for the given attribute (if any)
+    fn link_section(&self, _f: &ItemFn) -> Option<TokenStream2> {
+        // Use this match to specify section names for different functions in the future
+        let section_name: Option<String> = match self {
+            Self::PostInit => None,
+        };
+
+        section_name.map(|section| quote! {
+            #[cfg_attr(any(target_arch = "riscv32", target_arch = "riscv64"), link_section = #section)]
+        })
+    }
+
+    /// Check that no arguments were provided to the macro attribute
+    fn check_args_empty(&self, args: TokenStream) -> Result<()> {
+        if args.is_empty() {
+            Ok(())
+        } else {
+            let args: TokenStream2 = args.into();
+            let attr = self.attr_name();
+            Err(Error::new(
+                args.span(),
+                format!("`#[{attr}]` function does not accept any arguments"),
+            ))
+        }
+    }
+
+    /// Iterates through the input arguments and checks that their types match the expected types
+    fn check_fn_args(
+        &self,
+        inputs: &Punctuated<FnArg, Comma>,
+        expected_types: &[&str],
+    ) -> Result<()> {
+        let mut expected_iter = expected_types.iter();
+        for arg in inputs.iter() {
+            match expected_iter.next() {
+                Some(expected) => check_arg_type(arg, expected)?,
+                None => {
+                    let attr = self.attr_name();
+                    return Err(Error::new(
+                        arg.span(),
+                        format!("`#[{attr}]` function has too many input arguments"),
+                    ));
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+/// Check if a function argument matches the expected type
+fn check_arg_type(arg: &FnArg, expected: &str) -> Result<()> {
+    match arg {
+        FnArg::Typed(argument) => {
+            if !is_correct_type(&argument.ty, expected) {
+                Err(Error::new(
+                    argument.ty.span(),
+                    format!("argument type must be {expected}"),
+                ))
+            } else {
+                Ok(())
+            }
+        }
+        FnArg::Receiver(_) => Err(Error::new(arg.span(), "invalid argument")),
+    }
+}
+
+/// Check if a type matches the expected type name
+fn is_correct_type(ty: &Type, expected: &str) -> bool {
+    let correct: Type = syn::parse_str(expected).unwrap();
+    if let Some(ty) = strip_type_path(ty) {
+        ty == correct
+    } else {
+        false
+    }
+}
+
+/// Strip the path of a type, returning only the last segment (e.g., `core::usize` -> `usize`)
+fn strip_type_path(ty: &Type) -> Option<Type> {
+    match ty {
+        Type::Ptr(ty) => {
+            let mut ty = ty.clone();
+            *ty.elem = strip_type_path(&ty.elem)?;
+            Some(Type::Ptr(ty))
+        }
+        Type::Path(ty) => {
+            let mut ty = ty.clone();
+            let last_segment = ty.path.segments.last().unwrap().clone();
+            ty.path.segments = Punctuated::new();
+            ty.path.segments.push_value(last_segment);
+            Some(Type::Path(ty))
+        }
+        _ => None,
+    }
+}

riscv-rt/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
 
+### Changed
+
+- Using `post_init` from `riscv-macros` instead of `riscv-rt-macros`.
+
 ## v0.17.0 - 2025-12-19
 
 ### Added

riscv-rt/Cargo.toml

@@ -25,6 +25,7 @@ riscv-target-parser = { path = "../riscv-target-parser", version = "0.1.2" }
 
 [dependencies]
 riscv = { path = "../riscv", version = "0.16.0", features = ["rt"] }
+riscv-macros = { path = "../riscv-macros", version = "0.4.1", features = ["riscv-rt"] }
 riscv-types = { path = "../riscv-types", version = "0.1.0" }
 riscv-rt-macros = { path = "macros", version = "0.6.1" }
 
@@ -40,7 +41,7 @@ pre-init = []
 post-init = []
 s-mode = ["riscv-rt-macros/s-mode"]
 single-hart = []
-v-trap = ["riscv-rt-macros/v-trap", "riscv/rt-v-trap"]
+v-trap = ["riscv-rt-macros/v-trap", "riscv/rt-v-trap", "riscv-macros/rt-v-trap"]
 u-boot = ["riscv-rt-macros/u-boot", "single-hart"]
 no-interrupts = []
 no-exceptions = []

riscv-rt/src/lib.rs

@@ -703,7 +703,7 @@ pub use riscv_rt_macros::{core_interrupt, entry, exception, external_interrupt};
 pub use riscv_types::*;
 
 #[cfg(feature = "post-init")]
-pub use riscv_rt_macros::post_init;
+pub use riscv_macros::post_init;
 #[cfg(feature = "pre-init")]
 pub use riscv_rt_macros::pre_init;
 

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

@@ -1,4 +1,4 @@
-error: `#[post_init]` function has too many arguments
+error: `#[post_init]` function has too many input 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_async.stderr

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