scattered_collect_proc_macro/
lib.rs

1//! Proc-macro shims for `scattered-collect` (`scatter` / `gather` attributes).
2
3use std::iter::FromIterator;
4
5use proc_macro::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
6
7/// Registers one value into a scattered collection.
8#[proc_macro_attribute]
9pub fn scatter(attribute: TokenStream, item: TokenStream) -> TokenStream {
10    generate("scatter", "scattered_collect", attribute, item)
11}
12
13/// Builds the sorted slice type over a scattered collection.
14#[proc_macro_attribute]
15pub fn gather(attribute: TokenStream, item: TokenStream) -> TokenStream {
16    generate("gather", "scattered_collect", attribute, item)
17}
18
19#[allow(unknown_lints, tail_expr_drop_order)]
20fn generate(
21    macro_type: &str,
22    macro_crate: &str,
23    attribute: TokenStream,
24    item: TokenStream,
25) -> TokenStream {
26    let mut inner = TokenStream::new();
27
28    let mut crate_path = None;
29    let mut tokens = attribute.clone().into_iter().peekable();
30
31    while let Some(token) = tokens.next() {
32        if let TokenTree::Ident(ident) = &token {
33            if ident.to_string() == "crate_path" {
34                if let Some(TokenTree::Punct(punct)) = tokens.next() {
35                    if punct.as_char() == '=' {
36                        let mut path = TokenStream::new();
37                        while let Some(token) = tokens.peek() {
38                            match token {
39                                TokenTree::Punct(p) if p.as_char() == ',' => {
40                                    tokens.next();
41                                    break;
42                                }
43                                _ => {
44                                    path.extend(std::iter::once(tokens.next().unwrap()));
45                                }
46                            }
47                        }
48                        crate_path = Some(path);
49                        break;
50                    }
51                }
52            }
53        }
54    }
55
56    if attribute.is_empty() {
57        inner.extend([
58            TokenTree::Punct(Punct::new('#', Spacing::Alone)),
59            TokenTree::Group(Group::new(
60                Delimiter::Bracket,
61                TokenStream::from_iter([TokenTree::Ident(Ident::new(
62                    macro_type,
63                    Span::call_site(),
64                ))]),
65            )),
66        ]);
67    } else {
68        inner.extend([
69            TokenTree::Punct(Punct::new('#', Spacing::Alone)),
70            TokenTree::Group(Group::new(
71                Delimiter::Bracket,
72                TokenStream::from_iter([
73                    TokenTree::Ident(Ident::new(macro_type, Span::call_site())),
74                    TokenTree::Group(Group::new(Delimiter::Parenthesis, attribute)),
75                ]),
76            )),
77        ]);
78    }
79
80    inner.extend(item);
81
82    let mut invoke = crate_path.unwrap_or_else(|| {
83        TokenStream::from_iter([
84            TokenTree::Punct(Punct::new(':', Spacing::Joint)),
85            TokenTree::Punct(Punct::new(':', Spacing::Alone)),
86            TokenTree::Ident(Ident::new(macro_crate, Span::call_site())),
87        ])
88    });
89
90    invoke.extend([
91        TokenTree::Punct(Punct::new(':', Spacing::Joint)),
92        TokenTree::Punct(Punct::new(':', Spacing::Alone)),
93        TokenTree::Ident(Ident::new("__support", Span::call_site())),
94        TokenTree::Punct(Punct::new(':', Spacing::Joint)),
95        TokenTree::Punct(Punct::new(':', Spacing::Alone)),
96        TokenTree::Ident(Ident::new(
97            &format!("{macro_type}_parse"),
98            Span::call_site(),
99        )),
100        TokenTree::Punct(Punct::new('!', Spacing::Alone)),
101        TokenTree::Group(Group::new(Delimiter::Parenthesis, inner)),
102        TokenTree::Punct(Punct::new(';', Spacing::Alone)),
103    ]);
104
105    invoke
106}
107
108/// Concatenate two identifiers.
109#[proc_macro]
110pub fn ident_concat(item: TokenStream) -> TokenStream {
111    let mut item = item.into_iter();
112    let Some(TokenTree::Group(pre_group)) = item.next() else {
113        panic!("pre_group: Expected a group");
114    };
115    let Some(TokenTree::Group(name_group)) = item.next() else {
116        panic!("name_group: Expected a group");
117    };
118    let Some(TokenTree::Group(post_group)) = item.next() else {
119        panic!("post_group: Expected a group");
120    };
121
122    let mut item = name_group.stream().into_iter();
123    let mut name = String::new();
124    while let Some(TokenTree::Ident(ident)) = item.next() {
125        name.push_str(&ident.to_string());
126    }
127
128    let mut output = pre_group.stream();
129    output.extend([TokenTree::Ident(Ident::new(&name, Span::call_site()))]);
130    output.extend(post_group.stream());
131    output
132}
scattered_collect_proc_macro/lib.rs

scattered_collect_proc_macro/
lib.rs
