rust: macros: add seq! macro

Signed-off-by: Fabien Parent <[email protected]>

Author: Fabo

rust/macros/lib.rs

@@ -10,6 +10,7 @@ mod module;
 mod paste;
 mod pin_data;
 mod pinned_drop;
+mod seq;
 mod vtable;
 mod zeroable;
 
@@ -387,6 +388,25 @@ pub fn paste(input: TokenStream) -> TokenStream {
     tokens.into_iter().collect()
 }
 
+/// Repeat a fragment of code and provide a numerical index for the current repetition
+///
+/// # Examples
+///
+/// ```rust,ignore
+/// seq!(i in 0..10) {
+///     func$i() {
+///     }
+/// }
+///
+/// seq!(i in 8..=15) {
+///     bit$i() {
+///     }
+/// }
+#[proc_macro]
+pub fn seq(input: TokenStream) -> TokenStream {
+    seq::expand(input)
+}
+
 /// Derives the [`Zeroable`] trait for the given struct.
 ///
 /// This can only be used for structs where every field implements the [`Zeroable`] trait.

rust/macros/seq.rs

@@ -0,0 +1,132 @@
+use proc_macro::{
+    Delimiter, TokenStream,
+    TokenTree::{self, Group, Ident, Literal},
+};
+use std::ops::Range;
+
+fn process_group(var: &str, i: isize, tokens: impl Iterator<Item = TokenTree>) -> Vec<TokenTree> {
+    let mut tt = Vec::<TokenTree>::new();
+
+    let mut tokens = tokens.peekable();
+
+    while let Some(token) = tokens.next() {
+        match token {
+            Group(ref group) => {
+                let group_tokens = process_group(var, i, group.stream().into_iter());
+                let stream = FromIterator::from_iter(group_tokens.into_iter());
+                let new_group = proc_macro::Group::new(group.delimiter(), stream);
+                tt.push(TokenTree::Group(new_group));
+            }
+            TokenTree::Punct(ref punct) => {
+                if punct.to_string() == "$" {
+                    if let Some(TokenTree::Ident(ident)) = tokens.peek() {
+                        if ident.to_string() == var {
+                            tt.push(TokenTree::Literal(proc_macro::Literal::isize_unsuffixed(i)));
+                            tokens.next();
+                            continue;
+                        }
+                    }
+                }
+
+                tt.push(token);
+            }
+            _ => tt.push(token),
+        }
+    }
+
+    tt
+}
+
+pub(crate) fn expand(input: TokenStream) -> TokenStream {
+    let mut tokens = input.into_iter().peekable();
+
+    let var = if let Some(Ident(i)) = tokens.next() {
+        i.to_string()
+    } else {
+        panic!("seq! first token should be an identifier");
+    };
+
+    let token = tokens.next().expect("missing token, expecting '.'");
+    assert!(matches!(token, TokenTree::Ident(x) if x.to_string() == "in"));
+
+    let token = tokens
+        .next()
+        .expect("seq!: missing token, expecting integer");
+    let token = if let Group(group) = token {
+        group
+            .stream()
+            .into_iter()
+            .next()
+            .expect("seq: missing token, expecting integer")
+    } else {
+        token
+    };
+
+    let start = if let Literal(lit) = token {
+        lit.to_string()
+            .parse::<isize>()
+            .expect("Failed to convert literal to isize")
+    } else {
+        panic!("seq!: unexpected token '{token}'");
+    };
+
+    let token = tokens.next().expect("seq!: missing token, expecting '.'");
+    assert!(matches!(token, TokenTree::Punct(x) if x == '.'));
+    let token = tokens.next().expect("seq!: missing token, expecting '.'");
+    assert!(matches!(token, TokenTree::Punct(x) if x == '.'));
+
+    let is_inclusive_range = if let Some(TokenTree::Punct(p)) = tokens.peek() {
+        if p.as_char() == '=' {
+            tokens.next();
+            true
+        } else {
+            false
+        }
+    } else {
+        false
+    };
+
+    let token = tokens
+        .next()
+        .expect("seq!: missing token, expecting integer");
+    let token = if let Group(group) = token {
+        group
+            .stream()
+            .into_iter()
+            .next()
+            .expect("seq: missing token, expecting integer")
+    } else {
+        token
+    };
+
+    let end = if let Literal(lit) = token {
+        lit.to_string()
+            .parse::<isize>()
+            .expect("Failed to convert literal to isize")
+    } else {
+        panic!("seq!: unexpected token '{token}'");
+    };
+    let range = Range {
+        start,
+        end: end + if is_inclusive_range { 1 } else { 0 },
+    };
+
+    let tokens = if let Some(Group(group)) = tokens.next() {
+        if group.delimiter() != Delimiter::Brace {
+            panic!("seq! expected brace");
+        }
+
+        group.stream().into_iter()
+    } else {
+        panic!("seq! missing opening brace");
+    };
+
+    let tokens: Vec<TokenTree> = tokens.collect();
+    let mut tt = Vec::<TokenTree>::new();
+
+    for i in range {
+        tt.extend_from_slice(&process_group(&var, i, tokens.clone().into_iter()));
+    }
+
+    FromIterator::from_iter(tt)
+}