Skip to main content

sort_macro/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use proc_macro::*;
4use Spacing::*;
5use tt_path::TtPath;
6
7fn stream<I>(iter: I) -> TokenStream
8where I: IntoIterator<Item = TokenTree>,
9{
10    TokenStream::from_iter(iter)
11}
12
13fn span_setter(span: Span) -> impl Fn(TokenTree) -> TokenTree {
14    move |mut tt| {
15        tt.set_span(span);
16        tt
17    }
18}
19
20fn err(msg: &str, span: Span) -> TokenStream {
21    let s = span_setter(span);
22    stream([
23        s(Punct::new(':', Joint).into()),
24        s(Punct::new(':', Joint).into()),
25        s(Ident::new("core", span).into()),
26        s(Punct::new(':', Joint).into()),
27        s(Punct::new(':', Joint).into()),
28        s(Ident::new("compile_error", span).into()),
29        s(Punct::new('!', Joint).into()),
30        s(Group::new(Delimiter::Brace, stream([
31            s(Literal::string(msg).into()),
32        ])).into()),
33    ])
34}
35
36struct Sorter {
37    key: TtPath,
38    prefix: TokenStream,
39    tokens: Vec<TokenTree>,
40    args: Group,
41}
42impl Sorter {
43    fn finish(self) -> TokenStream {
44        let args = self.args.stream().into_iter()
45            .map(|tt| {
46                let key = self.key.apply(tt.clone())
47                    .into_iter()
48                    .map(|it| it.map(|it| it.to_string()))
49                    .collect::<Result<Vec<_>, _>>();
50                match key {
51                    Ok(key) => Ok((tt, key)),
52                    Err(e) => Err(e),
53                }
54            })
55            .collect::<Result<Vec<_>, _>>();
56        let mut args = match args {
57            Ok(args) => args,
58            Err(e) => return err(e.msg(), e.span()),
59        };
60        args.sort_by(|(_, k1), (_, k2)| k1.cmp(k2));
61
62        let sorted_args = self.prefix.into_iter().chain(args.into_iter().map(|it| it.0));
63        let mut sorted_args = Group::new(self.args.delimiter(), TokenStream::from_iter(sorted_args));
64        sorted_args.set_span(self.args.span());
65
66        self.tokens.into_iter()
67            .chain([sorted_args.into()])
68            .collect()
69    }
70}
71
72#[doc = include_str!("../README.md")]
73#[proc_macro]
74pub fn sort_in(tokens: TokenStream) -> TokenStream {
75    let mut tokens = tokens.into_iter().collect::<Vec<_>>();
76    let Some(args) = tokens.pop() else {
77        return err("must input a `()` `{}` `[]`", Span::call_site());
78    };
79    let TokenTree::Group(args) = args else {
80        return err("last token must is group", args.span());
81    };
82
83    let mut prefix = TokenStream::new();
84    let mut key = TtPath::from_tokens(&mut None.into_iter().peekable()).unwrap();
85
86    if let Some(TokenTree::Group(first)) = tokens.first()
87        && first.delimiter() == Delimiter::Bracket
88    {
89        let iter = &mut first.stream().into_iter().peekable();
90        match TtPath::from_tokens(iter) {
91            Ok(tt_path) => key = tt_path,
92            Err(e) => return err(e.msg(), e.span()),
93        }
94        if let Some(peek) = iter.peek() {
95            return err("unexpected key token", peek.span());
96        }
97        tokens.remove(0);
98    }
99
100    if let Some(TokenTree::Group(first)) = tokens.first()
101        && first.delimiter() == Delimiter::Brace
102    {
103        prefix = first.stream();
104        tokens.remove(0);
105    }
106
107    Sorter { key, prefix, tokens, args }.finish()
108}