1 use crate::error::{Error, Result};
2 use proc_macro::{token_stream, Delimiter, Ident, Span, TokenTree};
3 use std::iter::Peekable;
4 
5 pub(crate) enum Segment {
6     String(LitStr),
7     Apostrophe(Span),
8     Env(LitStr),
9     Modifier(Colon, Ident),
10 }
11 
12 pub(crate) struct LitStr {
13     pub value: String,
14     pub span: Span,
15 }
16 
17 pub(crate) struct Colon {
18     pub span: Span,
19 }
20 
parse(tokens: &mut Peekable<token_stream::IntoIter>) -> Result<Vec<Segment>>21 pub(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 
paste(segments: &[Segment]) -> Result<String>141 pub(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 }
234