From 6cfc6c7f5c5e5d2a9b27574209b0204985cf7b7b Mon Sep 17 00:00:00 2001
From: mekosko <mekosko@morojenoye.com>
Date: Fri, 23 Aug 2024 21:51:45 +0500
Subject: [PATCH] Make support for u-boot

---
 riscv-rt/Cargo.toml        |  1 +
 riscv-rt/macros/Cargo.toml |  1 +
 riscv-rt/macros/src/lib.rs | 46 ++++++++++++++++++++++++++++++++++++++
 riscv-rt/src/asm.rs        |  5 +++--
 4 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/riscv-rt/Cargo.toml b/riscv-rt/Cargo.toml
index 2910e355..440a41bb 100644
--- a/riscv-rt/Cargo.toml
+++ b/riscv-rt/Cargo.toml
@@ -23,3 +23,4 @@ panic-halt = "0.2.0"
 s-mode = ["riscv-rt-macros/s-mode"]
 single-hart = []
 v-trap = ["riscv-rt-macros/v-trap"]
+u-boot = ["riscv-rt-macros/u-boot"]
diff --git a/riscv-rt/macros/Cargo.toml b/riscv-rt/macros/Cargo.toml
index 19f6690d..b7dba6fb 100644
--- a/riscv-rt/macros/Cargo.toml
+++ b/riscv-rt/macros/Cargo.toml
@@ -23,3 +23,4 @@ syn = { version = "2.0", features = ["extra-traits", "full"] }
 [features]
 s-mode = []
 v-trap = []
+u-boot = []
diff --git a/riscv-rt/macros/src/lib.rs b/riscv-rt/macros/src/lib.rs
index c34e3a25..25d571f4 100644
--- a/riscv-rt/macros/src/lib.rs
+++ b/riscv-rt/macros/src/lib.rs
@@ -53,6 +53,7 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
     let f = parse_macro_input!(input as ItemFn);
 
     // check the function arguments
+    #[cfg(not(feature = "u-boot"))]
     if f.sig.inputs.len() > 3 {
         return parse::Error::new(
             f.sig.inputs.last().unwrap().span(),
@@ -61,6 +62,17 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
         .to_compile_error()
         .into();
     }
+    #[cfg(feature = "u-boot")]
+    if f.sig.inputs.len() != 2 {
+        return parse::Error::new(
+            f.sig.inputs.last().unwrap().span(),
+            "`#[entry]` function must have exactly two arguments",
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    #[cfg(not(feature = "u-boot"))]
     for arg in &f.sig.inputs {
         match arg {
             FnArg::Receiver(_) => {
@@ -77,6 +89,40 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
             }
         }
     }
+    #[cfg(feature = "u-boot")]
+    {
+        // We have checked that there are two arguments above.
+        let (a1, a2) = (&f.sig.inputs[0], &f.sig.inputs[1]);
+
+        match a1 {
+            FnArg::Receiver(_) => {
+                return parse::Error::new(a1.span(), "invalid argument")
+                    .to_compile_error()
+                    .into();
+            }
+            FnArg::Typed(t) => {
+                if !is_simple_type(&t.ty, "c_int") {
+                    return parse::Error::new(t.ty.span(), "argument type must be c_int")
+                        .to_compile_error()
+                        .into();
+                }
+            }
+        }
+        match a2 {
+            FnArg::Receiver(_) => {
+                return parse::Error::new(a2.span(), "invalid argument")
+                    .to_compile_error()
+                    .into();
+            }
+            FnArg::Typed(t) => {
+                if !is_simple_type(&t.ty, "usize") {
+                    return parse::Error::new(t.ty.span(), "argument type must be usize")
+                        .to_compile_error()
+                        .into();
+                }
+            }
+        }
+    }
 
     // check the function signature
     let valid_signature = f.sig.constness.is_none()
diff --git a/riscv-rt/src/asm.rs b/riscv-rt/src/asm.rs
index 7883c446..46fb112a 100644
--- a/riscv-rt/src/asm.rs
+++ b/riscv-rt/src/asm.rs
@@ -125,6 +125,7 @@ cfg_global_asm!(
 );
 
 // STORE A0..A2 IN THE STACK, AS THEY WILL BE NEEDED LATER BY main
+#[cfg(not(feature = "u-boot"))]
 cfg_global_asm!(
     #[cfg(riscv32)]
     "addi sp, sp, -4 * 3
@@ -213,12 +214,12 @@ riscv_rt_macros::loop_global_asm!("    fmv.w.x f{}, x0", 32);
 // SET UP INTERRUPTS, RESTORE a0..a2, AND JUMP TO MAIN RUST FUNCTION
 cfg_global_asm!(
     "call _setup_interrupts",
-    #[cfg(riscv32)]
+    #[cfg(all(riscv32, not(feature = "u-boot")))]
     "lw a0, 4 * 0(sp)
     lw a1, 4 * 1(sp)
     lw a2, 4 * 2(sp)
     addi sp, sp, 4 * 3",
-    #[cfg(riscv64)]
+    #[cfg(all(riscv64, not(feature = "u-boot")))]
     "ld a0, 8 * 0(sp)
     ld a1, 8 * 1(sp)
     ld a2, 8 * 2(sp)