pastey/
segment.rs

1use crate::error::{Error, Result};
2use proc_macro::{token_stream, Delimiter, Group, Ident, Literal, 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    Replace(Colon, Group),
11}
12
13pub(crate) struct LitStr {
14    pub value: String,
15    pub span: Span,
16}
17
18pub(crate) struct Colon {
19    pub span: Span,
20}
21
22pub(crate) fn parse(tokens: &mut Peekable<token_stream::IntoIter>) -> Result<Vec<Segment>> {
23    let mut segments = Vec::new();
24    while match tokens.peek() {
25        None => false,
26        Some(TokenTree::Punct(punct)) => punct.as_char() != '>',
27        Some(_) => true,
28    } {
29        match tokens.next().unwrap() {
30            TokenTree::Ident(ident) => {
31                let mut fragment = ident.to_string();
32                if fragment.starts_with("r#") {
33                    fragment = fragment.split_off(2);
34                }
35                if fragment == "env"
36                    && match tokens.peek() {
37                        Some(TokenTree::Punct(punct)) => punct.as_char() == '!',
38                        _ => false,
39                    }
40                {
41                    let bang = tokens.next().unwrap(); // `!`
42                    let expect_group = tokens.next();
43                    let parenthesized = match &expect_group {
44                        Some(TokenTree::Group(group))
45                            if group.delimiter() == Delimiter::Parenthesis =>
46                        {
47                            group
48                        }
49                        Some(wrong) => return Err(Error::new(wrong.span(), "expected `(`")),
50                        None => {
51                            return Err(Error::new2(
52                                ident.span(),
53                                bang.span(),
54                                "expected `(` after `env!`",
55                            ));
56                        }
57                    };
58                    let mut inner = parenthesized.stream().into_iter();
59                    let lit = match inner.next() {
60                        Some(TokenTree::Literal(lit)) => lit,
61                        Some(wrong) => {
62                            return Err(Error::new(wrong.span(), "expected string literal"))
63                        }
64                        None => {
65                            return Err(Error::new2(
66                                ident.span(),
67                                parenthesized.span(),
68                                "expected string literal as argument to env! macro",
69                            ))
70                        }
71                    };
72
73                    segments.push(Segment::Env(LitStr {
74                        value: get_literal_string_value(&lit, false, false)?,
75                        span: lit.span(),
76                    }));
77
78                    if let Some(unexpected) = inner.next() {
79                        return Err(Error::new(
80                            unexpected.span(),
81                            "unexpected token in env! macro",
82                        ));
83                    }
84                } else {
85                    segments.push(Segment::String(LitStr {
86                        value: fragment,
87                        span: ident.span(),
88                    }));
89                }
90            }
91            TokenTree::Literal(lit) => {
92                segments.push(Segment::String(LitStr {
93                    value: lit.to_string(),
94                    span: lit.span(),
95                }));
96            }
97            TokenTree::Punct(punct) => match punct.as_char() {
98                '_' => segments.push(Segment::String(LitStr {
99                    value: "_".to_owned(),
100                    span: punct.span(),
101                })),
102                '\'' => segments.push(Segment::Apostrophe(punct.span())),
103                ':' => {
104                    let colon_span = punct.span();
105                    let colon = Colon { span: colon_span };
106                    let ident = match tokens.next() {
107                        Some(TokenTree::Ident(ident)) => ident,
108                        wrong => {
109                            let span = wrong.as_ref().map_or(colon_span, TokenTree::span);
110                            return Err(Error::new(span, "expected identifier after `:`"));
111                        }
112                    };
113
114                    if ident.to_string().as_str() == "replace" {
115                        let replace = tokens.next();
116
117                        match replace {
118                            Some(TokenTree::Group(group))
119                                if group.delimiter() == Delimiter::Parenthesis =>
120                            {
121                                segments.push(Segment::Replace(colon, group));
122                            }
123                            _ => {
124                                return Err(Error::new2(
125                                    colon.span,
126                                    ident.span(),
127                                    "expected `(` after replace modifier",
128                                ));
129                            }
130                        }
131                    } else {
132                        segments.push(Segment::Modifier(colon, ident));
133                    }
134                }
135                '#' => segments.push(Segment::String(LitStr {
136                    value: "#".to_string(),
137                    span: punct.span(),
138                })),
139                _ => return Err(Error::new(punct.span(), "unexpected punct")),
140            },
141            TokenTree::Group(group) => {
142                if group.delimiter() == Delimiter::None {
143                    let mut inner = group.stream().into_iter().peekable();
144                    let nested = parse(&mut inner)?;
145                    if let Some(unexpected) = inner.next() {
146                        return Err(Error::new(unexpected.span(), "unexpected token"));
147                    }
148                    segments.extend(nested);
149                } else {
150                    return Err(Error::new(group.span(), "unexpected token"));
151                }
152            }
153        }
154    }
155    Ok(segments)
156}
157
158pub(crate) fn paste(segments: &[Segment]) -> Result<String> {
159    let mut evaluated = Vec::new();
160    let mut is_lifetime = false;
161
162    for (i, segment) in segments.iter().enumerate() {
163        match segment {
164            Segment::String(segment) => {
165                if segment.value.as_str() == "#" {
166                    if i == 0 {
167                        // Enable Raw mode
168                        evaluated.push(String::from("r#"));
169                        continue;
170                    }
171                    return Err(Error::new(
172                        segment.span,
173                        "`#` is reserved keyword and it enables the raw mode \
174                            (i.e. generate Raw Identifiers) and it is only allowed in \
175                            the beginning like `[< # ... >]`",
176                    ));
177                }
178                evaluated.push(segment.value.clone());
179            }
180            Segment::Apostrophe(span) => {
181                if is_lifetime {
182                    return Err(Error::new(*span, "unexpected lifetime"));
183                }
184                is_lifetime = true;
185            }
186            Segment::Env(var) => {
187                let resolved = match std::env::var(&var.value) {
188                    Ok(resolved) => resolved,
189                    Err(_) => {
190                        return Err(Error::new(
191                            var.span,
192                            &format!("no such env var: {:?}", var.value),
193                        ));
194                    }
195                };
196                let resolved = resolved.replace('-', "_");
197                evaluated.push(resolved);
198            }
199            Segment::Modifier(colon, ident) => {
200                let last = match evaluated.pop() {
201                    Some(last) => last,
202                    None => {
203                        return Err(Error::new2(colon.span, ident.span(), "unexpected modifier"))
204                    }
205                };
206                match ident.to_string().as_str() {
207                    "lower" => {
208                        evaluated.push(last.to_lowercase());
209                    }
210                    "upper" => {
211                        evaluated.push(last.to_uppercase());
212                    }
213                    "snake" => {
214                        let mut acc = String::new();
215                        let mut prev = '_';
216                        for ch in last.chars() {
217                            if ch.is_uppercase() && prev != '_' {
218                                acc.push('_');
219                            }
220                            acc.push(ch);
221                            prev = ch;
222                        }
223                        evaluated.push(acc.to_lowercase());
224                    }
225                    "camel" | "upper_camel" | "lower_camel" => {
226                        let mut is_lower_camel = ident.to_string().as_str() == "lower_camel";
227                        let mut acc = String::new();
228                        let mut prev = '_';
229                        for ch in last.chars() {
230                            if ch != '_' {
231                                if prev == '_' {
232                                    if is_lower_camel {
233                                        for chl in ch.to_lowercase() {
234                                            acc.push(chl);
235                                        }
236                                        is_lower_camel = false;
237                                    } else {
238                                        for chu in ch.to_uppercase() {
239                                            acc.push(chu);
240                                        }
241                                    }
242                                } else if prev.is_uppercase() {
243                                    for chl in ch.to_lowercase() {
244                                        acc.push(chl);
245                                    }
246                                } else {
247                                    acc.push(ch);
248                                }
249                            }
250                            prev = ch;
251                        }
252                        evaluated.push(acc);
253                    }
254                    "camel_edge" => {
255                        let mut acc = String::new();
256                        let mut prev = '_';
257                        for ch in last.chars() {
258                            if ch != '_' {
259                                if prev == '_' {
260                                    for chu in ch.to_uppercase() {
261                                        acc.push(chu);
262                                    }
263                                } else if prev.is_uppercase() {
264                                    for chl in ch.to_lowercase() {
265                                        acc.push(chl);
266                                    }
267                                } else {
268                                    acc.push(ch);
269                                }
270                            } else if prev == '_' {
271                                acc.push(ch);
272                            }
273                            prev = ch;
274                        }
275                        evaluated.push(acc);
276                    }
277                    _ => {
278                        return Err(Error::new2(
279                            colon.span,
280                            ident.span(),
281                            "unsupported modifier",
282                        ));
283                    }
284                }
285            }
286            Segment::Replace(colon, group) => {
287                let mut inner_stream = group.stream().into_iter();
288                let from = inner_stream.next();
289                let punct = inner_stream.next();
290                let to = inner_stream.next();
291
292                if let Some(unexpected_token) = inner_stream.next() {
293                    return Err(Error::new(unexpected_token.span(), "expected `)`"));
294                }
295
296                match (from, punct, to) {
297                    (Some(from), Some(TokenTree::Punct(punct)), Some(to))
298                        if punct.as_char() == ',' =>
299                    {
300                        let last =
301                            match evaluated.pop() {
302                                Some(last) => last,
303                                None => return Err(Error::new2(
304                                    colon.span,
305                                    group.span(),
306                                    "replace modifier requires a preceding value to operate on.",
307                                )),
308                            };
309
310                        let from_str = get_token_tree_string_value(&from)?;
311                        let to_str = get_token_tree_string_value(&to)?;
312
313                        let new_ident = last.replace(&from_str, &to_str);
314
315                        evaluated.push(new_ident);
316                    }
317                    _ => {
318                        return Err(Error::new(
319                            group.span(),
320                            "expected replace modifier format: `:replace(\"from\", \"to\")`",
321                        ))
322                    }
323                }
324            }
325        }
326    }
327
328    let mut pasted = evaluated.into_iter().collect::<String>();
329    if is_lifetime {
330        pasted.insert(0, '\'');
331    }
332    Ok(pasted)
333}
334
335fn get_literal_string_value(l: &Literal, parse_char: bool, parse_numbers: bool) -> Result<String> {
336    let l_str = l.to_string();
337
338    if ((l_str.starts_with('"') && l_str.ends_with('"'))
339        || (parse_char && l_str.starts_with('\'') && l_str.ends_with('\'')))
340        && l_str.len() >= 2
341    {
342        // TODO: maybe handle escape sequences in the string if
343        // someone has a use case.
344        Ok(String::from(&l_str[1..l_str.len() - 1]))
345    } else if parse_numbers {
346        Ok(l_str)
347    } else {
348        Err(Error::new(l.span(), "expected string literal"))
349    }
350}
351
352fn get_token_tree_string_value(t: &TokenTree) -> Result<String> {
353    match t {
354        TokenTree::Ident(ident) => Ok(ident.to_string()),
355        TokenTree::Literal(literal) => get_literal_string_value(literal, true, true),
356        _ => Err(Error::new(t.span(), "Expected either Ident, or Literal.")),
357    }
358}