syntactic_for/
lib.rs

1//! A syntactic for loop.
2//!
3//! For example, the following takes the sum of the bit-length of four integer
4//! types:
5//! ```
6//! # use syntactic_for::syntactic_for;
7//! let sum = syntactic_for!{ ty in [ u8, u16, u32, u64 ] {
8//!     [$( <$ty>::BITS ),*].into_iter().sum::<u32>()
9//! }};
10//! assert_eq!(sum, 120);
11//! ```
12//!
13//! # Usage
14//!
15//! The syntax is as follows:
16//! ```ignore
17//! syntactic_for!{ IDENTIFIER in [ EXPRESSION, EXPRESSION, ... ] {
18//!     BODY
19//! }}
20//! ```
21//! where `BODY` works similarly to `macro_rules!`, that is:
22//! `$($IDENTIFIER)SEPARATOR*` will expand and substitute `IDENTIFIER` with
23//! each `EXPRESSION`, separating the expansions with `SEPARATOR`.
24//!
25//! `SEPARATOR` can be any non-`*` punctuation.  Hence, the example from above
26//! could also be written without an iterator:
27//! ```
28//! # use syntactic_for::syntactic_for;
29//! # let sum =
30//! # syntactic_for!{ ty in [ u8, u16, u32, u64 ] {
31//! $( <$ty>::BITS )+*
32//! # }};
33//! # assert_eq!(sum, 120);
34//! ```
35//!
36//! # Examples
37//!
38//! ## Loop unrolling
39//!
40//! Sum the elements of an array with
41//! [loop unrolling](https://en.wikipedia.org/wiki/Loop_unrolling):
42//! ```
43//! # use syntactic_for::syntactic_for;
44//! let array = b"oh my, I am getting summed!";
45//! let mut acc = 0u32;
46//! let mut i = 0;
47//! while i <= array.len()-4 {
48//!     syntactic_for!{ offset in [ 0, 1, 2, 3 ] {$(
49//!         acc += array[i + $offset] as u32;
50//!     )*}}
51//!     i += 4;
52//! }
53//! for j in i..array.len() {
54//!     acc += array[j] as u32;
55//! }
56//! assert_eq!(acc, 2366);
57//! ```
58//!
59//! ## Matching
60//!
61//! Find the maximum value of an integer type of the given bit size:
62//! ```
63//! # use syntactic_for::syntactic_for;
64//! # let bit_size = 16;
65//! let max_size = syntactic_for!{ ty in [ u8, u16, u32, u64, u128 ] {
66//!     match bit_size {
67//!         $(<$ty>::BITS => <$ty>::MAX as u128,)*
68//!         other => panic!("No integer of size {other}"),
69//!     }
70//! }};
71//! # assert_eq!(max_size, u16::MAX as u128)
72//! ```
73//!
74//! ## `impl` blocks
75//!
76//! Implement a trait for a set of types:
77//! ```
78//! # use syntactic_for::syntactic_for;
79//! # trait MyTrait {}
80//! syntactic_for!{ ty in [ u8, u16, u32, u64, u128 ] {$(
81//!     impl MyTrait for $ty {
82//!         // snip.
83//!     }
84//! )*}}
85//! ```
86//!
87//! ## Custom syntactic loop
88//!
89//! A useful design pattern is to define a custom macro that expands to a
90//! syntactic loop over a given set of expressions:
91//! ```
92//! # struct CustomType1;
93//! # struct CustomType2;
94//! #[doc(hidden)]
95//! pub extern crate syntactic_for;
96//!
97//! #[macro_export]
98//! macro_rules! for_each_custom_type {
99//!     ($ident:ident { $($tt:tt)* }) => {
100//!         $crate::syntactic_for::syntactic_for! { $ident in [
101//!             $crate::CustomType1,
102//!             $crate::CustomType2,
103//!             // etc.
104//!         ] { $($tt)* } }
105//!     }
106//! }
107//! ```
108//!
109//! For example, a library could expose `for_each_custom_type` as a way of
110//! letting its users write syntactic loops over a set of types defined in the
111//! library.  Then, it becomes possible to add types to that loop inside the
112//! library, whithout requiring any change on the user's end:
113//!
114//! ```
115//! # struct CustomType1;
116//! # impl CustomType1 { fn parse(i: &str) -> Result<(), ()> { Err(()) }}
117//! # struct CustomType2;
118//! # impl CustomType2 { fn parse(i: &str) -> Result<(), ()> { Ok(()) }}
119//! # pub extern crate syntactic_for;
120//! # #[macro_export]
121//! # macro_rules! for_each_custom_type {
122//! #     ($ident:ident { $($tt:tt)* }) => {
123//! #         ::syntactic_for::syntactic_for! { $ident in [
124//! #             CustomType1,
125//! #             CustomType2,
126//! #             // etc.
127//! #         ] { $($tt)* } }
128//! #     }
129//! # }
130//! # mod my_library { pub use for_each_custom_type; }
131//! // Try and parse each library type in succession, stopping at the first
132//! // success:
133//! fn can_parse(input: &str) -> bool {
134//!     my_library::for_each_custom_type! { ty {
135//!         $(if let Ok(parsed) = <$ty>::parse(input) {
136//!             return true;
137//!         })*
138//!     }}
139//!     return false;
140//! }
141//! # assert_eq!(can_parse("foo"), true);
142//! ```
143use proc_macro2::{TokenStream, TokenTree};
144use quote::{ToTokens, TokenStreamExt};
145use syn::{
146    braced, bracketed,
147    parse::{Parse, ParseStream},
148    parse_macro_input,
149    punctuated::Punctuated,
150    token, Expr, Ident, Token,
151};
152
153struct SyntacticFor {
154    ident: Ident,
155    _in_token: Token![in],
156    _bracket_token: token::Bracket,
157    exprs: Punctuated<Expr, Token![,]>,
158    _brace_token: token::Brace,
159    body: TokenStream,
160}
161
162impl Parse for SyntacticFor {
163    fn parse(input: ParseStream) -> syn::Result<Self> {
164        let brackets;
165        let braces;
166        Ok(SyntacticFor {
167            ident: input.parse()?,
168            _in_token: input.parse()?,
169            _bracket_token: bracketed!(brackets in input),
170            exprs: Punctuated::parse_terminated(&brackets)?,
171            _brace_token: braced!(braces in input),
172            body: braces.parse()?,
173        })
174    }
175}
176
177fn subs_group<'a, S, IntoIter>(
178    pattern: &Ident,
179    subs: &'a S,
180    tokens: TokenStream,
181) -> syn::Result<TokenStream>
182where
183    &'a S: IntoIterator<IntoIter = IntoIter> + Clone + 'a,
184    IntoIter: ExactSizeIterator,
185    <IntoIter as Iterator>::Item: ToTokens,
186{
187    let mut output = TokenStream::new();
188    let mut tokens = tokens.into_iter().peekable();
189    while let Some(token) = tokens.next() {
190        match token {
191            TokenTree::Punct(punct) if punct.as_char() == '$' => match tokens.next() {
192                Some(TokenTree::Group(group)) => {
193                    let separator = match tokens.peek() {
194                        Some(TokenTree::Punct(punct)) if punct.as_char() == '*' => None,
195                        Some(TokenTree::Punct(_)) => {
196                            if let TokenTree::Punct(punct) = tokens.next().unwrap() {
197                                Some(punct)
198                            } else {
199                                unreachable!()
200                            }
201                        }
202                        Some(token) => {
203                            return Err(syn::Error::new_spanned(
204                                token,
205                                format!("expected punctuation or `*`, found {}", token),
206                            ))
207                        }
208                        None => panic!("unexpected end of stream after group"),
209                    };
210                    match tokens.next() {
211                        Some(TokenTree::Punct(punct)) if punct.as_char() == '*' => {}
212                        Some(token) => {
213                            return Err(syn::Error::new_spanned(
214                                &token,
215                                format!("expected `*`, found {}", token),
216                            ))
217                        }
218                        None => panic!("unexpected end of stream after group"),
219                    }
220
221                    let subs = subs.into_iter();
222                    let len = subs.len();
223                    for (i, sub) in subs.enumerate() {
224                        output.extend(subs_ident(pattern, &sub, group.stream())?);
225                        if i + 1 < len {
226                            if let Some(separator) = &separator {
227                                output.append(separator.clone());
228                            }
229                        }
230                    }
231                }
232                Some(token) => {
233                    return Err(syn::Error::new_spanned(
234                        &token,
235                        format!("expected group after `$`, found `{}`", token),
236                    ))
237                }
238                None => {
239                    panic!("unexpected end of stream after `$`")
240                }
241            },
242            TokenTree::Group(group) => {
243                output.append(proc_macro2::Group::new(
244                    group.delimiter(),
245                    subs_group(pattern, subs, group.stream())?,
246                ));
247            }
248            token => output.append(token),
249        }
250    }
251    Ok(output)
252}
253
254fn subs_ident<'a, I>(pattern: &Ident, sub: &'a I, tokens: TokenStream) -> syn::Result<TokenStream>
255where
256    &'a I: ToTokens,
257{
258    let mut output = TokenStream::new();
259    let mut tokens = tokens.into_iter();
260    while let Some(token) = tokens.next() {
261        match token {
262            TokenTree::Punct(punct) if punct.as_char() == '$' => match tokens.next() {
263                Some(TokenTree::Ident(ident)) if &ident == pattern => {
264                    sub.to_tokens(&mut output);
265                }
266                Some(token) => {
267                    return Err(syn::Error::new_spanned(
268                        &token,
269                        format!("expected `{}` after `$`, found `{}`", pattern, token),
270                    ))
271                }
272                None => {
273                    panic!("unexpected end of stream after `$`")
274                }
275            },
276            TokenTree::Group(group) => {
277                output.append(proc_macro2::Group::new(
278                    group.delimiter(),
279                    subs_ident(pattern, sub, group.stream())?,
280                ));
281            }
282            token => output.append(token),
283        }
284    }
285    Ok(output)
286}
287
288/// Iterate over a list of (syntactic) expressions.
289///
290/// For details, see [top level documentation][crate].
291#[proc_macro]
292pub fn syntactic_for(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
293    let SyntacticFor {
294        ident, exprs, body, ..
295    } = parse_macro_input!(input as SyntacticFor);
296
297    subs_group(&ident, &exprs, body)
298        .unwrap_or_else(syn::Error::into_compile_error)
299        .into()
300}