]>
Commit | Line | Data |
---|---|---|
5869c6ff XL |
1 | use proc_macro::{Delimiter, Ident, Literal, Span, TokenStream, TokenTree}; |
2 | ||
3 | use crate::{ | |
4 | iter::TokenIter, | |
5 | to_tokens::ToTokens, | |
6 | utils::{parse_as_empty, tt_span}, | |
7 | Result, | |
8 | }; | |
9 | ||
10 | pub(crate) struct Func { | |
11 | pub(crate) attrs: Vec<Attribute>, | |
12 | pub(crate) sig: Vec<TokenTree>, | |
13 | pub(crate) body: TokenTree, | |
14 | pub(crate) print_const: bool, | |
15 | } | |
16 | ||
17 | pub(crate) fn parse_input(input: TokenStream) -> Result<Func> { | |
18 | let mut input = TokenIter::new(input); | |
19 | ||
20 | let attrs = parse_attrs(&mut input)?; | |
21 | let sig = parse_signature(&mut input); | |
22 | let body = input.next(); | |
23 | parse_as_empty(input)?; | |
24 | ||
25 | if body.is_none() | |
26 | || !sig | |
27 | .iter() | |
28 | .any(|tt| if let TokenTree::Ident(i) = tt { i.to_string() == "fn" } else { false }) | |
29 | { | |
30 | return Err(error!( | |
31 | Span::call_site(), | |
32 | "#[const_fn] attribute may only be used on functions" | |
33 | )); | |
34 | } | |
35 | ||
36 | Ok(Func { attrs, sig, body: body.unwrap(), print_const: true }) | |
37 | } | |
38 | ||
39 | impl ToTokens for Func { | |
40 | fn to_tokens(&self, tokens: &mut TokenStream) { | |
41 | self.attrs.iter().for_each(|attr| attr.to_tokens(tokens)); | |
42 | if self.print_const { | |
43 | self.sig.iter().for_each(|attr| attr.to_tokens(tokens)); | |
44 | } else { | |
45 | self.sig | |
46 | .iter() | |
47 | .filter( | |
48 | |tt| if let TokenTree::Ident(i) = tt { i.to_string() != "const" } else { true }, | |
49 | ) | |
50 | .for_each(|tt| tt.to_tokens(tokens)); | |
51 | } | |
52 | self.body.to_tokens(tokens); | |
53 | } | |
54 | } | |
55 | ||
56 | fn parse_signature(input: &mut TokenIter) -> Vec<TokenTree> { | |
57 | let mut sig = Vec::new(); | |
58 | let mut has_const = false; | |
59 | loop { | |
60 | match input.peek() { | |
61 | Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => break, | |
62 | None => break, | |
63 | Some(TokenTree::Ident(i)) if !has_const => { | |
64 | match &*i.to_string() { | |
65 | "const" => has_const = true, | |
66 | "async" | "unsafe" | "extern" | "fn" => { | |
67 | sig.push(TokenTree::Ident(Ident::new("const", i.span()))); | |
68 | has_const = true; | |
69 | } | |
70 | _ => {} | |
71 | } | |
72 | sig.push(input.next().unwrap()); | |
73 | } | |
74 | Some(_) => sig.push(input.next().unwrap()), | |
75 | } | |
76 | } | |
77 | sig | |
78 | } | |
79 | ||
80 | fn parse_attrs(input: &mut TokenIter) -> Result<Vec<Attribute>> { | |
81 | let mut attrs = Vec::new(); | |
82 | loop { | |
83 | let pound_token = match input.peek() { | |
84 | Some(TokenTree::Punct(p)) if p.as_char() == '#' => input.next().unwrap(), | |
85 | _ => break, | |
86 | }; | |
87 | let group = match input.peek() { | |
88 | Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => { | |
89 | input.next().unwrap() | |
90 | } | |
91 | tt => return Err(error!(tt_span(tt), "expected `[`")), | |
92 | }; | |
93 | attrs.push(Attribute { pound_token, group }); | |
94 | } | |
95 | Ok(attrs) | |
96 | } | |
97 | ||
98 | pub(crate) struct Attribute { | |
99 | // `#` | |
100 | pub(crate) pound_token: TokenTree, | |
101 | // `[...]` | |
102 | pub(crate) group: TokenTree, | |
103 | } | |
104 | ||
105 | impl ToTokens for Attribute { | |
106 | fn to_tokens(&self, tokens: &mut TokenStream) { | |
107 | self.pound_token.to_tokens(tokens); | |
108 | self.group.to_tokens(tokens); | |
109 | } | |
110 | } | |
111 | ||
112 | pub(crate) struct LitStr { | |
113 | pub(crate) token: Literal, | |
114 | value: String, | |
115 | } | |
116 | ||
117 | impl LitStr { | |
118 | pub(crate) fn new(token: Literal) -> Result<Self> { | |
119 | let value = token.to_string(); | |
120 | // unlike `syn::LitStr`, only accepts `"..."` | |
121 | if value.starts_with('"') && value.ends_with('"') { | |
122 | Ok(Self { token, value }) | |
123 | } else { | |
124 | Err(error!(token.span(), "expected string literal")) | |
125 | } | |
126 | } | |
127 | ||
128 | pub(crate) fn value(&self) -> &str { | |
129 | &self.value[1..self.value.len() - 1] | |
130 | } | |
131 | ||
132 | pub(crate) fn span(&self) -> Span { | |
133 | self.token.span() | |
134 | } | |
135 | } | |
136 | ||
137 | impl ToTokens for LitStr { | |
138 | fn to_tokens(&self, tokens: &mut TokenStream) { | |
139 | self.token.to_tokens(tokens); | |
140 | } | |
141 | } |