pastey/
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                '#' => segments.push(Segment::String(LitStr {
123                    value: "#".to_string(),
124                    span: punct.span(),
125                })),
126                _ => return Err(Error::new(punct.span(), "unexpected punct")),
127            },
128            TokenTree::Group(group) => {
129                if group.delimiter() == Delimiter::None {
130                    let mut inner = group.stream().into_iter().peekable();
131                    let nested = parse(&mut inner)?;
132                    if let Some(unexpected) = inner.next() {
133                        return Err(Error::new(unexpected.span(), "unexpected token"));
134                    }
135                    segments.extend(nested);
136                } else {
137                    return Err(Error::new(group.span(), "unexpected token"));
138                }
139            }
140        }
141    }
142    Ok(segments)
143}
144
145pub(crate) fn paste(segments: &[Segment]) -> Result<String> {
146    let mut evaluated = Vec::new();
147    let mut is_lifetime = false;
148
149    for (i, segment) in segments.iter().enumerate() {
150        match segment {
151            Segment::String(segment) => {
152                if segment.value.as_str() == "#" {
153                    if i == 0 {
154                        // Enable Raw mode
155                        evaluated.push(String::from("r#"));
156                        continue;
157                    }
158                    return Err(Error::new(
159                        segment.span,
160                        "`#` is reserved keyword and it enables the raw mode \
161                            (i.e. generate Raw Identifiers) and it is only allowed in \
162                            the beginning like `[< # ... >]`",
163                    ));
164                }
165                evaluated.push(segment.value.clone());
166            }
167            Segment::Apostrophe(span) => {
168                if is_lifetime {
169                    return Err(Error::new(*span, "unexpected lifetime"));
170                }
171                is_lifetime = true;
172            }
173            Segment::Env(var) => {
174                let resolved = match std::env::var(&var.value) {
175                    Ok(resolved) => resolved,
176                    Err(_) => {
177                        return Err(Error::new(
178                            var.span,
179                            &format!("no such env var: {:?}", var.value),
180                        ));
181                    }
182                };
183                let resolved = resolved.replace('-', "_");
184                evaluated.push(resolved);
185            }
186            Segment::Modifier(colon, ident) => {
187                let last = match evaluated.pop() {
188                    Some(last) => last,
189                    None => {
190                        return Err(Error::new2(colon.span, ident.span(), "unexpected modifier"))
191                    }
192                };
193                match ident.to_string().as_str() {
194                    "lower" => {
195                        evaluated.push(last.to_lowercase());
196                    }
197                    "upper" => {
198                        evaluated.push(last.to_uppercase());
199                    }
200                    "snake" => {
201                        let mut acc = String::new();
202                        let mut prev = '_';
203                        for ch in last.chars() {
204                            if ch.is_uppercase() && prev != '_' {
205                                acc.push('_');
206                            }
207                            acc.push(ch);
208                            prev = ch;
209                        }
210                        evaluated.push(acc.to_lowercase());
211                    }
212                    "camel" | "upper_camel" | "lower_camel" => {
213                        let mut is_lower_camel = ident.to_string().as_str() == "lower_camel";
214                        let mut acc = String::new();
215                        let mut prev = '_';
216                        for ch in last.chars() {
217                            if ch != '_' {
218                                if prev == '_' {
219                                    if is_lower_camel {
220                                        for chl in ch.to_lowercase() {
221                                            acc.push(chl);
222                                        }
223                                        is_lower_camel = false;
224                                    } else {
225                                        for chu in ch.to_uppercase() {
226                                            acc.push(chu);
227                                        }
228                                    }
229                                } else if prev.is_uppercase() {
230                                    for chl in ch.to_lowercase() {
231                                        acc.push(chl);
232                                    }
233                                } else {
234                                    acc.push(ch);
235                                }
236                            }
237                            prev = ch;
238                        }
239                        evaluated.push(acc);
240                    }
241                    "camel_edge" => {
242                        let mut acc = String::new();
243                        let mut prev = '_';
244                        for ch in last.chars() {
245                            if ch != '_' {
246                                if prev == '_' {
247                                    for chu in ch.to_uppercase() {
248                                        acc.push(chu);
249                                    }
250                                } else if prev.is_uppercase() {
251                                    for chl in ch.to_lowercase() {
252                                        acc.push(chl);
253                                    }
254                                } else {
255                                    acc.push(ch);
256                                }
257                            } else if prev == '_' {
258                                acc.push(ch);
259                            }
260                            prev = ch;
261                        }
262                        evaluated.push(acc);
263                    }
264                    _ => {
265                        return Err(Error::new2(
266                            colon.span,
267                            ident.span(),
268                            "unsupported modifier",
269                        ));
270                    }
271                }
272            }
273        }
274    }
275
276    let mut pasted = evaluated.into_iter().collect::<String>();
277    if is_lifetime {
278        pasted.insert(0, '\'');
279    }
280    Ok(pasted)
281}