1mod lit_cast;
2
3use pmutil::ToTokensExt;
4use std::collections::HashMap;
5use syn::punctuated::Punctuated;
6use syn::{AttrStyle, Attribute, Expr, ExprLit, Lit, Meta, MetaList, MetaNameValue, Result, Token};
7
8use crate::lit_cast::FromLit;
9
10fn check_and_pop_hd<'a>(meta: &Meta, id: &'a [&'a str]) -> Option<&'a [&'a str]> {
11 id.split_first().and_then(|(hd, tl)| {
12 if meta.path().is_ident(hd) {
13 Some(tl)
14 } else {
15 None
16 }
17 })
18}
19
20fn iter_meta_list<T, F>(meta_list: &MetaList, mut f: F) -> Result<T>
21where
22 F: FnMut(&mut syn::punctuated::Iter<Meta>) -> T,
23{
24 meta_list
25 .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
26 .map(|nested_metas| f(&mut nested_metas.iter()))
27}
28
29pub fn contains_attribute(attrs: &[Attribute], id: &[&str]) -> bool {
30 attrs.iter().any(|Attribute { style, meta, .. }| {
31 *style == AttrStyle::Outer && contains_attribute_impl(meta, id)
32 })
33}
34
35fn contains_attribute_impl(meta: &Meta, id: &[&str]) -> bool {
36 let id = match check_and_pop_hd(meta, id) {
37 Some(id) => id,
38 None => {
39 return false;
40 }
41 };
42
43 match meta {
44 Meta::Path(..) => id.is_empty(),
45 Meta::List(meta_list) => iter_meta_list(meta_list, |iter| {
46 iter.any(|meta| contains_attribute_impl(meta, id))
47 })
48 .unwrap_or(false),
49 Meta::NameValue(..) => false,
50 }
51}
52
53pub fn get_attribute_value<T: FromLit>(attrs: &[Attribute], id: &[&str]) -> Option<T> {
54 attrs.iter().find_map(|Attribute { style, meta, .. }| {
55 if *style != AttrStyle::Outer {
56 return None;
57 }
58 get_attribute_value_impl(meta, id).and_then(|value| T::from_lit(value).ok())
59 })
60}
61
62fn get_attribute_value_impl(meta: &Meta, id: &[&str]) -> Option<Lit> {
63 let id = match check_and_pop_hd(meta, id) {
64 Some(id) => id,
65 None => {
66 return None;
67 }
68 };
69
70 match meta {
71 Meta::NameValue(MetaNameValue {
72 value: Expr::Lit(ExprLit { lit, .. }),
73 ..
74 }) if id.is_empty() => Some(lit.clone()),
75 Meta::List(meta_list) => iter_meta_list(meta_list, |iter| {
76 iter.find_map(|meta| get_attribute_value_impl(meta, id))
77 })
78 .unwrap_or(None),
79 Meta::Path(..) | Meta::NameValue(..) => None,
80 }
81}
82
83pub fn get_attribute_map(attrs: &[Attribute], separator: &str) -> HashMap<String, Vec<Lit>> {
84 let mut result = HashMap::new();
85 attrs.iter().for_each(|Attribute { style, meta, .. }| {
86 if *style == AttrStyle::Outer {
87 get_attribute_map_impl(&mut result, meta, "", separator);
88 }
89 });
90 result
91}
92
93fn get_attribute_map_impl(
94 map: &mut HashMap<String, Vec<Lit>>,
95 meta: &Meta,
96 prefix: &str,
97 separator: &str,
98) -> () {
99 let key = {
100 let path = meta.path().dump();
101 if prefix.is_empty() {
102 path.to_string()
103 } else {
104 format!("{}{}{}", prefix, separator, path)
105 }
106 };
107
108 match meta {
109 Meta::Path(..) => {
110 assert!(!map.contains_key(&key), "{} already exists.", key);
111 map.insert(key, vec![]);
112 }
113 Meta::NameValue(MetaNameValue {
114 value: Expr::Lit(ExprLit { lit, .. }),
115 ..
116 }) => map
117 .get_mut(&key)
118 .map(|values| {
119 assert!(!values.is_empty(), "conflicts among `{}` attributes.", key);
120 values.push(lit.clone());
121 })
122 .unwrap_or_else(|| {
123 map.insert(key, vec![lit.clone()]);
124 }),
125 Meta::NameValue(..) => (),
126 Meta::List(meta_list) => iter_meta_list(meta_list, |iter| {
127 iter.for_each(|meta| get_attribute_map_impl(map, meta, &key, separator))
128 })
129 .unwrap_or(()),
130 }
131}
132
133#[cfg(test)]
134mod test {
135 use super::*;
136 use proc_macro2::Span;
137 use syn::{parse_quote, LitStr};
138
139 fn lit_str(s: &str) -> Lit {
140 Lit::Str(LitStr::new(s, Span::call_site()))
141 }
142
143 #[test]
144 fn test_contains_attribute_impl() {
145 let attr: Attribute = parse_quote!(#[level0]);
146 assert!(contains_attribute(&[attr], &["level0"]));
147
148 let attr: Attribute = parse_quote!(#[level0(level1, level1_1(level2, level2_1 = "hello", level2_2), level1_2)]);
149 let attr = [attr];
150
151 assert!(!contains_attribute(&attr, &[]));
152
153 assert!(!contains_attribute(&attr, &["not"]));
154
155 assert!(!contains_attribute(&attr, &["level0"]));
156
157 assert!(contains_attribute(&attr, &["level0", "level1"]));
158
159 assert!(!contains_attribute(&attr, &["level0", "level1_1"]));
160
161 assert!(contains_attribute(&attr, &["level0", "level1_2"]));
162
163 assert!(contains_attribute(&attr, &["level0", "level1_1", "level2"]));
164
165 assert!(contains_attribute(
166 &attr,
167 &["level0", "level1_1", "level2_2"],
168 ),);
169
170 assert!(!contains_attribute(
171 &attr,
172 &["level0", "level1_1", "level2_1"],
173 ),);
174 }
175
176 #[test]
177 fn test_get_attribute_value_impl() {
178 let attr: Attribute = parse_quote!(#[level0(level1 = "hi", level1_1(level2 = "bye"))]);
179
180 let meta = attr.meta;
181
182 assert_eq!(get_attribute_value_impl(&meta, &[]), None);
183
184 assert_eq!(get_attribute_value_impl(&meta, &["not"]), None);
185
186 assert_eq!(get_attribute_value_impl(&meta, &["level0"]), None);
187
188 assert_eq!(
189 get_attribute_value_impl(&meta, &["level0", "level1"]),
190 Some(lit_str("hi"))
191 );
192
193 assert_eq!(
194 get_attribute_value_impl(&meta, &["level0", "level1_1"]),
195 None
196 );
197
198 assert_eq!(
199 get_attribute_value_impl(&meta, &["level0", "level1_1", "level2"]),
200 Some(lit_str("bye"))
201 );
202
203 let attr: Attribute = parse_quote!(#[doc = "hi"]);
204
205 let meta = attr.meta;
206
207 assert_eq!(
208 get_attribute_value_impl(&meta, &["doc"]),
209 Some(lit_str("hi"))
210 );
211 }
212
213 #[test]
214 fn test_get_attribute_value() {
215 let attr: Attribute = parse_quote!(#[level0 = "hi"]);
216 assert_eq!(
217 get_attribute_value(&[attr], &["level0"]),
218 Some(lit_str("hi"))
219 );
220
221 let attr: Attribute = parse_quote!(#[level0(level1 = "hi", level1_1(level2 = false))]);
222 let attr = [attr];
223
224 assert_eq!(get_attribute_value::<String>(&attr, &[""]), None);
225
226 assert_eq!(get_attribute_value::<String>(&attr, &["not"]), None);
227
228 assert_eq!(get_attribute_value::<String>(&attr, &["level0"]), None);
229
230 assert_eq!(
231 get_attribute_value(&attr, &["level0", "level1"]),
232 Some("hi".to_string())
233 );
234
235 assert_eq!(
236 get_attribute_value::<Lit>(&attr, &["level0", "level1_1"]),
237 None
238 );
239
240 assert_eq!(
241 get_attribute_value(&attr, &["level0", "level1_1", "level2"]),
242 Some(false)
243 );
244 }
245
246 #[test]
247 fn test_get_attribute_map_impl() {
248 let attr: Attribute =
249 parse_quote!(#[level0(level1 = "hi", level1 = "hi", level1_1(level2 = "bye"))]);
250
251 let meta = attr.meta;
252
253 let mut result = HashMap::new();
254 get_attribute_map_impl(&mut result, &meta, "", ".");
255 assert_eq!(
256 result,
257 vec![
258 (
259 "level0.level1".to_string(),
260 vec![lit_str("hi"), lit_str("hi")],
261 ),
262 ("level0.level1_1.level2".to_string(), vec![lit_str("bye")]),
263 ]
264 .into_iter()
265 .collect()
266 );
267 }
268
269 #[test]
270 fn test_get_attribute_map() {
271 assert_eq!(
272 get_attribute_map(
273 &[
274 parse_quote!(#[level9]),
275 parse_quote!(#[level0_0 = "greeting"]),
276 parse_quote!(#[level0(level8)]),
277 parse_quote!(#[level0(level1 = "hi", level1_1(level2 = "bye"))]),
278 parse_quote!(#[level0(level1 = "hi", level1_1(level2 = "bye"))]),
279 parse_quote!(#[gen0(gen1 = "amoeba", gen1_1 = "monad", gen1_2(gen2 = "monoid"))])
280 ],
281 ".",
282 ),
283 vec![
284 ("level0_0".to_string(), vec![lit_str("greeting")]),
285 ("level9".to_string(), vec![]),
286 ("level0.level8".to_string(), vec![]),
287 (
288 "level0.level1".to_string(),
289 vec![lit_str("hi"), lit_str("hi")],
290 ),
291 (
292 "level0.level1_1.level2".to_string(),
293 vec![lit_str("bye"), lit_str("bye")],
294 ),
295 ("gen0.gen1".to_string(), vec![lit_str("amoeba")]),
296 ("gen0.gen1_1".to_string(), vec![lit_str("monad")]),
297 ("gen0.gen1_2.gen2".to_string(), vec![lit_str("monoid")]),
298 ]
299 .into_iter()
300 .collect()
301 );
302 }
303}