seq_macro/
lib.rs

1//! [![github]](https://github.com/dtolnay/seq-macro) [![crates-io]](https://crates.io/crates/seq-macro) [![docs-rs]](https://docs.rs/seq-macro)
2//!
3//! [github]: https://img.shields.io/badge/github-8da0cb?style=for-the-badge&labelColor=555555&logo=github
4//! [crates-io]: https://img.shields.io/badge/crates.io-fc8d62?style=for-the-badge&labelColor=555555&logo=rust
5//! [docs-rs]: https://img.shields.io/badge/docs.rs-66c2a5?style=for-the-badge&labelColor=555555&logo=docs.rs
6//!
7//! <br>
8//!
9//! # Imagine for-loops in a macro
10//!
11//! This crate provides a `seq!` macro to repeat a fragment of source code and
12//! substitute into each repetition a sequential numeric counter.
13//!
14//! ```
15//! use seq_macro::seq;
16//!
17//! fn main() {
18//!     let tuple = (1000, 100, 10);
19//!     let mut sum = 0;
20//!
21//!     // Expands to:
22//!     //
23//!     //     sum += tuple.0;
24//!     //     sum += tuple.1;
25//!     //     sum += tuple.2;
26//!     //
27//!     // This cannot be written using an ordinary for-loop because elements of
28//!     // a tuple can only be accessed by their integer literal index, not by a
29//!     // variable.
30//!     seq!(N in 0..=2 {
31//!         sum += tuple.N;
32//!     });
33//!
34//!     assert_eq!(sum, 1110);
35//! }
36//! ```
37//!
38//! - If the input tokens contain a section surrounded by `#(` ... `)*` then
39//!   only that part is repeated.
40//!
41//! - The numeric counter can be pasted onto the end of some prefix to form
42//!   sequential identifiers.
43//!
44//! ```
45//! use seq_macro::seq;
46//!
47//! seq!(N in 64..=127 {
48//!     #[derive(Debug)]
49//!     enum Demo {
50//!         // Expands to Variant64, Variant65, ...
51//!         ##(
52//!             Variant~N,
53//!         )*
54//!     }
55//! });
56//!
57//! fn main() {
58//!     assert_eq!("Variant99", format!("{:?}", Demo::Variant99));
59//! }
60//! ```
61//!
62//! - Byte and character ranges are supported: `b'a'..=b'z'`, `'a'..='z'`.
63//!
64//! - If the range bounds are written in binary, octal, hex, or with zero
65//!   padding, those features are preserved in any generated tokens.
66//!
67//! ```
68//! use seq_macro::seq;
69//!
70//! seq!(P in 0x000..=0x00F {
71//!     // expands to structs Pin000, ..., Pin009, Pin00A, ..., Pin00F
72//!     struct Pin~P;
73//! });
74//! ```
75
76#![doc(html_root_url = "https://docs.rs/seq-macro/0.3.5")]
77#![allow(
78    clippy::cast_lossless,
79    clippy::cast_possible_truncation,
80    clippy::derive_partial_eq_without_eq,
81    clippy::let_underscore_untyped,
82    clippy::needless_doctest_main,
83    clippy::single_match_else,
84    clippy::wildcard_imports
85)]
86
87mod parse;
88
89use crate::parse::*;
90use proc_macro::{Delimiter, Group, Ident, Literal, Span, TokenStream, TokenTree};
91use std::char;
92use std::iter::{self, FromIterator};
93
94#[proc_macro]
95pub fn seq(input: TokenStream) -> TokenStream {
96    match seq_impl(input) {
97        Ok(expanded) => expanded,
98        Err(error) => error.into_compile_error(),
99    }
100}
101
102struct Range {
103    begin: u64,
104    end: u64,
105    inclusive: bool,
106    kind: Kind,
107    suffix: String,
108    width: usize,
109    radix: Radix,
110}
111
112struct Value {
113    int: u64,
114    kind: Kind,
115    suffix: String,
116    width: usize,
117    radix: Radix,
118    span: Span,
119}
120
121struct Splice<'a> {
122    int: u64,
123    kind: Kind,
124    suffix: &'a str,
125    width: usize,
126    radix: Radix,
127}
128
129#[derive(Copy, Clone, PartialEq)]
130enum Kind {
131    Int,
132    Byte,
133    Char,
134}
135
136#[derive(Copy, Clone, PartialEq)]
137enum Radix {
138    Binary,
139    Octal,
140    Decimal,
141    LowerHex,
142    UpperHex,
143}
144
145impl<'a> IntoIterator for &'a Range {
146    type Item = Splice<'a>;
147    type IntoIter = Box<dyn Iterator<Item = Splice<'a>> + 'a>;
148
149    fn into_iter(self) -> Self::IntoIter {
150        let splice = move |int| Splice {
151            int,
152            kind: self.kind,
153            suffix: &self.suffix,
154            width: self.width,
155            radix: self.radix,
156        };
157        match self.kind {
158            Kind::Int | Kind::Byte => {
159                if self.inclusive {
160                    Box::new((self.begin..=self.end).map(splice))
161                } else {
162                    Box::new((self.begin..self.end).map(splice))
163                }
164            }
165            Kind::Char => {
166                let begin = char::from_u32(self.begin as u32).unwrap();
167                let end = char::from_u32(self.end as u32).unwrap();
168                let int = |ch| u64::from(u32::from(ch));
169                if self.inclusive {
170                    Box::new((begin..=end).map(int).map(splice))
171                } else {
172                    Box::new((begin..end).map(int).map(splice))
173                }
174            }
175        }
176    }
177}
178
179fn seq_impl(input: TokenStream) -> Result<TokenStream, SyntaxError> {
180    let mut iter = input.into_iter();
181    let var = require_ident(&mut iter)?;
182    require_keyword(&mut iter, "in")?;
183    let begin = require_value(&mut iter)?;
184    require_punct(&mut iter, '.')?;
185    require_punct(&mut iter, '.')?;
186    let inclusive = require_if_punct(&mut iter, '=')?;
187    let end = require_value(&mut iter)?;
188    let body = require_braces(&mut iter)?;
189    require_end(&mut iter)?;
190
191    let range = validate_range(begin, end, inclusive)?;
192
193    let mut found_repetition = false;
194    let expanded = expand_repetitions(&var, &range, body.clone(), &mut found_repetition);
195    if found_repetition {
196        Ok(expanded)
197    } else {
198        // If no `#(...)*`, repeat the entire body.
199        Ok(repeat(&var, &range, &body))
200    }
201}
202
203fn repeat(var: &Ident, range: &Range, body: &TokenStream) -> TokenStream {
204    let mut repeated = TokenStream::new();
205    for value in range {
206        repeated.extend(substitute_value(var, &value, body.clone()));
207    }
208    repeated
209}
210
211fn substitute_value(var: &Ident, splice: &Splice, body: TokenStream) -> TokenStream {
212    let mut tokens = Vec::from_iter(body);
213
214    let mut i = 0;
215    while i < tokens.len() {
216        // Substitute our variable by itself, e.g. `N`.
217        let replace = match &tokens[i] {
218            TokenTree::Ident(ident) => ident.to_string() == var.to_string(),
219            _ => false,
220        };
221        if replace {
222            let original_span = tokens[i].span();
223            let mut literal = splice.literal();
224            literal.set_span(original_span);
225            tokens[i] = TokenTree::Literal(literal);
226            i += 1;
227            continue;
228        }
229
230        // Substitute our variable concatenated onto some prefix, `Prefix~N`.
231        if i + 3 <= tokens.len() {
232            let prefix = match &tokens[i..i + 3] {
233                [first, TokenTree::Punct(tilde), TokenTree::Ident(ident)]
234                    if tilde.as_char() == '~' && ident.to_string() == var.to_string() =>
235                {
236                    match first {
237                        TokenTree::Ident(ident) => Some(ident.clone()),
238                        TokenTree::Group(group) => {
239                            let mut iter = group.stream().into_iter().fuse();
240                            match (iter.next(), iter.next()) {
241                                (Some(TokenTree::Ident(ident)), None) => Some(ident),
242                                _ => None,
243                            }
244                        }
245                        _ => None,
246                    }
247                }
248                _ => None,
249            };
250            if let Some(prefix) = prefix {
251                let number = match splice.kind {
252                    Kind::Int => match splice.radix {
253                        Radix::Binary => format!("{0:01$b}", splice.int, splice.width),
254                        Radix::Octal => format!("{0:01$o}", splice.int, splice.width),
255                        Radix::Decimal => format!("{0:01$}", splice.int, splice.width),
256                        Radix::LowerHex => format!("{0:01$x}", splice.int, splice.width),
257                        Radix::UpperHex => format!("{0:01$X}", splice.int, splice.width),
258                    },
259                    Kind::Byte | Kind::Char => {
260                        char::from_u32(splice.int as u32).unwrap().to_string()
261                    }
262                };
263                let concat = format!("{}{}", prefix, number);
264                let ident = Ident::new(&concat, prefix.span());
265                tokens.splice(i..i + 3, iter::once(TokenTree::Ident(ident)));
266                i += 1;
267                continue;
268            }
269        }
270
271        // Recursively substitute content nested in a group.
272        if let TokenTree::Group(group) = &mut tokens[i] {
273            let original_span = group.span();
274            let content = substitute_value(var, splice, group.stream());
275            *group = Group::new(group.delimiter(), content);
276            group.set_span(original_span);
277        }
278
279        i += 1;
280    }
281
282    TokenStream::from_iter(tokens)
283}
284
285fn enter_repetition(tokens: &[TokenTree]) -> Option<TokenStream> {
286    assert!(tokens.len() == 3);
287    match &tokens[0] {
288        TokenTree::Punct(punct) if punct.as_char() == '#' => {}
289        _ => return None,
290    }
291    match &tokens[2] {
292        TokenTree::Punct(punct) if punct.as_char() == '*' => {}
293        _ => return None,
294    }
295    match &tokens[1] {
296        TokenTree::Group(group) if group.delimiter() == Delimiter::Parenthesis => {
297            Some(group.stream())
298        }
299        _ => None,
300    }
301}
302
303fn expand_repetitions(
304    var: &Ident,
305    range: &Range,
306    body: TokenStream,
307    found_repetition: &mut bool,
308) -> TokenStream {
309    let mut tokens = Vec::from_iter(body);
310
311    // Look for `#(...)*`.
312    let mut i = 0;
313    while i < tokens.len() {
314        if let TokenTree::Group(group) = &mut tokens[i] {
315            let content = expand_repetitions(var, range, group.stream(), found_repetition);
316            let original_span = group.span();
317            *group = Group::new(group.delimiter(), content);
318            group.set_span(original_span);
319            i += 1;
320            continue;
321        }
322        if i + 3 > tokens.len() {
323            i += 1;
324            continue;
325        }
326        let template = match enter_repetition(&tokens[i..i + 3]) {
327            Some(template) => template,
328            None => {
329                i += 1;
330                continue;
331            }
332        };
333        *found_repetition = true;
334        let mut repeated = Vec::new();
335        for value in range {
336            repeated.extend(substitute_value(var, &value, template.clone()));
337        }
338        let repeated_len = repeated.len();
339        tokens.splice(i..i + 3, repeated);
340        i += repeated_len;
341    }
342
343    TokenStream::from_iter(tokens)
344}
345
346impl Splice<'_> {
347    fn literal(&self) -> Literal {
348        match self.kind {
349            Kind::Int | Kind::Byte => {
350                let repr = match self.radix {
351                    Radix::Binary => format!("0b{0:02$b}{1}", self.int, self.suffix, self.width),
352                    Radix::Octal => format!("0o{0:02$o}{1}", self.int, self.suffix, self.width),
353                    Radix::Decimal => format!("{0:02$}{1}", self.int, self.suffix, self.width),
354                    Radix::LowerHex => format!("0x{0:02$x}{1}", self.int, self.suffix, self.width),
355                    Radix::UpperHex => format!("0x{0:02$X}{1}", self.int, self.suffix, self.width),
356                };
357                let tokens = repr.parse::<TokenStream>().unwrap();
358                let mut iter = tokens.into_iter();
359                let literal = match iter.next() {
360                    Some(TokenTree::Literal(literal)) => literal,
361                    _ => unreachable!(),
362                };
363                assert!(iter.next().is_none());
364                literal
365            }
366            Kind::Char => {
367                let ch = char::from_u32(self.int as u32).unwrap();
368                Literal::character(ch)
369            }
370        }
371    }
372}