rayon_macro_hack/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro_hack::proc_macro_hack;
5use quote::quote;
6use syn::fold::{self, Fold};
7use syn::parse_macro_input;
8
9#[proc_macro_hack]
10pub fn parallel(input: TokenStream) -> TokenStream {
11    // for #pat in #expr { #body }
12    let syn::ExprForLoop {
13        attrs,
14        label,
15        pat,
16        expr,
17        body,
18        ..
19    } = parse_macro_input!(input);
20
21    assert!(attrs.is_empty(), "loop attributes are not supported");
22
23    let mut transform = TransformBody::new(label.as_ref());
24    let body = transform.fold_block(body);
25
26    let imp = quote! { $crate::rayon_macro_impl };
27    TokenStream::from(quote! {
28        #imp::ParallelIterator::for_each(
29            #imp::IntoParallelIterator::into_par_iter(#expr),
30            |#pat| #label while { #body; false } {},
31        )
32    })
33}
34
35struct TransformBody {
36    for_ident: Option<syn::Ident>,
37    loop_level: usize,
38    try_level: usize,
39}
40
41impl TransformBody {
42    fn new(for_label: Option<&syn::Label>) -> Self {
43        let for_ident = for_label.map(|l| l.name.ident.clone());
44        TransformBody {
45            for_ident,
46            loop_level: 0,
47            try_level: 0,
48        }
49    }
50
51    fn is_controlled(&self, label: &Option<syn::Lifetime>) -> bool {
52        if let Some(label) = label {
53            self.for_ident.as_ref() == Some(&label.ident)
54        } else {
55            self.loop_level == 0
56        }
57    }
58
59    fn take_shadowed(&mut self, label: &Option<syn::Label>) -> Option<syn::Ident> {
60        if let Some(label) = label {
61            if self.for_ident.as_ref() == Some(&label.name.ident) {
62                return self.for_ident.take();
63            }
64        }
65        None
66    }
67}
68
69impl Fold for TransformBody {
70    // We don't care about other items defined in our block
71    fn fold_item(&mut self, i: syn::Item) -> syn::Item {
72        i
73    }
74
75    fn fold_expr(&mut self, i: syn::Expr) -> syn::Expr {
76        use syn::Expr::*;
77        match i {
78            // Control flow doesn't escape closures, so skip them.
79            Async(_) | Closure(_) => return i,
80
81            // Inner loops might intercept `break` and `continue`
82            ForLoop(syn::ExprForLoop { ref label, .. })
83            | Loop(syn::ExprLoop { ref label, .. })
84            | While(syn::ExprWhile { ref label, .. }) => {
85                let shadowed = self.take_shadowed(label);
86                self.loop_level += 1;
87                let ret = fold::fold_expr(self, i);
88                self.loop_level -= 1;
89                if shadowed.is_some() {
90                    self.for_ident = shadowed;
91                }
92                ret
93            }
94
95            Continue(c) => {
96                if self.is_controlled(&c.label) {
97                    Return(syn::ExprReturn {
98                        attrs: c.attrs,
99                        return_token: syn::token::Return {
100                            span: c.continue_token.span,
101                        },
102                        expr: None,
103                    })
104                } else {
105                    Continue(fold::fold_expr_continue(self, c))
106                }
107            }
108
109            Break(b) => {
110                if self.is_controlled(&b.label) {
111                    panic!("`break` is not supported yet")
112                } else {
113                    Break(fold::fold_expr_break(self, b))
114                }
115            }
116
117            Return(_) => panic!("`return` is not supported yet"),
118
119            Try(_) if self.try_level == 0 => panic!("`?` is not supported yet"),
120            TryBlock(_) => {
121                self.try_level += 1;
122                let ret = fold::fold_expr(self, i);
123                self.try_level -= 1;
124                ret
125            }
126
127            Await(_) => panic!("`.await` is not supported"),
128            Yield(_) => panic!("`yield` is not supported"),
129
130            _ => fold::fold_expr(self, i),
131        }
132    }
133}
rayon_macro_hack/lib.rs

rayon_macro_hack/
lib.rs
