syn_dissect_closure/
lib.rs

1use std::{collections::HashSet, fmt};
2
3use quote::ToTokens;
4use syn::{parse_quote, token::Brace, Block, Expr, ExprBlock, ExprClosure, Ident, Pat, Stmt};
5
6mod state;
7use state::{EnvStruct, State};
8
9struct DissectedClosure {
10	env_variables: HashSet<Ident>,
11	closure: ExprClosure,
12}
13impl fmt::Debug for DissectedClosure {
14	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
15		f.debug_struct("DissectedClosure")
16			.field("env_variables", &self.env_variables)
17			.field("closure", &self.closure.to_token_stream())
18			.finish()
19	}
20}
21
22#[derive(Default)]
23struct ClosureDissector {
24	known_env: HashSet<Ident>,
25	known_not_env: HashSet<Ident>,
26	env: Option<EnvStruct>,
27}
28impl ClosureDissector {
29	fn dissect(mut self, mut closure: ExprClosure) -> DissectedClosure {
30		// Convert closure to use block so any not_env_variables can be asserted.
31		closure.body = Box::new(match *closure.body {
32			Expr::Block(block) => Expr::Block(block),
33			// Avoid clippy::unused_unit
34			Expr::Tuple(tuple) if tuple.elems.is_empty() => Expr::Tuple(tuple),
35			expr => Expr::Block(ExprBlock {
36				attrs: vec![],
37				label: None,
38				block: Block {
39					brace_token: Brace::default(),
40					stmts: vec![Stmt::Expr(expr, None)],
41				},
42			}),
43		});
44
45		State::new(
46			&mut self.known_env,
47			&mut self.known_not_env,
48			self.env.as_ref(),
49		)
50		.expr_closure(&mut closure);
51
52		DissectedClosure {
53			env_variables: self.known_env,
54			closure,
55		}
56	}
57}
58
59/// Split closure into raw closure, accepting its environment as the first argument,
60/// and the environment collection code.
61pub fn split_env(closure: ExprClosure) -> (Expr, ExprClosure, Vec<Ident>) {
62	let mut dissected = ClosureDissector::default().dissect(closure);
63
64	let args = dissected.env_variables.iter().collect::<Vec<_>>();
65
66	let pat: Pat = parse_quote!((#(#args,)*));
67
68	dissected.closure.inputs.insert(0, pat);
69
70	let make_env: Expr = parse_quote! {
71		(#(#args,)*)
72	};
73
74	(
75		make_env,
76		dissected.closure,
77		args.into_iter().cloned().collect(),
78	)
79}
syn_dissect_closure/lib.rs

syn_dissect_closure/
lib.rs
