derive_termination/
lib.rs

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