esp-hal-procmacros 0.23.0

use std::{cell::RefCell, fmt::Display, thread};

use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, quote};
use syn::{Attribute, ItemFn, ReturnType, Type};

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Unified `#[main]` dispatcher.
///
/// Parses the function once, runs all validation in a single pass, then
/// dispatches to the appropriate code generator.
pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
    let f: ItemFn = match syn::parse2(item) {
        Ok(f) => f,
        Err(e) => return e.to_compile_error(),
    };

    let ctxt = Ctxt::new();

    // ---- common checks (both blocking and async) ----

    if !args.is_empty() {
        ctxt.error_at_callsite("This attribute accepts no arguments");
    }
    if !f.sig.generics.params.is_empty() {
        ctxt.error_spanned_by(&f.sig, "main function must not be generic");
    }
    if f.sig.generics.where_clause.is_some() {
        ctxt.error_spanned_by(&f.sig, "main function must not have `where` clauses");
    }
    if f.sig.abi.is_some() {
        ctxt.error_spanned_by(&f.sig, "main function must not have an ABI qualifier");
    }
    if f.sig.variadic.is_some() {
        ctxt.error_spanned_by(&f.sig, "main function must not be variadic");
    }
    match &f.sig.output {
        ReturnType::Default => {}
        ReturnType::Type(_, ty) => match &**ty {
            Type::Tuple(tuple) if tuple.elems.is_empty() => {}
            Type::Never(_) => {}
            _ => ctxt.error_spanned_by(
                &f.sig,
                "main function must either not return a value, return `()` or return `!`",
            ),
        },
    }

    // ---- async-specific check ----

    if f.sig.asyncness.is_some() && f.sig.inputs.len() != 1 {
        ctxt.error_spanned_by(&f.sig, "main function must have 1 argument: the spawner.");
    }

    if let Err(errors) = ctxt.check() {
        return errors;
    }

    if f.sig.asyncness.is_some() {
        gen_async(f)
    } else {
        gen_blocking(f)
    }
}

// ---------------------------------------------------------------------------
// Code generators (no validation)
// ---------------------------------------------------------------------------

fn gen_blocking(f: ItemFn) -> TokenStream {
    let root = esp_hal_crate();
    quote! {
        #[doc = "The main entry point of the firmware, generated by the `#[main]` macro."]
        #[#root::__macro_implementation::__entry]
        #f
    }
}

fn gen_async(f: ItemFn) -> TokenStream {
    let ItemFn {
        attrs: fattrs,
        block: f_body,
        sig,
        ..
    } = f;
    let fargs = sig.inputs;
    let out = sig.output;

    let lint_attrs: Vec<&Attribute> = fattrs
        .iter()
        .filter(|a| {
            a.path().is_ident("deny") || a.path().is_ident("allow") || a.path().is_ident("warn")
        })
        .collect();

    let root = esp_hal_crate();

    quote! {
        #(#lint_attrs)*
        #[doc(hidden)]
        pub(crate) mod __main {
            use super::*;

            #[doc(hidden)]
            #(#fattrs)*
            #[::embassy_executor::task()]
            async fn __embassy_main(#fargs) #out {
                #f_body
            }

            #[doc(hidden)]
            unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
                ::core::mem::transmute(t)
            }

            #(#fattrs)*
            #[#root::main]
            fn main() -> ! {
                let mut executor = ::esp_rtos::embassy::Executor::new();
                let executor = unsafe { __make_static(&mut executor) };
                executor.run(|spawner| {
                    spawner.spawn(__embassy_main(spawner).unwrap());
                })
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Shared helpers
// ---------------------------------------------------------------------------

fn esp_hal_crate() -> syn::Ident {
    match proc_macro_crate::crate_name("esp-hal") {
        Ok(proc_macro_crate::FoundCrate::Name(ref name)) => quote::format_ident!("{name}"),
        _ => quote::format_ident!("esp_hal"),
    }
}

/// Collects multiple errors and emits them all at once.
///
/// Must be consumed via [`Ctxt::check`]; dropping without checking panics.
struct Ctxt {
    errors: RefCell<Option<Vec<syn::Error>>>,
}

impl Ctxt {
    fn new() -> Self {
        Ctxt {
            errors: RefCell::new(Some(Vec::new())),
        }
    }

    fn error_spanned_by<A: ToTokens, T: Display>(&self, obj: A, msg: T) {
        self.errors
            .borrow_mut()
            .as_mut()
            .unwrap()
            .push(syn::Error::new_spanned(obj.into_token_stream(), msg));
    }

    fn error_at_callsite<T: Display>(&self, msg: T) {
        self.errors
            .borrow_mut()
            .as_mut()
            .unwrap()
            .push(syn::Error::new(Span::call_site(), msg));
    }

    fn check(self) -> Result<(), TokenStream> {
        let errors = self.errors.borrow_mut().take().unwrap();
        match errors.len() {
            0 => Ok(()),
            _ => Err(errors.iter().map(syn::Error::to_compile_error).collect()),
        }
    }
}

impl Drop for Ctxt {
    fn drop(&mut self) {
        if !thread::panicking() && self.errors.borrow().is_some() {
            panic!("forgot to check for errors");
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    // ---- blocking path ----

    #[test]
    fn test_blocking_basic() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { fn main() -> ! { loop {} } }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                #[doc = "The main entry point of the firmware, generated by the `#[main]` macro."]
                #[esp_hal::__macro_implementation::__entry]
                fn main() -> ! { loop {} }
            }
            .to_string()
        );
    }

    #[test]
    fn test_blocking_rejects_args() {
        let result = main(
            quote::quote! { non_empty }.into(),
            quote::quote! { fn main() -> ! {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "This attribute accepts no arguments" }
            }
            .to_string()
        );
    }

    #[test]
    fn test_blocking_rejects_bad_return_type() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { fn main() -> u32 { 0 } }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "main function must either not return a value, return `()` or return `!`" }
            }
            .to_string()
        );
    }

    // ---- async path ----

    #[test]
    fn test_async_basic() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo(spawner: Spawner) {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                #[doc(hidden)]
                pub(crate) mod __main {
                    use super::*;
                    #[doc(hidden)]
                    #[::embassy_executor::task()]
                    async fn __embassy_main(spawner: Spawner) {
                        {}
                    }
                    #[doc(hidden)]
                    unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
                        ::core::mem::transmute(t)
                    }
                    #[esp_hal::main]
                    fn main() -> ! {
                        let mut executor = ::esp_rtos::embassy::Executor::new();
                        let executor = unsafe { __make_static(&mut executor) };
                        executor.run(|spawner| {
                            spawner.spawn(__embassy_main(spawner).unwrap());
                        })
                    }
                }
            }
            .to_string()
        );
    }

    #[test]
    fn test_blocking_fn_routes_to_entry() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { fn foo(spawner: Spawner) {} }.into(),
        );

        // non-async → blocking generator → wraps with __entry
        assert!(
            result.to_string().contains("__entry"),
            "expected __entry in:\n{result}"
        );
    }

    #[test]
    fn test_async_no_arg() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo() {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "main function must have 1 argument: the spawner." }
            }
            .to_string()
        );
    }

    #[test]
    fn test_not_generic() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo<S>(spawner: S) {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "main function must not be generic" }
            }
            .to_string()
        );
    }

    #[test]
    fn test_not_abi() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async extern "C" fn foo(spawner: Spawner) {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "main function must not have an ABI qualifier" }
            }
            .to_string()
        );
    }

    #[test]
    fn test_not_variadic() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo(spawner: ...) {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "main function must not be variadic" }
                ::core::compile_error! { "main function must have 1 argument: the spawner." }
            }
            .to_string()
        );
    }

    #[test]
    fn test_not_return_value() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo(spawner: Spawner) -> u32 {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                ::core::compile_error! { "main function must either not return a value, return `()` or return `!`" }
            }
            .to_string()
        );
    }

    #[test]
    fn test_async_return_never() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo(spawner: Spawner) -> ! {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                #[doc(hidden)]
                pub(crate) mod __main {
                    use super::*;
                    #[doc(hidden)]
                    #[::embassy_executor::task()]
                    async fn __embassy_main(spawner: Spawner) -> ! {
                        {}
                    }
                    #[doc(hidden)]
                    unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
                        ::core::mem::transmute(t)
                    }
                    #[esp_hal::main]
                    fn main() -> ! {
                        let mut executor = ::esp_rtos::embassy::Executor::new();
                        let executor = unsafe { __make_static(&mut executor) };
                        executor.run(|spawner| {
                            spawner.spawn(__embassy_main(spawner).unwrap());
                        })
                    }
                }
            }
            .to_string()
        );
    }

    #[test]
    fn test_async_return_unit_tuple() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { async fn foo(spawner: Spawner) -> () {} }.into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                #[doc(hidden)]
                pub(crate) mod __main {
                    use super::*;
                    #[doc(hidden)]
                    #[::embassy_executor::task()]
                    async fn __embassy_main(spawner: Spawner) -> () {
                        {}
                    }
                    #[doc(hidden)]
                    unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
                        ::core::mem::transmute(t)
                    }
                    #[esp_hal::main]
                    fn main() -> ! {
                        let mut executor = ::esp_rtos::embassy::Executor::new();
                        let executor = unsafe { __make_static(&mut executor) };
                        executor.run(|spawner| {
                            spawner.spawn(__embassy_main(spawner).unwrap());
                        })
                    }
                }
            }
            .to_string()
        );
    }

    #[test]
    fn test_propagate_lint_attrs() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! {
                #[allow(allowed)]
                #[deny(denied)]
                #[warn(warning)]
                #[ram]
                async fn foo(spawner: Spawner) -> () {}
            }
            .into(),
        );

        assert_eq!(
            result.to_string(),
            quote::quote! {
                #[allow(allowed)]
                #[deny(denied)]
                #[warn(warning)]
                #[doc(hidden)]
                pub(crate) mod __main {
                    use super::*;
                    #[doc(hidden)]
                    #[allow(allowed)]
                    #[deny(denied)]
                    #[warn(warning)]
                    #[ram]
                    #[::embassy_executor::task()]
                    async fn __embassy_main(spawner: Spawner) -> () {
                        {}
                    }
                    #[doc(hidden)]
                    unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
                        ::core::mem::transmute(t)
                    }
                    #[allow(allowed)]
                    #[deny(denied)]
                    #[warn(warning)]
                    #[ram]
                    #[esp_hal::main]
                    fn main() -> ! {
                        let mut executor = ::esp_rtos::embassy::Executor::new();
                        let executor = unsafe { __make_static(&mut executor) };
                        executor.run(|spawner| {
                            spawner.spawn(__embassy_main(spawner).unwrap());
                        })
                    }
                }
            }
            .to_string()
        );
    }

    // ---- parse error ----

    #[test]
    fn test_invalid_input() {
        let result = main(
            quote::quote! {}.into(),
            quote::quote! { not a function at all @@ }.into(),
        );

        assert!(
            result.to_string().contains("compile_error"),
            "expected compile_error in:\n{result}"
        );
    }
}