compose_idents/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#![doc = include_str!("../README.md")]

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use std::collections::HashMap;
use syn::{
    bracketed,
    parse::{Parse, ParseStream},
    parse_macro_input,
    visit_mut::VisitMut,
    Block, Ident, LitStr, Token,
};

struct IdentSpecItem {
    alias: Ident,
    parts: Vec<String>,
}

impl IdentSpecItem {
    fn replacement(&self) -> Ident {
        let ident = self
            .parts
            .iter()
            .fold("".to_string(), |acc, item| format!("{}{}", acc, item));
        format_ident!("{}", ident)
    }
}

struct IdentSpec {
    items: Vec<IdentSpecItem>,
}

impl Parse for IdentSpec {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut items = Vec::new();
        loop {
            let alias: Ident = input.parse()?;
            input.parse::<Token![=]>()?;
            let content;
            bracketed!(content in input);
            let mut parts = Vec::new();
            loop {
                if let Ok(part) = content.parse::<Ident>() {
                    parts.push(part.to_string());
                } else if content.parse::<Token![_]>().is_ok() {
                    parts.push("_".to_string());
                } else if let Ok(part) = content.parse::<LitStr>() {
                    parts.push(part.value());
                } else {
                    return Err(content.error("Expected identifier or _"));
                }
                if content.is_empty() {
                    break;
                }
                content.parse::<Token![,]>()?;
            }
            items.push(IdentSpecItem { alias, parts });
            input.parse::<Token![;]>()?;
            if !input.peek(Ident) {
                break;
            }
        }
        Ok(IdentSpec { items })
    }
}

impl IdentSpec {
    fn replacements(&self) -> HashMap<Ident, Ident> {
        self.items
            .iter()
            .map(|item| (item.alias.clone(), item.replacement()))
            .collect()
    }
}

struct ComposeIdentsArgs {
    spec: IdentSpec,
    block: Block,
}

impl Parse for ComposeIdentsArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let spec: IdentSpec = input.parse()?;
        let block: Block = input.parse()?;
        Ok(ComposeIdentsArgs { spec, block })
    }
}

struct ComposeIdentsVisitor {
    replacements: HashMap<Ident, Ident>,
}

impl VisitMut for ComposeIdentsVisitor {
    fn visit_ident_mut(&mut self, ident: &mut Ident) {
        if let Some(replacement) = self.replacements.get(ident) {
            *ident = replacement.clone();
        }
    }
}

/// Compose identifiers from the provided parts and replace their aliases in the code block.
///
/// # Example
///
/// ```rust
/// use compose_idents::compose_idents;
///
/// compose_idents!(my_fn_1 = [foo, _, "baz"]; my_fn_2 = [spam, _, eggs]; {
///     fn my_fn_1() -> u32 {
///         111
///     }
///
///     fn my_fn_2() -> u32 {
///         999
///     }
/// });
///
/// assert_eq!(foo_baz(), 111);
/// assert_eq!(spam_eggs(), 999);
/// ```
#[proc_macro]
pub fn compose_idents(input: TokenStream) -> TokenStream {
    // Parse the input into our `Assignments` structure
    let args = parse_macro_input!(input as ComposeIdentsArgs);
    let mut visitor = ComposeIdentsVisitor {
        replacements: args.spec.replacements(),
    };
    let mut block = args.block;
    visitor.visit_block_mut(&mut block);

    let block_content = block.stmts;

    let expanded = quote! { #(#block_content)* };
    TokenStream::from(expanded)
}