Skip to main content

auto_error_into_macro/
lib.rs

1use proc_macro::{TokenStream, TokenTree};
2use quote::{format_ident, ToTokens};
3use syn::{parse_macro_input, parse_quote, FnArg, ItemFn, Pat, PatIdent, ReturnType};
4
5/// Wraps a function that returns a Result<T, E> in such a way that will always return the T type, relying on an `Into<T>` implemetation to exist for the E type.
6///
7/// # Remarks
8/// Use `#[auto_error_into(force_inline)]` to force the wrapped function to be inlined. This will only work on a subset of functions, and will never work on methods.
9#[proc_macro_attribute]
10pub fn auto_error_into(args: TokenStream, input: TokenStream) -> TokenStream {
11    let mut func = parse_macro_input!(input as ItemFn);
12
13    let mut original_signature = func.sig.clone();
14    let ReturnType::Type(_, t) = func.sig.output else { panic!("Unable to use auto_error_into on functions that return ()"); };
15
16    let block = func.block;
17    func.sig.output = parse_quote!(-> <#t as ::auto_error_into::__::ResultResolver>::Ok);
18
19    let mut args = args.into_iter();
20
21    let first = args.next();
22
23    if args.next().is_some() {
24        panic!("More than two arguments to macro");
25    }
26
27    match first {
28        Some(TokenTree::Ident(arg)) if arg.to_string() == "force_inline" => {
29            let parameter_names: Vec<_> = (0u32..)
30                .map(|num| format_ident!("param{}", num))
31                .take(original_signature.inputs.len())
32                .collect();
33
34            let new_args = original_signature
35                .inputs
36                .clone()
37                .into_iter()
38                .zip(&parameter_names)
39                .map(|(a, ident)| match a {
40                    FnArg::Typed(mut typ) => {
41                        typ.pat = Box::new(Pat::Ident(PatIdent {
42                            attrs: Vec::new(),
43                            by_ref: None,
44                            mutability: None,
45                            ident: ident.clone(),
46                            subpat: None,
47                        }));
48                        FnArg::Typed(typ)
49                    }
50                    _ => panic!("Cannot force inline on methods"),
51                });
52
53            func.sig.inputs = parse_quote!(#(#new_args),*);
54
55            original_signature.ident = format_ident!("__internal_invoke");
56            original_signature.abi = None;
57            func.block = Box::new(parse_quote!({
58                #[inline(always)] #original_signature #block
59                match __internal_invoke(#(#parameter_names),*) {
60                    Ok(o) => o,
61                    Err(e) => e.into(),
62                }
63            }));
64
65            func.into_token_stream().into()
66        }
67        Some(_) => panic!("Only supported argument is force_inline"),
68        _ => {
69            func.block = Box::new(parse_quote!({
70                match (move || -> #t #block)() {
71                    Ok(o) => o,
72                    Err(e) => e.into(),
73                }
74            }));
75            func.into_token_stream().into()
76        }
77    }
78}