use proc_macro::TokenStream;
use proc_macro2::{Spacing, Span, TokenStream as TokenStream2, TokenTree};
use quote::{format_ident, quote};
use syn::{BinOp, Expr};
#[proc_macro]
pub fn checked(input: TokenStream) -> TokenStream {
let input2: TokenStream2 = input.into();
match parse_input(input2) {
Ok((expr, error)) => match expand(&expr, error.as_ref(), 0) {
Ok(tokens) => tokens.into(),
Err(err) => err.to_compile_error().into(),
},
Err(err) => err.to_compile_error().into(),
}
}
fn parse_input(tokens: TokenStream2) -> syn::Result<(Expr, Option<Expr>)> {
let mut expr_tokens = TokenStream2::new();
let mut error_tokens = TokenStream2::new();
let mut found_at = false;
for tt in tokens {
if !found_at {
if let TokenTree::Punct(ref punct) = tt {
if punct.as_char() == '@' && punct.spacing() == Spacing::Alone {
found_at = true;
continue;
}
}
expr_tokens.extend(std::iter::once(tt));
} else {
error_tokens.extend(std::iter::once(tt));
}
}
if expr_tokens.is_empty() {
return Err(syn::Error::new(
Span::call_site(),
"expected an arithmetic expression",
));
}
let expr: Expr = syn::parse2(expr_tokens)?;
let error = if found_at {
if error_tokens.is_empty() {
return Err(syn::Error::new(
Span::call_site(),
"expected an error expression after `@`",
));
}
Some(syn::parse2::<Expr>(error_tokens)?)
} else {
None
};
Ok((expr, error))
}
fn expand(expr: &Expr, error: Option<&Expr>, depth: usize) -> syn::Result<TokenStream2> {
match expr {
Expr::Binary(binary) => {
let method = method_name(&binary.op)?;
let lhs = expand(&binary.left, error, depth + 1)?;
let rhs = expand(&binary.right, error, depth + 1)?;
let l = format_ident!("__l{}", depth);
let r = format_ident!("__r{}", depth);
let combine = match error {
Some(err) => quote! { #l.#method(#r).ok_or(#err) },
None => quote! { #l.#method(#r) },
};
Ok(quote! {
(#lhs).and_then(|#l| (#rhs).and_then(|#r| #combine))
})
}
Expr::Paren(paren) => expand(&paren.expr, error, depth),
Expr::Group(group) => expand(&group.expr, error, depth),
leaf => Ok(match error {
Some(_) => quote! { ::core::result::Result::<_, _>::Ok(#leaf) },
None => quote! { ::core::option::Option::Some(#leaf) },
}),
}
}
fn method_name(op: &BinOp) -> syn::Result<proc_macro2::Ident> {
let name = match op {
BinOp::Add(_) => "checked_add",
BinOp::Sub(_) => "checked_sub",
BinOp::Mul(_) => "checked_mul",
BinOp::Div(_) => "checked_div",
BinOp::Rem(_) => "checked_rem",
other => {
return Err(syn::Error::new_spanned(
other,
"`checked!` only supports the `+`, `-`, `*`, `/`, and `%` operators",
));
}
};
Ok(format_ident!("{}", name))
}