view src/lib.rs @ 1:babb2f9cb5ad

All features working
author Lewin Bormann <lbo@spheniscida.de>
date Fri, 20 Nov 2020 23:20:03 +0100
parents 7401c615042c
children 2b71ed370156
line wrap: on
line source

use std::ops::Deref;
use std::iter::IntoIterator;

use proc_macro;
use syn::parse_macro_input;

fn tuple_type_from_types<TT: quote::ToTokens, T: Deref<Target=TT>>(
    types: &[T]) -> syn::TypeTuple {

    if types.len() == 1 {
        let typ = types[0].deref();
        syn::parse_quote! { (#typ,) }
    } else {
        let types = types.iter().map(|t| t.deref());
        syn::parse_quote! {
            (#(#types),*)
        }
    }
}

fn tuple_type_from_sig<'a, Iter: IntoIterator<Item = &'a syn::FnArg>>(
    args: Iter,
) -> syn::TypeTuple {
    let types = args
        .into_iter()
        .map(|a| match &a {
            &syn::FnArg::Receiver(_) => panic!("unexpected self in function argument list"),
            &syn::FnArg::Typed(ref typ) => typ.ty.as_ref(),
        })
        .collect::<Vec<&syn::Type>>();

    tuple_type_from_types(types.as_slice())
}

fn lazied_method_name(id: &syn::Ident) -> syn::Ident {
    syn::Ident::new((id.to_string() + "_lazy").as_str(), id.span())
}

/// Enable a function to be lazily called.
///
/// This produces both the function itself and a wrapper enabling the function to be called lazily.
///
/// ```
/// #[lazy]
/// fn myfunc(arg: i32) -> bool {
///     arg%2 == 0
/// }
///
/// fn main() {
///     println!("{} vs {}", myfunc(32), lazy!(myfunc(expensive_function())));
/// }
#[proc_macro_attribute]
pub fn lazy(_: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let parsedfn = parse_macro_input!(item as syn::ItemFn);
    let sig = &parsedfn.sig;
    let fnname = &sig.ident;
    let returntype = match &sig.output {
        syn::ReturnType::Default => {
            let emptytuplets = proc_macro::TokenStream::from(quote::quote! { () });
            let emptytuple: syn::Type = parse_macro_input!(emptytuplets as syn::Type);
            syn::ReturnType::Type(<syn::Token![->]>::default(), Box::new(emptytuple))
        }
        r => r.clone(),
    };

    let newname = lazied_method_name(fnname);
    let mut newsig = sig.clone();
    newsig.ident = newname.clone();

    match sig.inputs.first() {
        Some(syn::FnArg::Receiver(_)) => {
            // this is an impl function

            // Skip `self` argument.
            let mut args_iter = (&newsig.inputs).into_iter();
            let receiver = args_iter.next().unwrap();
            let tuty = tuple_type_from_sig(args_iter);
            // We are not evaluating for the receiver, thus one fewer argument.
            let args_ixs: Vec<syn::Index> = (0..newsig.inputs.len()-1).map(syn::Index::from).collect();
            (quote::quote! {
                #parsedfn

                fn #newname<'lazyvallt>(#receiver, lzy: ::lazycall_support::LazyVal<'lazyvallt, #tuty>) #returntype {
                    use ::lazycall_support::Evaluate;
                    let evald = lzy.eval();
                    self.#fnname(#(evald.#args_ixs),*)
                }
            }).into()
        }
        Some(syn::FnArg::Typed(_)) => {
            // This is a freestanding function
            let tuty = tuple_type_from_sig(&newsig.inputs);
            let args_ixs: Vec<syn::Index> = (0..newsig.inputs.len()).map(syn::Index::from).collect();
            (quote::quote! {
                #parsedfn

                fn #newname<'lazyvallt>(lzy: ::lazycall_support::LazyVal<'lazyvallt, #tuty>) #returntype {
                    use ::lazycall_support::Evaluate;
                    let evald = lzy.eval();
                    #fnname(#(evald.#args_ixs),*)
                }
            }).into()
        }
        None => {
            // This is a freestanding function without arguments.
            // fn f() -> ResultT { ... } => fn f_lazy() { ... } (no body modification)
            let mut updatedfn = parsedfn.clone();
            updatedfn.sig = newsig;
            (quote::quote! {
                #updatedfn

                #parsedfn
            })
            .into()
        }
    }
}

#[proc_macro]
pub fn lazycall(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
    use syn::spanned::Spanned;
    use quote::ToTokens;
    let funcall = parse_macro_input!(item as syn::Expr);

    match funcall {
        syn::Expr::MethodCall(mc) => {
            let receiver = &mc.receiver;
            let args = (&mc.args).into_iter().collect::<Vec<&syn::Expr>>();
            let renamed_method = lazied_method_name(&mc.method);

            let tuple: syn::Expr;
            if args.len() == 1 {
                let typ = args[0].deref();
                tuple = syn::parse_quote! { (#typ,) };
            } else {
                let args = args.iter().map(|t| t.deref());
                tuple = syn::parse_quote! {
                    (#(#args),*)
                };
            }

            (quote::quote! { #receiver.#renamed_method(::lazycall_support::LazyVal::new(&mut || #tuple )) }).into()
        },
        syn::Expr::Call(c) => {
            let func = &c.func;
            let renamed_func;
            match func.deref() {
                &syn::Expr::Path(ref ep) => {
                    let path = &ep.path;
                    let last = path.segments.last().unwrap();
                    renamed_func = lazied_method_name(&last.ident);
                },
                other => { println!("{:?}", other.to_token_stream()); unimplemented!() },
            }
            let args = (&c.args).into_iter().collect::<Vec<&syn::Expr>>();

            // We should be able to do this:
            //let tuple = tuple_type_from_types(&args);
            // but rustc has a bug: https://gist.github.com/dermesser/902c08436bb7241f1d560f42880e0376
            // so we do it manually.

            let tuple: syn::Expr;
            if args.len() == 1 {
                let typ = args[0].deref();
                tuple = syn::parse_quote! { (#typ,) };
            } else {
                let args = args.iter().map(|t| t.deref());
                tuple = syn::parse_quote! {
                    (#(#args),*)
                };
            }
            let result = (quote::quote! { #renamed_func(::lazycall_support::LazyVal::new(&mut || #tuple )) }).into();
            result
        },
        _ => syn::Error::new(funcall.span(), "syntax for lazycall!: lazycall!([receiver.]myfunction(a, b, c));").to_compile_error().into(),
    }
}