From 6b149bb4c795b396b1066173110b81a8333dafee Mon Sep 17 00:00:00 2001
From: Nicholas Bishop <nbishop@nbishop.net>
Date: Sun, 7 Jul 2024 13:43:01 -0400
Subject: [PATCH] uefi-macros: Allow zero-param function in the entry macro

---
 uefi-macros/CHANGELOG.md                   |  5 +++
 uefi-macros/src/lib.rs                     | 44 +++++++++++++++++-----
 uefi-macros/tests/ui/pass/entry_no_args.rs |  9 +++++
 3 files changed, 49 insertions(+), 9 deletions(-)
 create mode 100644 uefi-macros/tests/ui/pass/entry_no_args.rs

diff --git a/uefi-macros/CHANGELOG.md b/uefi-macros/CHANGELOG.md
index ab6c8829b..2b3ffb624 100644
--- a/uefi-macros/CHANGELOG.md
+++ b/uefi-macros/CHANGELOG.md
@@ -1,5 +1,10 @@
 # uefi-macros - [Unreleased]
 
+## Changed
+
+- The `entry` macro now accepts a function with zero arguments in addition to
+  the two-argument form.
+
 
 # uefi-macros - 0.14.0 (2024-07-02)
 
diff --git a/uefi-macros/src/lib.rs b/uefi-macros/src/lib.rs
index 09a67716e..08be64f54 100644
--- a/uefi-macros/src/lib.rs
+++ b/uefi-macros/src/lib.rs
@@ -8,8 +8,8 @@ use proc_macro2::TokenStream as TokenStream2;
 use quote::{quote, quote_spanned, TokenStreamExt};
 use syn::spanned::Spanned;
 use syn::{
-    parse_macro_input, parse_quote, Error, Expr, ExprLit, ExprPath, FnArg, Ident, ItemFn,
-    ItemStruct, Lit, Pat, Visibility,
+    parse_macro_input, parse_quote, parse_quote_spanned, Error, Expr, ExprLit, ExprPath, FnArg,
+    Ident, ItemFn, ItemStruct, Lit, Pat, Visibility,
 };
 
 macro_rules! err {
@@ -121,18 +121,34 @@ fn get_function_arg_name(f: &ItemFn, arg_index: usize, errors: &mut TokenStream2
 /// Custom attribute for a UEFI executable entry point.
 ///
 /// This attribute modifies a function to mark it as the entry point for
-/// a UEFI executable. The function must have two parameters, [`Handle`]
-/// and [`SystemTable<Boot>`], and return a [`Status`]. The function can
-/// optionally be `unsafe`.
+/// a UEFI executable. The function:
+/// * Must return [`Status`].
+/// * Must have either zero parameters or two: [`Handle`] and [`SystemTable<Boot>`].
+/// * Can optionally be `unsafe`.
 ///
 /// Due to internal implementation details the parameters must both be
 /// named, so `arg` or `_arg` are allowed, but not `_`.
 ///
-/// The [`BootServices::set_image_handle`] function will be called
-/// automatically with the image [`Handle`] argument.
+/// The global system table pointer and global image handle will be set
+/// automatically.
 ///
 /// # Examples
 ///
+/// With no arguments:
+///
+/// ```no_run
+/// #![no_main]
+///
+/// use uefi::prelude::*;
+///
+/// #[entry]
+/// fn main() -> Status {
+///     Status::SUCCESS
+/// }
+/// ```
+///
+/// With two arguments:
+///
 /// ```no_run
 /// #![no_main]
 ///
@@ -180,6 +196,18 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
         ));
     }
 
+    let signature_span = f.sig.span();
+
+    // If the user doesn't specify any arguments to the entry function, fill in
+    // the image handle and system table arguments automatically.
+    if f.sig.inputs.is_empty() {
+        f.sig.inputs = parse_quote_spanned!(
+            signature_span=>
+                image_handle: ::uefi::Handle,
+                system_table: ::uefi::table::SystemTable<::uefi::table::Boot>
+        );
+    }
+
     let image_handle_ident = get_function_arg_name(&f, 0, &mut errors);
     let system_table_ident = get_function_arg_name(&f, 1, &mut errors);
 
@@ -188,8 +216,6 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
         return errors.into();
     }
 
-    let signature_span = f.sig.span();
-
     f.sig.abi = Some(syn::parse2(quote_spanned! (signature_span=> extern "efiapi")).unwrap());
 
     // allow the entry function to be unsafe (by moving the keyword around so that it actually works)
diff --git a/uefi-macros/tests/ui/pass/entry_no_args.rs b/uefi-macros/tests/ui/pass/entry_no_args.rs
new file mode 100644
index 000000000..e2271d5a2
--- /dev/null
+++ b/uefi-macros/tests/ui/pass/entry_no_args.rs
@@ -0,0 +1,9 @@
+use uefi::{entry, Status};
+
+#[entry]
+fn efi_main() -> Status {
+    Status::SUCCESS
+}
+
+// trybuild requires a `main` function.
+fn main() {}