Use callee-saved registers for preserving arguments

Author: romancardenas

riscv-rt/CHANGELOG.md

@@ -31,6 +31,15 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Now, `_start_rust` jumps to `hal_main` instead of `main` directly. At linker level,
   `hal_main` maps to `main` if not defined. However, we now allow HALs to inject
   additional configuration code before jumping to the final user's `main` function.
+- Now, `a0` is preserved in `s1` during the startup process.
+
+### Removed
+
+- Removed usage of the stack before `_start_rust`. This was unsound, as in the `.init`
+  section RAM is still uninitialized.
+- Now, `a1` and `a2` are **not** preserved during the startup process. New documentation
+  of startup functions (`_mp_hook` and `__pre_init`) now provide additional implementation
+  guidelines to ensure a correct behavior of the runtime.
 
 ### Fixed
 

riscv-rt/src/asm.rs

@@ -85,7 +85,7 @@ _abs_start:
     bgeu t0, a0, 1f
     la t0, abort // If hart_id > _max_hart_id, jump to abort
     jr t0
-1:", // only valid harts reach this point
+1:  mv s1, a0", // If hart ID is valid, preserve it in s1
 
 // INITIALIZE GLOBAL POINTER, STACK POINTER, AND FRAME POINTER
     ".option push
@@ -113,18 +113,6 @@ _abs_start:
     "sub t1, t1, t0",
     "andi sp, t1, -16 // align stack to 16-bytes
     add s0, sp, zero",
-// STORE A0..A2 IN THE STACK, AS THEY WILL BE NEEDED LATER BY _start_rust
-    #[cfg(target_arch = "riscv32")]
-    "addi sp, sp, -4 * 4 // we must keep stack aligned to 16-bytes
-    sw a0, 4 * 0(sp)
-    sw a1, 4 * 1(sp)
-    sw a2, 4 * 2(sp)",
-    #[cfg(target_arch = "riscv64")]
-    "addi sp, sp, -8 * 4 // we must keep stack aligned to 16-bytes
-    sd a0, 8 * 0(sp)
-    sd a1, 8 * 1(sp)
-    sd a2, 8 * 2(sp)",
-
 // CALL __pre_init (IF ENABLED) AND INITIALIZE RAM
     #[cfg(not(feature = "single-hart"))]
     // Skip RAM initialization if current hart is not the boot hart
@@ -183,18 +171,9 @@ _abs_start:
     "fscsr x0",
 }
 
-// RESTORE a0..a2, AND JUMP TO _start_rust FUNCTION
-    #[cfg(target_arch = "riscv32")]
-    "lw a0, 4 * 0(sp)
-    lw a1, 4 * 1(sp)
-    lw a2, 4 * 2(sp)
-    addi sp, sp, 4 * 4",
-    #[cfg(target_arch = "riscv64")]
-    "ld a0, 8 * 0(sp)
-    ld a1, 8 * 1(sp)
-    ld a2, 8 * 2(sp)
-    addi sp, sp, 8 * 4",
-    "la t0, _start_rust
+// RESTORE a0 AND JUMP TO _start_rust FUNCTION
+    "mv a0, s1
+    la t0, _start_rust
     jr t0
     .cfi_endproc",
 

riscv-rt/src/lib.rs

@@ -5,7 +5,7 @@
 //!
 //! # Features
 //!
-//! This crates takes care of:
+//! This crate takes care of:
 //!
 //! - The memory layout of the program.
 //!
@@ -327,21 +327,11 @@
 //! Furthermore, as this function is expected to behave like a trap handler, it is
 //! necessary to make it be 4-byte aligned.
 //!
-//! ## `_mp_hook`
+//! ## `_mp_hook` (for multi-core targets only)
 //!
 //! This function is called from all the harts and must return true only for one hart,
 //! which will perform memory initialization. For other harts it must return false
 //! and implement wake-up in platform-dependent way (e.g., after waiting for a user interrupt).
-//! The parameter `hartid` specifies the hartid of the caller.
-//!
-//! This function can be redefined in the following way:
-//!
-//! ``` no_run
-//! #[export_name = "_mp_hook"]
-//! pub extern "Rust" fn mp_hook(hartid: usize) -> bool {
-//!    // ...
-//! }
-//! ```
 //!
 //! Default implementation of this function wakes hart 0 and busy-loops all the other harts.
 //!
@@ -350,6 +340,36 @@
 //! `_mp_hook` is only necessary in multi-core targets. If the `single-hart` feature is enabled,
 //! `_mp_hook` is not included in the binary.
 //!
+//! ### Important implementation guidelines
+//!
+//! This function is called during the early boot process. Thus, when implementing it, you **MUST** follow these guidelines:
+//!
+//! - Implement it in assembly (no Rust code is allowed at this point).
+//! - Allocate this function within the `.init` section.
+//! - You can get the hart id from the `a0` register.
+//! - You must set the return value in the `a0` register.
+//! - You must **NOT** use the `a1-a2` registers, as they are required later by the runtime.
+//!
+//! **Violating these constraints will result in incorrect arguments being passed to `main()`.**
+//!
+//! ### Implementation example
+//!
+//! The following example shows how to implement the `_mp_hook` function in assembly.
+//!
+//! ``` no_run
+//! core::arch::global_asm!(
+//!     r#".section .init.mp_hook, "ax"
+//!     .global _mp_hook
+//! _mp_hook:
+//!     beqz a0, 2f // check if hartid is 0
+//! 1:  wfi         // If not, wait for interrupt in a loop
+//!     j 1b
+//! 2:  li a0, 1    // Otherwise, return true
+//!     ret
+//!     "#
+//! );
+//! ```
+//!
 //! ## `_setup_interrupts`
 //!
 //! This function is called right before the main function and is responsible for setting up
@@ -506,12 +526,30 @@
 //! If the feature is enabled, the `__pre_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 `__pre_init` function is not required.
 //!
-//! As `__pre_init` runs before RAM is initialised, it is not sound to use a Rust function for `__pre_init`, and
-//! instead it should typically be written in assembly using `global_asm` or an external assembly file.
+//! ### Important implementation guidelines
+//!
+//! This function is called during the early boot process. Thus, when implementing it, you **MUST** follow these guidelines:
+//!
+//! - Implement it in assembly (no Rust code is allowed at this point).
+//! - Allocate this function within the `.init` section.
+//! - You must **NOT** modify the `a1-a2` registers, as they are required later by the runtime.
 //!
-//! Alternatively, you can use the [`#[pre_init]`][attr-pre-init] attribute to define a pre-init function with Rust.
-//! Note that using this macro is discouraged, as it may lead to undefined behavior.
-//! We left this option for backwards compatibility, but it is subject to removal in the future.
+//! **Violating these constraints will result in incorrect arguments being passed to `main()`.**
+//!
+//! ### Implementation example
+//!
+//! The following example shows how to implement the `__pre_init` function in assembly.
+//!
+//! ``` no_run
+//! core::arch::global_asm!(
+//!     r#".section .init.__pre_init, "ax"
+//!     .global __pre_init
+//! __pre_init:
+//!     // Do some pre-initialization work here and return
+//!     ret
+//!     "#
+//! );
+//! ```
 //!
 //! ## `single-hart`
 //!