async_trait/
receiver.rs

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

async_trait/
receiver.rs
