Feature: Add macro expand!() to expand a template

drmingdrmer · drmingdrmer · commit b172dc8e261f · 2024-04-14T18:59:05.000+08:00
`expand!()` renders a template with arguments multiple times.

### Example:

```rust,ignore
expand!(KEYED, // ignore duplicate by `K`
        (K, T, V) =&gt; {let K: T = V;},
        (a, u64, 1),
        (a, u32, 2), // duplicate `a` will be ignored
        (c, Vec&lt;u8&gt;, vec![1,2])
);
```

The above code will be transformed into:

```rust,ignore
let a: u64 = 1;
let c: Vec&lt;u8&gt; = vec![1, 2];
```
diff --git a/macros/src/expand.rs b/macros/src/expand.rs
@@ -0,0 +1,216 @@
+use std::collections::HashSet;
+
+use proc_macro2::Ident;
+use quote::quote;
+use quote::ToTokens;
+use syn::parenthesized;
+use syn::parse::Parse;
+use syn::parse::ParseStream;
+use syn::Attribute;
+use syn::Expr;
+use syn::ExprTuple;
+use syn::Token;
+use syn::Type;
+use syn::__private::TokenStream2;
+
+/// A type or an expression which is used as an argument in the `expand` macro.
+#[derive(Debug, Clone, Hash, Eq, PartialEq)]
+enum TypeOrExpr {
+    Attribute(Vec<Attribute>),
+    Type(Type),
+    Expr(Expr),
+    Empty,
+}
+
+impl ToTokens for TypeOrExpr {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        match self {
+            TypeOrExpr::Attribute(attrs) => {
+                for a in attrs {
+                    a.to_tokens(tokens)
+                }
+            }
+            TypeOrExpr::Type(t) => t.to_tokens(tokens),
+            TypeOrExpr::Expr(e) => e.to_tokens(tokens),
+            TypeOrExpr::Empty => {}
+        }
+    }
+}
+
+impl Parse for TypeOrExpr {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let res = input.call(Attribute::parse_outer);
+        if let Ok(r) = res {
+            if !r.is_empty() {
+                return Ok(Self::Attribute(r));
+            }
+        }
+
+        let res = input.parse::<Type>();
+        if let Ok(t) = res {
+            return Ok(Self::Type(t));
+        }
+
+        let res = input.parse::<Expr>();
+        if let Ok(e) = res {
+            return Ok(Self::Expr(e));
+        }
+
+        let l = input.lookahead1();
+        if l.peek(Token![,]) {
+            Ok(Self::Empty)
+        } else {
+            Err(l.error())
+        }
+    }
+}
+
+pub(crate) struct Expand {
+    /// Whether to deduplicate by the first argument as key.
+    pub(crate) keyed: bool,
+
+    /// The template variables
+    pub(crate) idents: Vec<String>,
+
+    /// Template in tokens
+    pub(crate) template: TokenStream2,
+
+    /// Multiple arguments lists for rendering the template
+    args_list: Vec<Vec<TypeOrExpr>>,
+
+    /// The keys that have been present in one of the `args_list`.
+    /// It is used for deduplication, if `keyed` is true.
+    present_keys: HashSet<TypeOrExpr>,
+}
+
+impl Expand {
+    pub(crate) fn render(&self) -> TokenStream2 {
+        let mut output_tokens = TokenStream2::new();
+
+        for values in self.args_list.iter() {
+            for t in self.template.clone().into_iter() {
+                if let proc_macro2::TokenTree::Ident(ident) = t {
+                    let ident_str = ident.to_string();
+
+                    let ident_index = self.idents.iter().position(|x| x == &ident_str);
+                    if let Some(ident_index) = ident_index {
+                        let replacement = &values[ident_index];
+                        output_tokens.extend(replacement.to_token_stream());
+                    } else {
+                        output_tokens.extend(ident.into_token_stream());
+                    }
+                } else {
+                    output_tokens.extend(t.into_token_stream());
+                }
+            }
+        }
+
+        quote! {
+            #output_tokens
+        }
+    }
+}
+
+impl Parse for Expand {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let mut b = Expand {
+            keyed: true,
+            idents: vec![],
+            template: Default::default(),
+            args_list: vec![],
+            present_keys: Default::default(),
+        };
+
+        // KEYED, or !KEYED,
+        {
+            let not = input.parse::<Token![!]>();
+            let not_keyed = not.is_ok();
+
+            let keyed_lit = input.parse::<Ident>()?;
+            if keyed_lit != "KEYED" {
+                return Err(syn::Error::new_spanned(&keyed_lit, "Expected KEYED"));
+            };
+            b.keyed = !not_keyed;
+        }
+
+        input.parse::<Token![,]>()?;
+
+        // Template variables:
+        // (K, V...)
+
+        let idents_tuple = input.parse::<ExprTuple>()?;
+
+        for expr in idents_tuple.elems.iter() {
+            let Expr::Path(p) = expr else {
+                return Err(syn::Error::new_spanned(expr, "Expected path"));
+            };
+
+            let segment = p.path.segments.first().ok_or_else(|| syn::Error::new_spanned(p, "Expected ident"))?;
+            let ident = segment.ident.to_string();
+
+            b.idents.push(ident);
+        }
+
+        // Template body
+        // => { ... }
+        {
+            input.parse::<Token![=>]>()?;
+
+            let brace_group = input.parse::<proc_macro2::TokenTree>()?;
+            let proc_macro2::TokenTree::Group(tree) = brace_group else {
+                return Err(syn::Error::new_spanned(brace_group, "Expected { ... }"));
+            };
+            b.template = tree.stream();
+        }
+
+        // List of arguments tuples for rendering the template
+        // , (K1, V1...)...
+
+        loop {
+            if input.is_empty() {
+                break;
+            }
+
+            input.parse::<Token![,]>()?;
+
+            if input.is_empty() {
+                break;
+            }
+
+            // A tuple of arguments for each rendering
+            // (K1, V1, V2...)
+            {
+                let content;
+                let _parenthesis = parenthesized!(content in input);
+
+                let k = content.parse::<TypeOrExpr>()?;
+                let mut args = vec![k.clone()];
+
+                loop {
+                    if content.is_empty() {
+                        break;
+                    }
+
+                    content.parse::<Token![,]>()?;
+
+                    if content.is_empty() {
+                        break;
+                    }
+
+                    let v = content.parse::<TypeOrExpr>()?;
+                    args.push(v);
+                }
+
+                // Ignore duplicates if keyed
+                if b.present_keys.contains(&k) && b.keyed {
+                    continue;
+                }
+
+                b.present_keys.insert(k);
+                b.args_list.push(args);
+            }
+        }
+
+        Ok(b)
+    }
+}
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
@@ -1,5 +1,6 @@
 #![doc = include_str!("lib_readme.md")]
 
+mod expand;
 mod since;
 pub(crate) mod utils;
 
@@ -157,3 +158,68 @@ fn do_since(args: TokenStream, item: TokenStream) -> Result<TokenStream, syn::Er
     let tokens = since.append_since_doc(item)?;
     Ok(tokens)
 }
+
+/// Render a template with arguments multiple times.
+///
+/// The template to expand is defined as `(K,V) => { ... }`, where `K` and `V` are tempalte
+/// variables.
+///
+/// - The template must contain at least 1 variable.
+/// - If the first macro argument is `KEYED`, the first variable serve as the key for deduplication.
+///   Otherwise, the first macro argument should be `!KEYED`, and no deduplication will be
+///   performed.
+///
+/// # Example: `KEYED` for deduplication
+///
+/// The following code builds a series of let statements:
+/// ```
+/// # use openraft_macros::expand;
+/// # fn foo () {
+/// expand!(
+///     KEYED,
+///     // Template with variables K and V, and template body, excluding the braces.
+///     (K, T, V) => {let K: T = V;},
+///     // Arguments for rendering the template
+///     (a, u64, 1),
+///     (b, String, "foo".to_string()),
+///     (a, u32, 2), // duplicate a will be ignored
+///     (c, Vec<u8>, vec![1,2])
+/// );
+/// # }
+/// ```
+///
+/// The above code will be transformed into:
+///
+/// ```
+/// # fn foo () {
+/// let a: u64 = 1;
+/// let b: String = "foo".to_string();
+/// let c: Vec<u8> = vec![1, 2];
+/// # }
+/// ```
+///
+/// # Example: `!KEYED` for no deduplication
+///
+/// ```
+/// # use openraft_macros::expand;
+/// # fn foo () {
+/// expand!(!KEYED, (K, T, V) => {let K: T = V;},
+///                 (c, u8, 8),
+///                 (c, u16, 16),
+/// );
+/// # }
+/// ```
+///
+/// The above code will be transformed into:
+///
+/// ```
+/// # fn foo () {
+/// let c: u8 = 8;
+/// let c: u16 = 16;
+/// # }
+/// ```
+#[proc_macro]
+pub fn expand(item: TokenStream) -> TokenStream {
+    let repeat = parse_macro_input!(item as expand::Expand);
+    repeat.render().into()
+}
diff --git a/macros/src/lib_readme.md b/macros/src/lib_readme.md
@@ -2,7 +2,7 @@ Supporting utils for [Openraft](https://crates.io/crates/openraft).
 
 # `add_async_trait`
 
-`#[add_async_trait]` adds `Send` bounds to an async trait.
+[`#[add_async_trait]`](`macro@crate::add_async_trait`) adds `Send` bounds to an async trait.
 
 ## Example
 
@@ -24,14 +24,12 @@ trait MyTrait {
 
 # `since`
 
-`#[since(version = "1.0.0")]` adds a doc line `/// Since: 1.0.0`.
+[`#[since(version = "1.0.0")]`](`macro@crate::since`) adds a doc line `/// Since: 1.0.0`.
 
 ## Example
 
 ```rust,ignore
 /// Foo function
-///
-/// Does something.
 #[since(version = "1.0.0")]
 fn foo() {}
 ```
@@ -41,8 +39,34 @@ The above code will be transformed into:
 ```rust,ignore
 /// Foo function
 ///
-/// Does something.
-///
 /// Since: 1.0.0
 fn foo() {}
 ```
+
+
+# `expand` a template
+
+[`expand!()`](`crate::expand!`) renders a template with arguments multiple times.
+
+# Example:
+
+```rust
+# use openraft_macros::expand;
+# fn foo () {
+expand!(KEYED, // ignore duplicate by `K`
+        (K, T, V) => {let K: T = V;},
+        (a, u64, 1),
+        (a, u32, 2), // duplicate `a` will be ignored
+        (c, Vec<u8>, vec![1,2])
+);
+# }
+```
+
+The above code will be transformed into:
+
+```rust
+# fn foo () {
+let a: u64 = 1;
+let c: Vec<u8> = vec![1, 2];
+# }
+```
diff --git a/macros/tests/test_default_types.rs b/macros/tests/test_default_types.rs
@@ -0,0 +1,37 @@
+use openraft_macros::expand;
+
+#[allow(dead_code)]
+#[allow(unused_variables)]
+fn foo() {
+    expand!(
+        KEYED,
+        // template with variables K and V
+        (K, T, V) => {let K: T = V;},
+        // arguments for rendering the template
+        (a, u64, 1),
+        (b, String, "foo".to_string()),
+        (a, u32, 2), // duplicate `a` will be ignored
+        (c, Vec<u8>, vec![1,2]),
+    );
+
+    let _x = 1;
+
+    expand!(
+        !KEYED,
+        (K, T, V) => {let K: T = V;},
+        (c, u8, 8),
+        (c, u16, 16),
+    );
+
+    expand!(
+        !KEYED,
+        (K, M, T) => {M let K: T;},
+        (c, , u8, ),
+        (c, #[allow(dead_code)] , u16),
+        (c, #[allow(dead_code)] #[allow(dead_code)] , u16),
+    );
+}
+
+// #[parse_it]
+// #[allow(dead_code)]
+// fn bar() {}
