yield_closures_impl/
lib.rs

1//! # yield-closures
2//!
3//! An implementation of [MCP-49](https://github.com/rust-lang/lang-team/issues/49).
4//!
5//! ```rust
6//! #[test]
7//! fn decode_escape_string() {
8//!     let escaped_text = "Hello,\x20world!\\n";
9//!     let text: String = escaped_text
10//!         .chars()
11//!         .filter_map(co!(|c| {
12//!             loop {
13//!                 if c != '\\' {
14//!                     Not escaped
15//!                     yield Some(c);
16//!                     continue;
17//!                 }
18//!
19//!                 Go past the \
20//!                 yield None;
21//!
22//!                 Unescaped-char
23//!                 match c {
24//!                     Hexadecimal
25//!                     'x' => {
26//!                         yield None; Go past the x
27//!                         let most = c.to_digit(16);
28//!                         yield None; Go past the first digit
29//!                         let least = c.to_digit(16);
30//!                         Yield the decoded char if valid
31//!                         yield (|| char::from_u32(most? << 4 | least?))()
32//!                     }
33//!                     Simple escapes
34//!                     'n' => yield Some('\n'),
35//!                     'r' => yield Some('\r'),
36//!                     't' => yield Some('\t'),
37//!                     '0' => yield Some('\0'),
38//!                     '\\' => yield Some('\\'),
39//!                     Unnecessary escape
40//!                     _ => yield Some(c),
41//!                 }
42//!             }
43//!         }))
44//!         .collect();
45//!     assert_eq!(text, "Hello, world!\n");
46//! }
47//! ```
48//!
49//! For the details of the proposal, see https://lang-team.rust-lang.org/design_notes/general_coroutines.html.
50//!
51//! Differences between this implementation and the proposal are summarized below:
52//!
53//! - This crate offers a macro implementation. It works with the stable Rust.
54//! - No `FnPin` is provided. Yield closures made with this crate use `Box::pin` internally and hence `FnMut`.
55//! - In yield closures, one cannot use `return` expressions.
56//! - The body of a yield closure must be explosive i.e. must not return and typed by the `!` type. Thus it is compatible with both of the two designs of yield closures discussed in the document of MCP-49: poisoning by default or not.
57
58use proc_macro::*;
59use syn::{fold::Fold, ReturnType};
60
61#[proc_macro]
62pub fn co(input: TokenStream) -> TokenStream {
63    let closure: syn::ExprClosure = syn::parse2(input.into()).unwrap();
64
65    if !closure.attrs.is_empty() {
66        unimplemented!("attributes");
67    }
68    if closure.asyncness.is_some() {
69        unimplemented!("async closure");
70    }
71    if closure.movability.is_some() {
72        unimplemented!("movability");
73    }
74
75    let mut types: Vec<Option<syn::Type>> = vec![];
76    let inputs = closure
77        .inputs
78        .iter()
79        .map(|input| match input {
80            syn::Pat::Ident(ident) => {
81                if !ident.attrs.is_empty() {
82                    unimplemented!("attributes");
83                }
84                if ident.subpat.is_some() {
85                    unimplemented!("subpatterns");
86                }
87                if ident.by_ref.is_some() {
88                    unimplemented!("reference patterns");
89                }
90                if ident.mutability.is_some() {
91                    unimplemented!("mutable parameter");
92                }
93                types.push(None);
94                ident.ident.clone()
95            }
96            syn::Pat::Type(syn::PatType {
97                attrs,
98                pat,
99                colon_token: _colon_token,
100                ty,
101            }) => {
102                if !attrs.is_empty() {
103                    unimplemented!("attributes");
104                }
105                match &**pat {
106                    syn::Pat::Ident(ident) => {
107                        if !ident.attrs.is_empty() {
108                            unimplemented!("attributes");
109                        }
110                        if ident.subpat.is_some() {
111                            unimplemented!("subpatterns");
112                        }
113                        if ident.by_ref.is_some() {
114                            unimplemented!("reference patterns");
115                        }
116                        if ident.mutability.is_some() {
117                            unimplemented!("mutable parameter");
118                        }
119                        types.push(Some((**ty).clone()));
120                        ident.ident.clone()
121                    }
122                    _ => {
123                        unimplemented!("patterns inside type patterns must be variable");
124                    }
125                }
126            }
127            p => {
128                unimplemented!("input pattern must be variable: {:?}", p);
129            }
130        })
131        .collect::<Vec<_>>();
132
133    let ret_ty;
134    match &closure.output {
135        ReturnType::Default => {
136            ret_ty = None;
137        }
138        ReturnType::Type(_, ty) => {
139            ret_ty = Some(ty.clone());
140        }
141    }
142
143    let body = ReplaceYields { inputs: &inputs }.fold_expr(*closure.body);
144
145    let capture = closure.capture;
146
147    let co_imp_fn = match inputs.len() {
148        0 => {
149            let ret_ty = ret_ty
150                .map(|t| quote::quote!(#t))
151                .unwrap_or(quote::quote!(_));
152            quote::quote!(::yield_closures::co0::<_, #ret_ty, _>)
153        }
154        1 => {
155            let ret_ty = ret_ty
156                .map(|t| quote::quote!(#t))
157                .unwrap_or(quote::quote!(_));
158            let a0 = types[0]
159                .as_ref()
160                .map(|t| quote::quote!(#t))
161                .unwrap_or(quote::quote!(_));
162            quote::quote!(::yield_closures::co::<_, #a0, #ret_ty, _>)
163        }
164        2 => {
165            let ret_ty = ret_ty
166                .map(|t| quote::quote!(#t))
167                .unwrap_or(quote::quote!(_));
168            let a0 = types[0]
169                .as_ref()
170                .map(|t| quote::quote!(#t))
171                .unwrap_or(quote::quote!(_));
172            let a1 = types[1]
173                .as_ref()
174                .map(|t| quote::quote!(#t))
175                .unwrap_or(quote::quote!(_));
176            quote::quote!(::yield_closures::co2::<_, #a0, #a1, #ret_ty, _>)
177        }
178        3 => {
179            let ret_ty = ret_ty
180                .map(|t| quote::quote!(#t))
181                .unwrap_or(quote::quote!(_));
182            let a0 = types[0]
183                .as_ref()
184                .map(|t| quote::quote!(#t))
185                .unwrap_or(quote::quote!(_));
186            let a1 = types[1]
187                .as_ref()
188                .map(|t| quote::quote!(#t))
189                .unwrap_or(quote::quote!(_));
190            let a2 = types[1]
191                .as_ref()
192                .map(|t| quote::quote!(#t))
193                .unwrap_or(quote::quote!(_));
194            quote::quote!(::yield_closures::co3::<_, #a0, #a1, #a2, #ret_ty, _>)
195        }
196        _ => {
197            unimplemented!("the number of inputs is too large");
198        }
199    };
200    quote::quote!(
201        #co_imp_fn(|__arg_rx, __yield_tx| async #capture {
202            let (__arg_rx, __yield_tx) = async move { (__arg_rx, __yield_tx) }.await; do partial move
203            let (#( mut #inputs ),*);
204            ::yield_closures::reassign_args!(__arg_rx, #( #inputs, )*);
205            #body
206        })
207    )
208    .into()
209}
210
211struct ReplaceYields<'a> {
212    inputs: &'a [syn::Ident],
213}
214
215impl<'a> Fold for ReplaceYields<'a> {
216    fn fold_expr(&mut self, i: syn::Expr) -> syn::Expr {
217        match i {
218            syn::Expr::Yield(i) => {
219                let expr = if let Some(expr) = i.expr {
220                    self.fold_expr(*expr)
221                } else {
222                    syn::parse2(quote::quote!(())).unwrap()
223                };
224                let inputs = self.inputs;
225                syn::parse2(quote::quote! {{
226                    __yield_tx.send(#expr).unwrap();
227                    ::yield_closures::drop_args!(#( #inputs, )*);
228                    ::yield_closures::pend_once().await;
229                    ::yield_closures::reassign_args!(__arg_rx, #( #inputs, )*);
230                }})
231                .unwrap()
232            }
233            syn::Expr::Await(_) => {
234                unimplemented!("await expressions in yield closures");
235            }
236            syn::Expr::Async(_) => {
237                unimplemented!("async blocks in yield closures");
238            }
239            syn::Expr::TryBlock(_) => {
240                unimplemented!("try blocks in yield closures");
241            }
242            syn::Expr::Return(_) => {
243                panic!("return expressions in yield closures are unsupported");
244            }
245            _ => syn::fold::fold_expr(self, i),
246        }
247    }
248}