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
95            match &*ident.to_string() {
96                "stringify" => {
97                    let out_span = param.stream().into_iter().next()
98                        .map_or(param.span(), |t| t.span());
99
100                    let s = param.stream().to_string();
101                    let mut tt = Literal::string(&s);
102                    tt.set_span(out_span);
103                    stream([TokenTree::from(tt)])
104                },
105                "concat" => {
106                    let param = expr_impl(param.stream())?;
107                    let mut s = Literal::string("");
108                    let mut iter = param.into_iter().peekable();
109                    while let Some(tt) = iter.next() {
110                        iter.next_if(|p| matches!(p,
111                                TokenTree::Punct(p) if p.as_char() == ','));
112                        let TokenTree::Literal(lit) = tt else {
113                            return err("is not a literal", tt.span());
114                        };
115                        s = merge_str(&s, &lit)?;
116                    }
117                    stream([TokenTree::from(s)])
118                },
119                _ => return err("unknown operator", ident.span()),
120            }
121        },
122    })
123}
124
125fn expr_impl(stream: TokenStream) -> Result<TokenStream, TokenStream> {
126    let mut result = TokenStream::new();
127    let mut iter = stream.into_iter();
128
129    while let Some(tok) = iter.next() {
130        match tok {
131            proc_macro::TokenTree::Group(group) => {
132                let mut new_group = Group::new(
133                    group.delimiter(),
134                    expr_impl(group.stream())?,
135                );
136                new_group.set_span(group.span());
137                result.extend([TokenTree::from(new_group)]);
138            },
139            proc_macro::TokenTree::Punct(ref punct)
140                if punct.as_char() == '#' =>
141            {
142                result.extend(do_operation(tok, &mut iter)?);
143            },
144            _ => result.extend([tok]),
145        }
146    }
147
148    Ok(result)
149}
150
151/// Run string expressions on expressions
152///
153/// - `#stringify(...)`: like `stringify!(...)`
154/// - `#concat(...)`: like `concat!(...)`
155/// - `##`: like `#`
156/// - `#[...]`: like `#[...]`
157///
158/// # Examples
159/// ```
160/// use stringify_inner::sexpr;
161///
162/// assert_eq!(sexpr!(#stringify(foo)), "foo");
163/// assert_eq!(sexpr!(&#stringify(foo)[1..]), "oo");
164/// assert_eq!(sexpr!(#concat(#stringify(foo), "bar")), "foobar");
165/// ```
166#[proc_macro]
167pub fn sexpr(stream: TokenStream) -> TokenStream {
168    expr_impl(stream).map_or_else(identity, identity)
169}
170
171/// Run string expressions on attribute
172///
173/// - `#stringify(...)`: like `stringify!(...)`
174/// - `#concat(...)`: like `concat!(...)`
175/// - `##`: like `#`
176/// - `#[...]`: like `#[...]`
177///
178/// # Examples
179/// ```
180/// use stringify_inner::sexpr_attr;
181///
182/// #[sexpr_attr(doc(alias = #stringify(bar)))]
183/// fn foo() {}
184/// ```
185#[proc_macro_attribute]
186pub fn sexpr_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
187    stream([
188        expr_impl(stream([
189            TokenTree::from(Punct::new('#', Joint)),
190            Group::new(Delimiter::Bracket, attr).into(),
191        ])).map_or_else(identity, identity),
192        item,
193    ])
194}