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}
234