propane_macros/
lib.rs

1extern crate proc_macro;
2
3mod elision;
4
5use proc_macro::*;
6use syn::fold::Fold;
7
8#[proc_macro_attribute]
9pub fn generator(_: TokenStream, input: TokenStream) -> TokenStream {
10    if let Ok(item_fn) = syn::parse(input) {
11        let item_fn = Generator { outer_fn: true, lifetimes: vec![] }.fold_item_fn(item_fn);
12        quote::quote!(#item_fn).into()
13    } else {
14        panic!("#[generator] atribute can only be applied to functions");
15    }
16}
17
18struct Generator {
19    outer_fn: bool,
20    lifetimes: Vec<syn::Lifetime>,
21}
22
23impl Fold for Generator {
24    fn fold_item_fn(&mut self, mut i: syn::ItemFn) -> syn::ItemFn {
25        if !self.outer_fn { return i }
26
27        let inputs = elision::unelide_lifetimes(&mut i.sig.generics.params, i.sig.inputs);
28        self.lifetimes = i.sig.generics.lifetimes().map(|l| l.lifetime.clone()).collect();
29        let output = self.fold_return_type(i.sig.output);
30        let sig = syn::Signature { output, inputs, ..i.sig };
31
32        self.outer_fn = false;
33
34        let inner = self.fold_block(*i.block);
35        let block = Box::new(make_fn_block(&inner));
36
37        syn::ItemFn { sig, block, ..i }
38    }
39
40    fn fold_return_type(&mut self, i: syn::ReturnType) -> syn::ReturnType {
41        use syn::{ReturnType, Token};
42        use proc_macro2::Span;
43        if !self.outer_fn { return i; }
44        let (arrow, ret) = match i {
45            ReturnType::Default => (Token![->](Span::call_site()), syn::parse_str("()").unwrap()),
46            ReturnType::Type(arrow, ty) => (arrow, *ty),
47        };
48        let lifetimes = std::mem::replace(&mut self.lifetimes, vec![]);
49        let ret = syn::parse2(quote::quote!(
50                (impl Iterator<Item = #ret> #(+ #lifetimes )*)
51        )).unwrap();
52        ReturnType::Type(arrow, Box::new(ret))
53    }
54
55    fn fold_expr(&mut self, i: syn::Expr) -> syn::Expr {
56        if let syn::Expr::Try(syn::ExprTry { expr, .. }) = i {
57            syn::parse2(quote::quote!(propane::gen_try!(#expr))).unwrap()
58        } else {
59            syn::fold::fold_expr(self, i)
60        }
61    }
62}
63
64fn make_fn_block(inner: &syn::Block) -> syn::Block {
65    syn::parse2(quote::quote! {{
66        let __ret = || {
67            #inner;
68            #[allow(unreachable_code)]
69            {
70                return;
71                yield panic!();
72            }
73        };
74
75        #[allow(unreachable_code)]
76        propane::__internal::GenIter(__ret)
77    }}).unwrap()
78}
propane_macros/lib.rs

propane_macros/
lib.rs
