1 use proc_macro2::{Group, Span, TokenStream, TokenTree};
2 use std::iter::FromIterator;
3 use syn::visit_mut::{self, VisitMut};
4 use syn::{
5     Block, ExprPath, Ident, Item, Macro, Pat, PatIdent, PatPath, Path, Receiver, Signature, Token,
6     TypePath,
7 };
8 
has_self_in_sig(sig: &mut Signature) -> bool9 pub fn has_self_in_sig(sig: &mut Signature) -> bool {
10     let mut visitor = HasSelf(false);
11     visitor.visit_signature_mut(sig);
12     visitor.0
13 }
14 
has_self_in_block(block: &mut Block) -> bool15 pub fn has_self_in_block(block: &mut Block) -> bool {
16     let mut visitor = HasSelf(false);
17     visitor.visit_block_mut(block);
18     visitor.0
19 }
20 
has_self_in_token_stream(tokens: TokenStream) -> bool21 fn has_self_in_token_stream(tokens: TokenStream) -> bool {
22     tokens.into_iter().any(|tt| match tt {
23         TokenTree::Ident(ident) => ident == "Self",
24         TokenTree::Group(group) => has_self_in_token_stream(group.stream()),
25         _ => false,
26     })
27 }
28 
mut_pat(pat: &mut Pat) -> Option<Token![mut]>29 pub fn mut_pat(pat: &mut Pat) -> Option<Token![mut]> {
30     let mut visitor = HasMutPat(None);
31     visitor.visit_pat_mut(pat);
32     visitor.0
33 }
34 
contains_fn(tokens: TokenStream) -> bool35 fn contains_fn(tokens: TokenStream) -> bool {
36     tokens.into_iter().any(|tt| match tt {
37         TokenTree::Ident(ident) => ident == "fn",
38         TokenTree::Group(group) => contains_fn(group.stream()),
39         _ => false,
40     })
41 }
42 
43 struct HasMutPat(Option<Token![mut]>);
44 
45 impl VisitMut for HasMutPat {
visit_pat_ident_mut(&mut self, i: &mut PatIdent)46     fn visit_pat_ident_mut(&mut self, i: &mut PatIdent) {
47         if let Some(m) = i.mutability {
48             self.0 = Some(m);
49         } else {
50             visit_mut::visit_pat_ident_mut(self, i);
51         }
52     }
53 }
54 
55 struct HasSelf(bool);
56 
57 impl VisitMut for HasSelf {
visit_expr_path_mut(&mut self, expr: &mut ExprPath)58     fn visit_expr_path_mut(&mut self, expr: &mut ExprPath) {
59         self.0 |= expr.path.segments[0].ident == "Self";
60         visit_mut::visit_expr_path_mut(self, expr);
61     }
62 
visit_pat_path_mut(&mut self, pat: &mut PatPath)63     fn visit_pat_path_mut(&mut self, pat: &mut PatPath) {
64         self.0 |= pat.path.segments[0].ident == "Self";
65         visit_mut::visit_pat_path_mut(self, pat);
66     }
67 
visit_type_path_mut(&mut self, ty: &mut TypePath)68     fn visit_type_path_mut(&mut self, ty: &mut TypePath) {
69         self.0 |= ty.path.segments[0].ident == "Self";
70         visit_mut::visit_type_path_mut(self, ty);
71     }
72 
visit_receiver_mut(&mut self, _arg: &mut Receiver)73     fn visit_receiver_mut(&mut self, _arg: &mut Receiver) {
74         self.0 = true;
75     }
76 
visit_item_mut(&mut self, _: &mut Item)77     fn visit_item_mut(&mut self, _: &mut Item) {
78         // Do not recurse into nested items.
79     }
80 
visit_macro_mut(&mut self, mac: &mut Macro)81     fn visit_macro_mut(&mut self, mac: &mut Macro) {
82         if !contains_fn(mac.tokens.clone()) {
83             self.0 |= has_self_in_token_stream(mac.tokens.clone());
84         }
85     }
86 }
87 
88 pub struct ReplaceSelf(pub Span);
89 
90 impl ReplaceSelf {
91     #[cfg_attr(not(self_span_hack), allow(clippy::unused_self))]
prepend_underscore_to_self(&self, ident: &mut Ident) -> bool92     fn prepend_underscore_to_self(&self, ident: &mut Ident) -> bool {
93         let modified = ident == "self";
94         if modified {
95             *ident = Ident::new("__self", ident.span());
96             #[cfg(self_span_hack)]
97             ident.set_span(self.0);
98         }
99         modified
100     }
101 
visit_token_stream(&mut self, tokens: &mut TokenStream) -> bool102     fn visit_token_stream(&mut self, tokens: &mut TokenStream) -> bool {
103         let mut out = Vec::new();
104         let mut modified = false;
105         visit_token_stream_impl(self, tokens.clone(), &mut modified, &mut out);
106         if modified {
107             *tokens = TokenStream::from_iter(out);
108         }
109         return modified;
110 
111         fn visit_token_stream_impl(
112             visitor: &mut ReplaceSelf,
113             tokens: TokenStream,
114             modified: &mut bool,
115             out: &mut Vec<TokenTree>,
116         ) {
117             for tt in tokens {
118                 match tt {
119                     TokenTree::Ident(mut ident) => {
120                         *modified |= visitor.prepend_underscore_to_self(&mut ident);
121                         out.push(TokenTree::Ident(ident));
122                     }
123                     TokenTree::Group(group) => {
124                         let mut content = group.stream();
125                         *modified |= visitor.visit_token_stream(&mut content);
126                         let mut new = Group::new(group.delimiter(), content);
127                         new.set_span(group.span());
128                         out.push(TokenTree::Group(new));
129                     }
130                     other => out.push(other),
131                 }
132             }
133         }
134     }
135 }
136 
137 impl VisitMut for ReplaceSelf {
visit_ident_mut(&mut self, i: &mut Ident)138     fn visit_ident_mut(&mut self, i: &mut Ident) {
139         self.prepend_underscore_to_self(i);
140     }
141 
visit_path_mut(&mut self, p: &mut Path)142     fn visit_path_mut(&mut self, p: &mut Path) {
143         if p.segments.len() == 1 {
144             // Replace `self`, but not `self::function`.
145             self.visit_ident_mut(&mut p.segments[0].ident);
146         }
147         for segment in &mut p.segments {
148             self.visit_path_arguments_mut(&mut segment.arguments);
149         }
150     }
151 
visit_item_mut(&mut self, i: &mut Item)152     fn visit_item_mut(&mut self, i: &mut Item) {
153         // Visit `macro_rules!` because locally defined macros can refer to
154         // `self`.
155         //
156         // Visit `futures::select` and similar select macros, which commonly
157         // appear syntactically like an item despite expanding to an expression.
158         //
159         // Otherwise, do not recurse into nested items.
160         if let Item::Macro(i) = i {
161             if i.mac.path.is_ident("macro_rules")
162                 || i.mac.path.segments.last().unwrap().ident == "select"
163             {
164                 self.visit_macro_mut(&mut i.mac);
165             }
166         }
167     }
168 
visit_macro_mut(&mut self, mac: &mut Macro)169     fn visit_macro_mut(&mut self, mac: &mut Macro) {
170         // We can't tell in general whether `self` inside a macro invocation
171         // refers to the self in the argument list or a different self
172         // introduced within the macro. Heuristic: if the macro input contains
173         // `fn`, then `self` is more likely to refer to something other than the
174         // outer function's self argument.
175         if !contains_fn(mac.tokens.clone()) {
176             self.visit_token_stream(&mut mac.tokens);
177         }
178     }
179 }
180