Accept 2 or less arguments

Author: mekosko

riscv-rt/macros/src/lib.rs

@@ -52,21 +52,16 @@ use proc_macro::TokenStream;
 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(),
-            "`#[entry]` function has too many arguments",
-        )
-        .to_compile_error()
-        .into();
-    }
+    let arguments_limit = 3;
     #[cfg(feature = "u-boot")]
-    if f.sig.inputs.len() != 2 {
+    let arguments_limit = 2;
+
+    // check the function arguments
+    if f.sig.inputs.len() > arguments_limit {
         return parse::Error::new(
             f.sig.inputs.last().unwrap().span(),
-            "`#[entry]` function must have exactly two arguments",
+            "`#[entry]` function has too many arguments",
         )
         .to_compile_error()
         .into();
@@ -90,10 +85,7 @@ 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]);
-
+    if let Some(a1) = f.sig.inputs.get(0) {
         match a1 {
             FnArg::Receiver(_) => {
                 return parse::Error::new(a1.span(), "invalid argument")
@@ -108,18 +100,20 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
                 }
             }
         }
-        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")
+        if let Some(a2) = f.sig.inputs.get(1) {
+            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();
+                    }
+                }
             }
         }
     }