derive_termination/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as TokenStream2;
3use quote::quote;
4use syn::{Data, DeriveInput, Fields, LitInt, parse_macro_input, spanned::Spanned};
5
6#[proc_macro_derive(Termination, attributes(exit_code))]
7pub fn derive_termination(input: TokenStream) -> TokenStream {
8    let input: DeriveInput = parse_macro_input!(input);
9    match build_termination(&input) {
10        Ok(tokens) => tokens.into(),
11        Err(e) => e.into_compile_error().into(),
12    }
13}
14
15fn build_termination(input: &DeriveInput) -> syn::Result<TokenStream2> {
16    let enum_name = &input.ident;
17    let Data::Enum(enum_data) = &input.data else {
18        return Err(syn::Error::new(
19            input.span(),
20            "Termination can only be derived for enum",
21        ));
22    };
23
24    let mut variant_to_exit_code = Vec::new();
25    for variant in &enum_data.variants {
26        for attr in &variant.attrs {
27            let path = attr.path();
28            if path.is_ident("exit_code") {
29                let code: LitInt = attr.parse_args()?;
30                variant_to_exit_code.push((variant, code));
31            }
32        }
33    }
34
35    let arms: TokenStream2 = variant_to_exit_code
36        .into_iter()
37        .map(|(variant, code)| {
38            let ident = &variant.ident;
39            let variant = match variant.fields {
40                Fields::Unit => quote! [ Self::#ident ],
41                Fields::Unnamed(_) => quote! [ Self::#ident ( .. ) ],
42                Fields::Named(_) => quote! [ Self::#ident { .. } ],
43            };
44            quote! {
45                #variant => ::std::process::ExitCode::from(#code),
46            }
47        })
48        .collect();
49
50    Ok(quote! {
51        impl ::std::process::Termination for #enum_name {
52            fn report(self) -> ::std::process::ExitCode {
53                match self {
54                    #arms
55                }
56            }
57        }
58    })
59}