paste/
segment.rs

1use crate::error::{Error, Result};
2use proc_macro::{token_stream, Delimiter, Ident, Span, TokenTree};
3use std::iter::Peekable;
4
5pub(crate) enum Segment {
6    String(LitStr),
7    Apostrophe(Span),
8    Env(LitStr),
9    Modifier(Colon, Ident),
10}
11
12pub(crate) struct LitStr {
13    pub value: String,
14    pub span: Span,
15}
16
17pub(crate) struct Colon {
18    pub span: Span,
19}
20
21pub(crate) fn parse(tokens: &mut Peekable<token_stream::IntoIter>) -> Result<Vec<Segment>> {
22    let mut segments = Vec::new();
23    while match tokens.peek() {
24        None => false,
25        Some(TokenTree::Punct(punct)) => punct.as_char() != '>',
26        Some(_) => true,
27    } {
28        match tokens.next().unwrap() {
29            TokenTree::Ident(ident) => {
30                let mut fragment = ident.to_string();
31                if fragment.starts_with("r#") {
32                    fragment = fragment.split_off(2);
33                }
34                if fragment == "env"
35                    && match tokens.peek() {
36                        Some(TokenTree::Punct(punct)) => punct.as_char() == '!',
37                        _ => false,
38                    }
39                {
40                    let bang = tokens.next().unwrap(); // `!`
41                    let expect_group = tokens.next();
42                    let parenthesized = match &expect_group {
43                        Some(TokenTree::Group(group))
44                            if group.delimiter() == Delimiter::Parenthesis =>
45                        {
46                            group
47                        }
48                        Some(wrong) => return Err(Error::new(wrong.span(), "expected `(`")),
49                        None => {
50                            return Err(Error::new2(
51                                ident.span(),
52                                bang.span(),
53                                "expected `(` after `env!`",
54                            ));
55                        }
56                    };
57                    let mut inner = parenthesized.stream().into_iter();
58                    let lit = match inner.next() {
59                        Some(TokenTree::Literal(lit)) => lit,
60                        Some(wrong) => {
61                            return Err(Error::new(wrong.span(), "expected string literal"))
62                        }
63                        None => {
64                            return Err(Error::new2(
65                                ident.span(),
66                                parenthesized.span(),
67                                "expected string literal as argument to env! macro",
68                            ))
69                        }
70                    };
71                    let lit_string = lit.to_string();
72                    if lit_string.starts_with('"')
73                        && lit_string.ends_with('"')
74                        && lit_string.len() >= 2
75                    {
76                        // TODO: maybe handle escape sequences in the string if
77                        // someone has a use case.
78                        segments.push(Segment::Env(LitStr {
79                            value: lit_string[1..lit_string.len() - 1].to_owned(),
80                            span: lit.span(),
81                        }));
82                    } else {
83                        return Err(Error::new(lit.span(), "expected string literal"));
84                    }
85                    if let Some(unexpected) = inner.next() {
86                        return Err(Error::new(
87                            unexpected.span(),
88                            "unexpected token in env! macro",
89                        ));
90                    }
91                } else {
92                    segments.push(Segment::String(LitStr {
93                        value: fragment,
94                        span: ident.span(),
95                    }));
96                }
97            }
98            TokenTree::Literal(lit) => {
99                segments.push(Segment::String(LitStr {
100                    value: lit.to_string(),
101                    span: lit.span(),
102                }));
103            }
104            TokenTree::Punct(punct) => match punct.as_char() {
105                '_' => segments.push(Segment::String(LitStr {
106                    value: "_".to_owned(),
107                    span: punct.span(),
108                })),
109                '\'' => segments.push(Segment::Apostrophe(punct.span())),
110                ':' => {
111                    let colon_span = punct.span();
112                    let colon = Colon { span: colon_span };
113                    let ident = match tokens.next() {
114                        Some(TokenTree::Ident(ident)) => ident,
115                        wrong => {
116                            let span = wrong.as_ref().map_or(colon_span, TokenTree::span);
117                            return Err(Error::new(span, "expected identifier after `:`"));
118                        }
119                    };
120                    segments.push(Segment::Modifier(colon, ident));
121                }
122                _ => return Err(Error::new(punct.span(), "unexpected punct")),
123            },
124            TokenTree::Group(group) => {
125                if group.delimiter() == Delimiter::None {
126                    let mut inner = group.stream().into_iter().peekable();
127                    let nested = parse(&mut inner)?;
128                    if let Some(unexpected) = inner.next() {
129                        return Err(Error::new(unexpected.span(), "unexpected token"));
130                    }
131                    segments.extend(nested);
132                } else {
133                    return Err(Error::new(group.span(), "unexpected token"));
134                }
135            }
136        }
137    }
138    Ok(segments)
139}
140
141pub(crate) fn paste(segments: &[Segment]) -> Result<String> {
142    let mut evaluated = Vec::new();
143    let mut is_lifetime = false;
144
145    for segment in segments {
146        match segment {
147            Segment::String(segment) => {
148                evaluated.push(segment.value.clone());
149            }
150            Segment::Apostrophe(span) => {
151                if is_lifetime {
152                    return Err(Error::new(*span, "unexpected lifetime"));
153                }
154                is_lifetime = true;
155            }
156            Segment::Env(var) => {
157                let resolved = match std::env::var(&var.value) {
158                    Ok(resolved) => resolved,
159                    Err(_) => {
160                        return Err(Error::new(
161                            var.span,
162                            &format!("no such env var: {:?}", var.value),
163                        ));
164                    }
165                };
166                let resolved = resolved.replace('-', "_");
167                evaluated.push(resolved);
168            }
169            Segment::Modifier(colon, ident) => {
170                let last = match evaluated.pop() {
171                    Some(last) => last,
172                    None => {
173                        return Err(Error::new2(colon.span, ident.span(), "unexpected modifier"))
174                    }
175                };
176                match ident.to_string().as_str() {
177                    "lower" => {
178                        evaluated.push(last.to_lowercase());
179                    }
180                    "upper" => {
181                        evaluated.push(last.to_uppercase());
182                    }
183                    "snake" => {
184                        let mut acc = String::new();
185                        let mut prev = '_';
186                        for ch in last.chars() {
187                            if ch.is_uppercase() && prev != '_' {
188                                acc.push('_');
189                            }
190                            acc.push(ch);
191                            prev = ch;
192                        }
193                        evaluated.push(acc.to_lowercase());
194                    }
195                    "camel" => {
196                        let mut acc = String::new();
197                        let mut prev = '_';
198                        for ch in last.chars() {
199                            if ch != '_' {
200                                if prev == '_' {
201                                    for chu in ch.to_uppercase() {
202                                        acc.push(chu);
203                                    }
204                                } else if prev.is_uppercase() {
205                                    for chl in ch.to_lowercase() {
206                                        acc.push(chl);
207                                    }
208                                } else {
209                                    acc.push(ch);
210                                }
211                            }
212                            prev = ch;
213                        }
214                        evaluated.push(acc);
215                    }
216                    _ => {
217                        return Err(Error::new2(
218                            colon.span,
219                            ident.span(),
220                            "unsupported modifier",
221                        ));
222                    }
223                }
224            }
225        }
226    }
227
228    let mut pasted = evaluated.into_iter().collect::<String>();
229    if is_lifetime {
230        pasted.insert(0, '\'');
231    }
232    Ok(pasted)
233}