closure_it/
lib.rs

1#![doc = include_str!("../README.md")]
2#![allow(clippy::needless_doctest_main)]
3
4use proc_macro::{
5    Delimiter, Group, Ident, Literal, Punct, Spacing::*, Span, TokenStream,
6    TokenTree,
7};
8
9fn span_setter(span: Span) -> impl Fn(TokenTree) -> TokenTree {
10    move |mut tt| {
11        tt.set_span(span);
12        tt
13    }
14}
15
16#[must_use]
17fn stream<I>(iter: I) -> TokenStream
18where I: IntoIterator<Item = TokenTree>,
19{
20    TokenStream::from_iter(iter)
21}
22
23fn err(msg: &str, span: Span) -> TokenStream {
24    let s = span_setter(span);
25    stream([
26        s(Punct::new(':', Joint).into()),
27        s(Punct::new(':', Joint).into()),
28        s(Ident::new("core", span).into()),
29        s(Punct::new(':', Joint).into()),
30        s(Punct::new(':', Joint).into()),
31        s(Ident::new("compile_error", span).into()),
32        s(Punct::new('!', Joint).into()),
33        s(Group::new(Delimiter::Brace, stream([
34            s(Literal::string(msg).into()),
35        ])).into()),
36    ])
37}
38
39#[derive(Default, Clone)]
40struct Closure<'a> {
41    it: Option<TokenTree>,
42    catch_it: &'a str,
43}
44impl Closure<'_> {
45    fn make_closure(&mut self) -> TokenStream {
46        let Some(it) = self.it.take() else {
47            return TokenStream::new();
48        };
49        let s = span_setter(it.span());
50        stream([
51            s(Punct::new('|', Joint).into()),
52            it,
53            s(Punct::new('|', Joint).into()),
54        ])
55    }
56
57    fn ext_proc_it(&mut self, input: TokenStream) -> TokenStream {
58        let ext = &mut Self { catch_it: self.catch_it, ..Default::default() };
59        let proc_it = ext.proc_it(input);
60        ext.make_closure().into_iter().chain(proc_it).collect()
61    }
62
63    fn proc_it(&mut self, input: TokenStream) -> TokenStream {
64        let iter = &mut input.into_iter().peekable();
65        let mut result = TokenStream::new();
66
67        while let Some(tt) = iter.next() {
68            match tt {
69                TokenTree::Group(group)
70                    if group.delimiter() == Delimiter::Parenthesis =>
71                {
72                    let grouped = self.ext_proc_it(group.stream());
73                    result.extend([
74                        Group::new(group.delimiter(), grouped).into(),
75                    ] as [TokenTree; 1]);
76                },
77                TokenTree::Group(group) => {
78                    let grouped = self.proc_it(group.stream());
79                    result.extend([
80                        Group::new(group.delimiter(), grouped).into(),
81                    ] as [TokenTree; 1]);
82                },
83                TokenTree::Ident(ref ident)
84                    if ident.to_string() == self.catch_it =>
85                {
86                    result.extend([self.it.get_or_insert(tt).clone()]);
87                },
88                TokenTree::Ident(_) | TokenTree::Literal(_) => {
89                    result.extend([tt]);
90                },
91                TokenTree::Punct(ref punct)
92                    if matches!(punct.as_char(), ',' | ';') =>
93                {
94                    result.extend([tt]);
95                    result.extend(self.ext_proc_it(iter.collect()));
96                },
97                TokenTree::Punct(ref punct)
98                    if punct.as_char() == '='
99                        && punct.spacing() == Joint
100                        && iter.peek().is_some_and(|p| {
101                            matches!(p, TokenTree::Punct(p)
102                                if p.as_char() == '>')
103                        }) =>
104                {
105                    result.extend([tt, iter.next().unwrap()]);
106                    result.extend(self.ext_proc_it(iter.collect()));
107                },
108                TokenTree::Punct(_) => {
109                    result.extend([tt]);
110                },
111            }
112        }
113
114        result
115    }
116}
117
118fn get_catch_it<F>(attr: TokenStream, f: F) -> TokenStream
119where F: FnOnce(&str) -> TokenStream,
120{
121    let iter = &mut attr.into_iter();
122    let catch_it = match iter.next() {
123        Some(TokenTree::Ident(ident)) => &*ident.to_string(),
124        Some(attr) => return err("invalid input", attr.span()),
125        _ => "it",
126    };
127    if let Some(extra) = iter.next() {
128        return err("invalid input", extra.span());
129    }
130    f(catch_it)
131}
132
133/// Replace `it` to closure body,
134/// expand the closure parameter after `,` `;` `=>` and `(`
135///
136/// # Examples
137/// ```
138/// #[closure_it::closure_it]
139/// fn main() {
140///     assert_eq!([0i32, 1, 2].map(it+2), [2, 3, 4]);
141///     assert_eq!([0i32, -1, 2].map(it.abs()), [0, 1, 2]);
142///     assert_eq!(Some(2).map_or(3, it*2), 4);
143/// }
144/// ```
145///
146/// ```
147/// #[closure_it::closure_it(this)]
148/// fn main() {
149///     assert_eq!([0i32, 1, 2].map(this+2), [2, 3, 4]);
150///     assert_eq!([0i32, -1, 2].map(this.abs()), [0, 1, 2]);
151///     assert_eq!(Some(2).map_or(3, this*2), 4);
152/// }
153/// ```
154#[proc_macro_attribute]
155pub fn closure_it(attr: TokenStream, item: TokenStream) -> TokenStream {
156    get_catch_it(attr, |catch_it| {
157        Closure { catch_it, ..Default::default() }
158            .ext_proc_it(item)
159    })
160}
closure_it/lib.rs

closure_it/
lib.rs
