vexide-macro 0.4.0

//! This crate provides procedural macros for [vexide](https://vexide.dev) crates.

use parse::{Attrs, MacroOpts};
use proc_macro::TokenStream;
use quote::quote;
use syn::{ItemFn, Signature, parse_macro_input};

mod parse;

const NO_SYNC_ERR: &str = "The vexide entrypoint must be marked `async`.";
const NO_UNSAFE_ERR: &str = "The vexide entrypoint must be not marked `unsafe`.";
const WRONG_ARGS_ERR: &str = "The vexide entrypoint must take a single parameter of type `vexide_devices::peripherals::Peripherals`";

fn verify_function_sig(sig: &Signature) -> Result<(), syn::Error> {
    let mut error = None;

    if sig.asyncness.is_none() {
        let message = syn::Error::new_spanned(sig, NO_SYNC_ERR);
        error.replace(message);
    }
    if sig.unsafety.is_some() {
        let message = syn::Error::new_spanned(sig, NO_UNSAFE_ERR);
        match error {
            Some(ref mut e) => e.combine(message),
            None => {
                error.replace(message);
            }
        }
    }
    if sig.inputs.len() != 1 {
        let message = syn::Error::new_spanned(sig, WRONG_ARGS_ERR);
        match error {
            Some(ref mut e) => e.combine(message),
            None => {
                error.replace(message);
            }
        }
    }

    match error {
        Some(e) => Err(e),
        None => Ok(()),
    }
}

fn make_code_sig(opts: MacroOpts) -> proc_macro2::TokenStream {
    let sig = if let Some(code_sig) = opts.code_sig {
        quote! { #code_sig }
    } else {
        quote! {  ::vexide::program::CodeSignature::new(
            ::vexide::program::ProgramType::User,
            ::vexide::program::ProgramOwner::Partner,
            ::vexide::program::ProgramOptions::empty(),
        ) }
    };

    quote! {
        #[cfg_attr(target_os = "vexos", unsafe(link_section = ".code_signature"))]
        #[used] // This is needed to prevent the linker from removing this object in release builds
        #[unsafe(no_mangle)]
        static __VEXIDE_CODE_SIGNATURE: ::vexide::program::CodeSignature = #sig;
    }
}

fn make_entrypoint(inner: &ItemFn, opts: MacroOpts) -> proc_macro2::TokenStream {
    match verify_function_sig(&inner.sig) {
        Ok(()) => {}
        Err(e) => return e.to_compile_error(),
    }
    let inner_ident = inner.sig.ident.clone();
    let ret_type = match &inner.sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    let banner_theme = if let Some(theme) = opts.banner_theme {
        quote! { #theme }
    } else {
        quote! { ::vexide::startup::banner::themes::THEME_DEFAULT }
    };

    let banner_print = if opts.banner_enabled {
        quote! {
            ::vexide::startup::banner::print(#banner_theme);
        }
    } else {
        quote! {}
    };

    quote! {
        fn main() -> #ret_type {
            unsafe {
                ::vexide::startup::startup();
            }

            #banner_print
            #inner

            ::vexide::runtime::block_on(
                #inner_ident(::vexide::peripherals::Peripherals::take().unwrap())
            )
        }
    }
}

/// vexide's entrypoint macro
///
/// Marks a function as the entrypoint for a vexide program. When the program is started, the `main`
/// function will be called with a single argument of type `Peripherals` which allows access to
/// device peripherals like motors, sensors, and the display.
///
/// The `main` function must be marked `async` and must not be marked `unsafe`. It may return any
/// type that implements `Termination`, which includes `()`, `!`, and `Result`.
///
/// # Parameters
///
/// The `main` attribute can be provided with parameters that alter the behavior of the program.
///
/// - `banner`: Allows for disabling or using a custom banner theme. When `enabled = false` the
///   banner will be disabled. `theme` can be set to a custom `BannerTheme` struct.
/// - `code_sig`: Allows using a custom `CodeSignature` struct to configure program behavior.
///
/// # Examples
///
/// The most basic usage of the `main` attribute is to mark an async function as the entrypoint for
/// a vexide program. The function must take a single argument of type `Peripherals`.
///
/// ```
/// use std::fmt::Write;
///
/// use vexide::prelude::*;
///
/// #[vexide::main]
/// async fn main(mut peripherals: Peripherals) {
///     write!(peripherals.display, "Hello, vexide!").unwrap();
/// }
/// ```
///
/// The `main` attribute can also be provided with parameters to customize the behavior of the
/// program.
///
/// This includes disabling the banner or using a custom banner theme:
///
/// ```
/// use vexide::prelude::*;
///
/// #[vexide::main(banner(enabled = false))]
/// async fn main(_p: Peripherals) {
///     println!("This is the only serial output from this program!")
/// }
/// ```
///
/// ```
/// use vexide::{prelude::*, startup::banner::themes::THEME_SYNTHWAVE};
///
/// #[vexide::main(banner(theme = THEME_SYNTHWAVE))]
/// async fn main(_p: Peripherals) {
///     println!("This program has a synthwave themed banner!")
/// }
/// ```
///
/// A custom code signature may be used to further configure the behavior of the program.
///
/// ```
/// use vexide::{
///     prelude::*,
///     program::{CodeSignature, ProgramOptions, ProgramOwner, ProgramType},
/// };
///
/// static CODE_SIG: CodeSignature = CodeSignature::new(
///     ProgramType::User,
///     ProgramOwner::Partner,
///     ProgramOptions::empty(),
/// );
///
/// #[vexide::main(code_sig = CODE_SIG)]
/// async fn main(_p: Peripherals) {
///     println!("Hello world!")
/// }
/// ```
#[proc_macro_attribute]
pub fn main(attrs: TokenStream, item: TokenStream) -> TokenStream {
    let item = parse_macro_input!(item as ItemFn);
    let opts = MacroOpts::from(parse_macro_input!(attrs as Attrs));

    let entrypoint = make_entrypoint(&item, opts.clone());
    let code_signature = make_code_sig(opts);

    quote! {
        const _: () = {
            #code_signature
        };

        #entrypoint
    }
    .into()
}

/// Prints a failure message indicating that the required features for the [`main`] macro are not
/// enabled.
#[proc_macro_attribute]
#[doc(hidden)]
pub fn main_fail(_args: TokenStream, _item: TokenStream) -> TokenStream {
    syn::Error::new(
        proc_macro2::Span::call_site(),
        "The #[vexide::main] macro requires the `core`, `async`, `startup`, and `devices` features to be enabled.",
    )
    .to_compile_error()
    .into()
}

/// Wraps a Rust unit test in vexide's async runtime.
///
/// This macro should be accompanied with an SDK provider capable of running vexide programs on a
/// host system for unit tests, such as `vex-sdk-mock`.
#[proc_macro_attribute]
pub fn test(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);

    // Ensure it's async
    if input.sig.asyncness.is_none() {
        return syn::Error::new_spanned(input.sig.fn_token, "#[vexide::test] requires an async fn")
            .to_compile_error()
            .into();
    }

    let vis = &input.vis;
    let ident = &input.sig.ident;
    let inputs = &input.sig.inputs;
    let block = &input.block;

    quote! {
        #[::core::prelude::v1::test]
        #vis fn #ident() {
            async fn #ident(#inputs) #block

            ::vexide::runtime::block_on(
                #ident(unsafe { ::vexide::peripherals::Peripherals::steal() })
            )
        }
    }
    .into()
}

/// Prints a failure message indicating that the required features for the [`test`] macro are not
/// enabled.
#[proc_macro_attribute]
#[doc(hidden)]
pub fn test_fail(_args: TokenStream, _item: TokenStream) -> TokenStream {
    syn::Error::new(
        proc_macro2::Span::call_site(),
        "The #[vexide::test] macro requires the `core`, `async`, `startup`, and `devices` features to be enabled.",
    )
    .to_compile_error()
    .into()
}

#[cfg(test)]
mod test {
    use quote::quote;
    use syn::{Ident, ItemFn};

    use super::{make_code_sig, make_entrypoint};
    use crate::{MacroOpts, NO_SYNC_ERR, NO_UNSAFE_ERR, WRONG_ARGS_ERR};

    #[test]
    fn wraps_main_fn() {
        let source = quote! {
            async fn main(_peripherals: Peripherals) {
                println!("Hello, world!");
            }
        };

        let input = syn::parse2::<ItemFn>(source.clone()).unwrap();
        let output = make_entrypoint(&input, MacroOpts::default());

        assert_eq!(
            output.to_string(),
            quote! {
                fn main() -> () {
                    unsafe {
                        ::vexide::startup::startup();
                    }

                    ::vexide::startup::banner::print(::vexide::startup::banner::themes::THEME_DEFAULT);
                    #source

                    ::vexide::runtime::block_on(
                        main(::vexide::peripherals::Peripherals::take().unwrap())
                    )
                }
            }
            .to_string()
        );
    }

    #[test]
    fn toggles_banner_using_parsed_opts() {
        let source = quote! {
            async fn main(_peripherals: Peripherals) {
                println!("Hello, world!");
            }
        };
        let input = syn::parse2::<ItemFn>(source.clone()).unwrap();
        let entrypoint = make_entrypoint(
            &input,
            MacroOpts {
                banner_enabled: false,
                banner_theme: None,
                code_sig: None,
            },
        );
        assert!(!entrypoint.to_string().contains("banner"));

        let entrypoint = make_entrypoint(
            &input,
            MacroOpts {
                banner_enabled: true,
                banner_theme: None,
                code_sig: None,
            },
        );
        assert!(entrypoint.to_string().contains("banner"));
    }

    #[test]
    fn uses_custom_code_sig_from_parsed_opts() {
        let code_sig = make_code_sig(MacroOpts {
            banner_enabled: false,
            banner_theme: None,
            code_sig: Some(Ident::new(
                "__custom_code_sig_ident__",
                proc_macro2::Span::call_site(),
            )),
        });

        assert!(code_sig.to_string().contains(
            "static __VEXIDE_CODE_SIGNATURE : :: vexide :: program :: CodeSignature = __custom_code_sig_ident__ ;"
        ));
    }

    #[test]
    fn requires_async() {
        let source = quote! {
            fn main(_peripherals: Peripherals) {
                println!("Hello, world!");
            }
        };

        let input = syn::parse2::<ItemFn>(source.clone()).unwrap();
        let output = make_entrypoint(&input, MacroOpts::default());

        assert!(output.to_string().contains(NO_SYNC_ERR));
    }

    #[test]
    fn requires_safe() {
        let source = quote! {
            async unsafe fn main(_peripherals: Peripherals) {
                println!("Hello, world!");
            }
        };

        let input = syn::parse2::<ItemFn>(source.clone()).unwrap();
        let output = make_entrypoint(&input, MacroOpts::default());

        assert!(output.to_string().contains(NO_UNSAFE_ERR));
    }

    #[test]
    fn disallows_0_args() {
        let source = quote! {
            async fn main() {
                println!("Hello, world!");
            }
        };

        let input = syn::parse2::<ItemFn>(source.clone()).unwrap();
        let output = make_entrypoint(&input, MacroOpts::default());

        assert!(output.to_string().contains(WRONG_ARGS_ERR));
    }

    #[test]
    fn disallows_2_args() {
        let source = quote! {
            async fn main(_peripherals: Peripherals, _other: Peripherals) {
                println!("Hello, world!");
            }
        };

        let input = syn::parse2::<ItemFn>(source.clone()).unwrap();
        let output = make_entrypoint(&input, MacroOpts::default());

        assert!(output.to_string().contains(WRONG_ARGS_ERR));
    }
}