stringify_inner/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use core::{convert::identity, str::FromStr as _};
4
5use proc_macro::{
6    Delimiter, Group, Ident, Literal, Punct, Spacing::*, Span, TokenStream,
7    TokenTree,
8};
9
10#[must_use]
11fn stream<I>(iter: I) -> TokenStream
12where I: IntoIterator,
13      TokenStream: FromIterator<I::Item>,
14{
15    TokenStream::from_iter(iter)
16}
17
18fn err<T>(msg: &str, span: Span) -> Result<T, TokenStream> {
19    let s = |mut t: TokenTree| {
20        t.set_span(span);
21        t
22    };
23    Err(stream([
24        s(Punct::new(':', Joint).into()),
25        s(Punct::new(':', Joint).into()),
26        s(Ident::new("core", span).into()),
27        s(Punct::new(':', Joint).into()),
28        s(Punct::new(':', Joint).into()),
29        s(Ident::new("compile_error", span).into()),
30        s(Punct::new('!', Joint).into()),
31        s(Group::new(Delimiter::Brace, stream([
32            s(Literal::string(msg).into()),
33        ])).into()),
34    ]))
35}
36
37fn extract_str(s: &str, span: Span) -> Result<String, TokenStream> {
38    if !s.ends_with(['"', '#']) {
39        return err("invalid string suffix", span);
40    }
41    if s.starts_with('"') {
42        return Ok(s[1..s.len()-1].to_owned());
43    }
44    if !s.starts_with('r') {
45        return err("invalid string literal", span);
46    }
47    let mut s = &s[1..];
48    while s.starts_with('#') {
49        s = &s[1..s.len()-1];
50    }
51    let escaped = format!("{:?}", &s[1..s.len()-1]);
52    debug_assert!(escaped.starts_with('"') && escaped.ends_with('"'), "{escaped}");
53    Ok(escaped[1..escaped.len()-1].to_string())
54}
55
56fn merge_str(a: &Literal, b: &Literal) -> Result<Literal, TokenStream> {
57    let (sa, sb) = (a.span(), b.span());
58    let (a, b) = (a.to_string(), b.to_string());
59    let (a, b) = (extract_str(&a, sa)?, extract_str(&b, sb)?);
60    let mut lit = Literal::from_str(&format!("\"{a}{b}\"")).unwrap();
61    lit.set_span(sa);
62    Ok(lit)
63}
64
65fn do_operation(
66    tok: TokenTree,
67    iter: &mut impl Iterator<Item = TokenTree>,
68) -> Result<TokenStream, TokenStream> {
69    let Some(op) = iter.next() else {
70        return err("unexpected end of input", tok.span());
71    };
72    Ok(match op {
73        TokenTree::Punct(ref punct) => {
74            if punct.as_char() == '#' {
75                stream([op])
76            } else {
77                return err("invalid operator", punct.span());
78            }
79        },
80        TokenTree::Group(_) => {
81            stream([
82                stream([tok]),
83                expr_impl(stream([op]))?,
84            ])
85        },
86        TokenTree::Literal(tt) => return err("invalid operator", tt.span()),
87        TokenTree::Ident(ident) => {
88            let Some(param) = iter.next() else {
89                return err("unexpected end of input", ident.span());
90            };
91            let TokenTree::Group(param) = param else {
92                return err("invalid operation param", param.span());
93            };
94            let gspan = param.span();
95
96            match &*ident.to_string() {
97                "stringify" => {
98                    let out_span = param.stream().into_iter().next()
99                        .map_or(param.span(), |t| t.span());
100
101                    let s = param.stream().to_string();
102                    let mut tt = Literal::string(&s);
103                    tt.set_span(out_span);
104                    stream([TokenTree::from(tt)])
105                },
106                "concat" => {
107                    let param = expr_impl(param.stream())?;
108                    let mut s = Literal::string("");
109                    let mut span = None;
110                    let mut iter = param.into_iter().peekable();
111
112                    while let Some(tt) = iter.next() {
113                        iter.next_if(|p| matches!(p,
114                                TokenTree::Punct(p) if p.as_char() == ','));
115                        let TokenTree::Literal(lit) = tt else {
116                            return err("is not a literal", tt.span());
117                        };
118                        span.get_or_insert(lit.span());
119                        s = merge_str(&s, &lit)?;
120                    }
121
122                    s.set_span(span.unwrap_or(gspan));
123                    stream([TokenTree::from(s)])
124                },
125                _ => return err("unknown operator", ident.span()),
126            }
127        },
128    })
129}
130
131fn expr_impl(stream: TokenStream) -> Result<TokenStream, TokenStream> {
132    let mut result = TokenStream::new();
133    let mut iter = stream.into_iter();
134
135    while let Some(tok) = iter.next() {
136        match tok {
137            proc_macro::TokenTree::Group(group) => {
138                let mut new_group = Group::new(
139                    group.delimiter(),
140                    expr_impl(group.stream())?,
141                );
142                new_group.set_span(group.span());
143                result.extend([TokenTree::from(new_group)]);
144            },
145            proc_macro::TokenTree::Punct(ref punct)
146                if punct.as_char() == '#' =>
147            {
148                result.extend(do_operation(tok, &mut iter)?);
149            },
150            _ => result.extend([tok]),
151        }
152    }
153
154    Ok(result)
155}
156
157/// Run string expressions on expressions
158///
159/// - `#stringify(...)`: like `stringify!(...)`
160/// - `#concat(...)`: like `concat!(...)`
161/// - `##`: like `#`
162/// - `#[...]`: like `#[...]`
163///
164/// # Examples
165/// ```
166/// use stringify_inner::sexpr;
167///
168/// assert_eq!(sexpr!(#stringify(foo)), "foo");
169/// assert_eq!(sexpr!(&#stringify(foo)[1..]), "oo");
170/// assert_eq!(sexpr!(#concat(#stringify(foo), "bar")), "foobar");
171/// ```
172#[proc_macro]
173pub fn sexpr(stream: TokenStream) -> TokenStream {
174    expr_impl(stream).map_or_else(identity, identity)
175}
176
177/// Run string expressions on attribute
178///
179/// - `#stringify(...)`: like `stringify!(...)`
180/// - `#concat(...)`: like `concat!(...)`
181/// - `##`: like `#`
182/// - `#[...]`: like `#[...]`
183///
184/// # Examples
185/// ```
186/// use stringify_inner::sexpr_attr;
187///
188/// #[sexpr_attr(doc(alias = #stringify(bar)))]
189/// fn foo() {}
190/// ```
191#[proc_macro_attribute]
192pub fn sexpr_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
193    stream([
194        expr_impl(stream([
195            TokenTree::from(Punct::new('#', Joint)),
196            Group::new(Delimiter::Bracket, attr).into(),
197        ])).map_or_else(identity, identity),
198        item,
199    ])
200}