1 // vim: tw=80
2 use super::*;
3 use std::collections::HashMap;
4 use syn::parse::{Parse, ParseStream};
5 
6 /// A single automock attribute
7 // This enum is very short-lived, so it's fine not to box it.
8 #[allow(clippy::large_enum_variant)]
9 enum Attr {
10     Type(TraitItemType),
11 }
12 
13 impl Parse for Attr {
parse(input: ParseStream) -> parse::Result<Self>14     fn parse(input: ParseStream) -> parse::Result<Self> {
15         let lookahead = input.lookahead1();
16         if lookahead.peek(Token![type]) {
17             input.parse().map(Attr::Type)
18         } else {
19             Err(lookahead.error())
20         }
21     }
22 }
23 
24 /// automock attributes
25 #[derive(Debug, Default)]
26 pub(crate) struct Attrs {
27     pub attrs: HashMap<Ident, Type>,
28 }
29 
30 impl Attrs {
get_path(&self, path: &Path) -> Option<Type>31     fn get_path(&self, path: &Path) -> Option<Type> {
32         if path.leading_colon.is_none() & (path.segments.len() == 2) {
33             if path.segments.first().unwrap().ident == "Self" {
34                 let ident = &path.segments.last().unwrap().ident;
35                 self.attrs.get(ident).cloned()
36             } else {
37                 None
38             }
39         } else {
40             None
41         }
42     }
43 
substitute_item_impl(&self, item_impl: &mut ItemImpl)44     pub(crate) fn substitute_item_impl(&self, item_impl: &mut ItemImpl) {
45         let (_, trait_path, _) = item_impl.trait_.as_ref()
46             .expect("Should only be called for trait item impls");
47         let trait_ident = find_ident_from_path(trait_path).0;
48         for item in item_impl.items.iter_mut() {
49             if let ImplItem::Fn(method) = item {
50                 let sig = &mut method.sig;
51                 for fn_arg in sig.inputs.iter_mut() {
52                     if let FnArg::Typed(arg) = fn_arg {
53                         self.substitute_type(&mut arg.ty, &trait_ident);
54                     }
55                 }
56                 if let ReturnType::Type(_, ref mut ty) = &mut sig.output {
57                     self.substitute_type(ty, &trait_ident);
58                 }
59             }
60         }
61     }
62 
substitute_path_segment(&self, seg: &mut PathSegment, traitname: &Ident)63     fn substitute_path_segment(&self, seg: &mut PathSegment, traitname: &Ident)
64     {
65         match &mut seg.arguments {
66             PathArguments::None => /* nothing to do */(),
67             PathArguments::Parenthesized(p) => {
68                 compile_error(p.span(),
69                     "Mockall does not support mocking Fn objects.  See https://github.com/asomers/mockall/issues/139");
70             },
71             PathArguments::AngleBracketed(abga) => {
72                 for arg in abga.args.iter_mut() {
73                     match arg {
74                         GenericArgument::Lifetime(_) => {
75                             /*
76                              * Nothing to do, as long as lifetimes can't be
77                              * associated types
78                              */
79                         }
80                         GenericArgument::Type(ty) => {
81                             self.substitute_type(ty, traitname)
82                         },
83                         GenericArgument::AssocConst(_) => {
84                             // Nothing to do
85                         }
86                         GenericArgument::AssocType(at) => {
87                             self.substitute_type(&mut at.ty, traitname);
88                         }
89                         // TODO: Constraints
90                         _ => {
91                             // Not handled. Hopefully doing nothing works.
92                         }
93                     }
94                 }
95             },
96         }
97     }
98 
99     /// Recursively substitute types in the input
substitute_type(&self, ty: &mut Type, traitname: &Ident)100     fn substitute_type(&self, ty: &mut Type, traitname: &Ident) {
101         match ty {
102             Type::Slice(s) => {
103                 self.substitute_type(s.elem.as_mut(), traitname)
104             },
105             Type::Array(a) => {
106                 self.substitute_type(a.elem.as_mut(), traitname)
107             },
108             Type::Ptr(p) => {
109                 self.substitute_type(p.elem.as_mut(), traitname)
110             },
111             Type::Reference(r) => {
112                 self.substitute_type(r.elem.as_mut(), traitname)
113             },
114             Type::BareFn(bfn) => {
115                 for fn_arg in bfn.inputs.iter_mut() {
116                     self.substitute_type(&mut fn_arg.ty, traitname);
117                 }
118                 if let ReturnType::Type(_, ref mut ty) = &mut bfn.output {
119                     self.substitute_type(ty, traitname);
120                 }
121             },
122             Type::Tuple(tuple) => {
123                 for elem in tuple.elems.iter_mut() {
124                     self.substitute_type(elem, traitname)
125                 }
126             }
127             Type::Path(path) => {
128                 if let Some(ref qself) = path.qself {
129                     let qp = if let Type::Path(p) = qself.ty.as_ref() {
130                         &p.path
131                     } else {
132                         panic!("QSelf's type isn't a path?")
133                     };
134                     let qident = &qp.segments.first().unwrap().ident;
135                     if qself.position != 1
136                         || qp.segments.len() != 1
137                         || path.path.segments.len() != 2
138                         || qident != "Self" {
139                         compile_error(path.span(),
140                             "QSelf is a work in progress");
141                     }
142 
143                     let mut seg_iter = path.path.segments.iter().rev();
144                     let last_seg = seg_iter.next().unwrap();
145                     let to_sub = &last_seg.ident;
146                     let penultimate_seg = seg_iter.next().unwrap();
147                     let qident = &penultimate_seg.ident;
148                     drop(seg_iter);
149 
150                     if qident != traitname {
151                         compile_error(qident.span(),
152                             "Mockall does not support QSelf substitutions except for the trait being mocked");
153                     }
154                     if let Some(new_type) = self.attrs.get(to_sub) {
155                         *ty = new_type.clone();
156                     } else {
157                         compile_error(to_sub.span(),
158                             "Unknown type substitution for QSelf");
159                     }
160                 } else if let Some(newty) = self.get_path(&path.path) {
161                     *ty = newty;
162                 } else {
163                     for seg in path.path.segments.iter_mut() {
164                         self.substitute_path_segment(seg, traitname);
165                     }
166                 }
167             },
168             Type::TraitObject(to) => {
169                 for bound in to.bounds.iter_mut() {
170                     self.substitute_type_param_bound(bound, traitname);
171                 }
172             },
173             Type::ImplTrait(it) => {
174                 for bound in it.bounds.iter_mut() {
175                     self.substitute_type_param_bound(bound, traitname);
176                 }
177             },
178             Type::Paren(p) => {
179                 self.substitute_type(p.elem.as_mut(), traitname)
180             },
181             Type::Group(g) => {
182                 self.substitute_type(g.elem.as_mut(), traitname)
183             },
184             Type::Macro(_) | Type::Verbatim(_) => {
185                 compile_error(ty.span(),
186                     "mockall_derive does not support this type when using associated types");
187             },
188             Type::Infer(_) | Type::Never(_) => {
189                 /* Nothing to do */
190             },
191             _ => compile_error(ty.span(), "Unsupported type"),
192         }
193     }
194 
substitute_type_param_bound(&self, bound: &mut TypeParamBound, traitname: &Ident)195     fn substitute_type_param_bound(&self,
196                                    bound: &mut TypeParamBound,
197                                    traitname: &Ident)
198     {
199         if let TypeParamBound::Trait(t) = bound {
200             match self.get_path(&t.path) {
201                 None => {
202                     for seg in t.path.segments.iter_mut() {
203                         self.substitute_path_segment(seg, traitname);
204                     }
205                 },
206                 Some(Type::Path(type_path)) => {
207                     t.path = type_path.path;
208                 },
209                 Some(_) => {
210                     compile_error(t.path.span(),
211                         "Can only substitute paths for trait bounds");
212                 }
213             }
214         }
215     }
216 
substitute_trait(&self, item: &ItemTrait) -> ItemTrait217     pub(crate) fn substitute_trait(&self, item: &ItemTrait) -> ItemTrait {
218         let mut output = item.clone();
219         for trait_item in output.items.iter_mut() {
220             match trait_item {
221                 TraitItem::Type(tity) => {
222                     if let Some(ty) = self.attrs.get(&tity.ident) {
223                         let span = tity.span();
224                         tity.default = Some((Token![=](span), ty.clone()));
225                         // Concrete associated types aren't allowed to have
226                         // bounds
227                         tity.bounds = Punctuated::new();
228                     } else {
229                         compile_error(tity.span(),
230                             "Default value not given for associated type");
231                     }
232                 },
233                 TraitItem::Fn(method) => {
234                     let sig = &mut method.sig;
235                     for fn_arg in sig.inputs.iter_mut() {
236                         if let FnArg::Typed(arg) = fn_arg {
237                             self.substitute_type(&mut arg.ty, &item.ident);
238                         }
239                     }
240                     if let ReturnType::Type(_, ref mut ty) = &mut sig.output {
241                         self.substitute_type(ty, &item.ident);
242                     }
243                 },
244                 _ => {
245                     // Nothing to do
246                 }
247             }
248         }
249         output
250     }
251 }
252 
253 impl Parse for Attrs {
parse(input: ParseStream) -> parse::Result<Self>254     fn parse(input: ParseStream) -> parse::Result<Self> {
255         let mut attrs = HashMap::new();
256         while !input.is_empty() {
257             let attr: Attr = input.parse()?;
258             match attr {
259                 Attr::Type(trait_item_type) => {
260                     let ident = trait_item_type.ident.clone();
261                     if let Some((_, ty)) = trait_item_type.default {
262                         attrs.insert(ident, ty.clone());
263                     } else {
264                         compile_error(trait_item_type.span(),
265                           "automock type attributes must have a default value");
266                     }
267                 }
268             }
269         }
270         Ok(Attrs{attrs})
271     }
272 }
273 
274 /// Unit tests for `Attrs`.
275 #[cfg(test)]
276 mod t {
277     use super::super::*;
278 
check_substitute_type( attrs: TokenStream, input: TokenStream, traitname: Ident, expected: TokenStream)279     fn check_substitute_type(
280         attrs: TokenStream,
281         input: TokenStream,
282         traitname: Ident,
283         expected: TokenStream)
284     {
285         let _self: super::Attrs = parse2(attrs).unwrap();
286         let mut in_ty: Type = parse2(input).unwrap();
287         let expect_ty: Type = parse2(expected).unwrap();
288         _self.substitute_type(&mut in_ty, &traitname);
289         assert_eq!(in_ty, expect_ty);
290     }
291 
292     #[test]
qself()293     fn qself() {
294         check_substitute_type(quote!(type T = u32;),
295                               quote!(<Self as Foo>::T),
296                               format_ident!("Foo"),
297                               quote!(u32));
298     }
299 
300     #[test]
301     #[should_panic(expected = "Mockall does not support QSelf substitutions except for the trait being mocked")]
qself_other()302     fn qself_other() {
303         check_substitute_type(quote!(type T = u32;),
304                               quote!(<Self as AsRef>::T),
305                               format_ident!("Foo"),
306                               quote!(u32));
307     }
308 
309     #[test]
310     #[should_panic(expected = "Unknown type substitution for QSelf")]
unknown_substitution()311     fn unknown_substitution() {
312         check_substitute_type(quote!(type T = u32;),
313                               quote!(<Self as Foo>::Q),
314                               format_ident!("Foo"),
315                               quote!(u32));
316     }
317 }
318