yield_closures_impl/
lib.rs1use 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}