Skip to content

Commit

Permalink
Feature: Add macro expand!() to expand a template
Browse files Browse the repository at this point in the history
`expand!()` renders a template with arguments multiple times.

### Example:

```rust,ignore
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,ignore
let a: u64 = 1;
let c: Vec<u8> = vec![1, 2];
```
  • Loading branch information
drmingdrmer committed Apr 14, 2024
1 parent 2322676 commit 01bb4da
Show file tree
Hide file tree
Showing 4 changed files with 349 additions and 6 deletions.
216 changes: 216 additions & 0 deletions macros/src/expand.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
66 changes: 66 additions & 0 deletions macros/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#![doc = include_str!("lib_readme.md")]

mod expand;
mod since;
pub(crate) mod utils;

Expand Down Expand Up @@ -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()
}
36 changes: 30 additions & 6 deletions macros/src/lib_readme.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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() {}
```
Expand All @@ -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];
# }
```
37 changes: 37 additions & 0 deletions macros/tests/test_default_types.rs
Original file line number Diff line number Diff line change
@@ -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() {}

0 comments on commit 01bb4da

Please sign in to comment.