1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

//! This crate provides procedural macros for [Quit].
//!
//! **Do not add this crate as a dependency.** It has no backward compatibility
//! guarantees. Use the macros re-exported from [Quit] instead.
//!
//! [Quit]: https://crates.io/crates/quit

#![doc(html_root_url = "https://docs.rs/quit_macros/*")]
#![forbid(unsafe_code)]
#![warn(unused_results)]

#[allow(unused_extern_crates)]
extern crate proc_macro;

use std::fmt::Display;

use proc_macro::TokenStream;

use quote::quote;
use quote::ToTokens;

use syn::parse_macro_input;
use syn::Error as SynError;
use syn::ItemFn;

/// Modifies the main function to exit with the code passed to [`with_code`].
///
/// This attribute should always be attached to the main function. Otherwise,
/// the exit code of the program may be incorrect.
///
/// # Examples
///
/// ```
/// #[quit::main]
/// fn main() {}
/// ```
///
/// [`with_code`]: fn.with_code.html
#[inline]
#[proc_macro_attribute]
pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
    if let Some(arg) = args.into_iter().next() {
        return SynError::new(arg.span().into(), "arguments are not accepted")
            .to_compile_error()
            .into();
    }

    let input = parse_macro_input!(item as ItemFn);
    let signature = &input.sig;

    let name = &signature.ident;
    if name != "main" {
        return error(name, "`quit::main` can only be attached to `main`");
    }

    let attrs = &input.attrs;
    let visibility = &input.vis;
    let body = &input.block;
    return quote! {
        #(#attrs)*
        #visibility #signature {
            match ::std::panic::catch_unwind(|| { #body }) {
                Ok(result) => result,
                Err(payload) => {
                    if let Some(&::quit::_ExitCode(exit_code)) =
                        payload.downcast_ref()
                    {
                        ::std::process::exit(exit_code);
                    };
                    ::std::panic::resume_unwind(payload);
                },
            }
        }
    }
    .into();

    fn error<TMessage, TTokens>(
        tokens: TTokens,
        message: TMessage,
    ) -> TokenStream
    where
        TMessage: Display,
        TTokens: ToTokens,
    {
        SynError::new_spanned(tokens, message)
            .to_compile_error()
            .into()
    }
}