1use crate::internals::respan::respan;
2use proc_macro2::Span;
3use quote::ToTokens;
4use std::mem;
5use syn::punctuated::Punctuated;
6use syn::{
7    parse_quote, Data, DeriveInput, Expr, ExprPath, GenericArgument, GenericParam, Generics, Macro,
8    Path, PathArguments, QSelf, ReturnType, Token, Type, TypeParamBound, TypePath, WherePredicate,
9};
10
11pub fn replace_receiver(input: &mut DeriveInput) {
12    let self_ty = {
13        let ident = &input.ident;
14        let ty_generics = input.generics.split_for_impl().1;
15        parse_quote!(#ident #ty_generics)
16    };
17    let mut visitor = ReplaceReceiver(&self_ty);
18    visitor.visit_generics_mut(&mut input.generics);
19    visitor.visit_data_mut(&mut input.data);
20}
21
22struct ReplaceReceiver<'a>(&'a TypePath);
23
24impl ReplaceReceiver<'_> {
25    fn self_ty(&self, span: Span) -> TypePath {
26        let tokens = self.0.to_token_stream();
27        let respanned = respan(tokens, span);
28        syn::parse2(respanned).unwrap()
29    }
30
31    fn self_to_qself(&self, qself: &mut Option<QSelf>, path: &mut Path) {
32        if path.leading_colon.is_some() || path.segments[0].ident != "Self" {
33            return;
34        }
35
36        if path.segments.len() == 1 {
37            self.self_to_expr_path(path);
38            return;
39        }
40
41        let span = path.segments[0].ident.span();
42        *qself = Some(QSelf {
43            lt_token: Token![<](span),
44            ty: Box::new(Type::Path(self.self_ty(span))),
45            position: 0,
46            as_token: None,
47            gt_token: Token![>](span),
48        });
49
50        path.leading_colon = Some(**path.segments.pairs().next().unwrap().punct().unwrap());
51
52        let segments = mem::replace(&mut path.segments, Punctuated::new());
53        path.segments = segments.into_pairs().skip(1).collect();
54    }
55
56    fn self_to_expr_path(&self, path: &mut Path) {
57        let self_ty = self.self_ty(path.segments[0].ident.span());
58        let variant = mem::replace(path, self_ty.path);
59        for segment in &mut path.segments {
60            if let PathArguments::AngleBracketed(bracketed) = &mut segment.arguments {
61                if bracketed.colon2_token.is_none() && !bracketed.args.is_empty() {
62                    bracketed.colon2_token = Some(<Token![::]>::default());
63                }
64            }
65        }
66        if variant.segments.len() > 1 {
67            path.segments.push_punct(<Token![::]>::default());
68            path.segments.extend(variant.segments.into_pairs().skip(1));
69        }
70    }
71}
72
73impl ReplaceReceiver<'_> {
74    // `Self` -> `Receiver`
75    fn visit_type_mut(&mut self, ty: &mut Type) {
76        let span = if let Type::Path(node) = ty {
77            if node.qself.is_none() && node.path.is_ident("Self") {
78                node.path.segments[0].ident.span()
79            } else {
80                self.visit_type_path_mut(node);
81                return;
82            }
83        } else {
84            self.visit_type_mut_impl(ty);
85            return;
86        };
87        *ty = self.self_ty(span).into();
88    }
89
90    // `Self::Assoc` -> `<Receiver>::Assoc`
91    fn visit_type_path_mut(&mut self, ty: &mut TypePath) {
92        if ty.qself.is_none() {
93            self.self_to_qself(&mut ty.qself, &mut ty.path);
94        }
95        self.visit_type_path_mut_impl(ty);
96    }
97
98    // `Self::method` -> `<Receiver>::method`
99    fn visit_expr_path_mut(&mut self, expr: &mut ExprPath) {
100        if expr.qself.is_none() {
101            self.self_to_qself(&mut expr.qself, &mut expr.path);
102        }
103        self.visit_expr_path_mut_impl(expr);
104    }
105
106    // Everything below is simply traversing the syntax tree.
107
108    fn visit_type_mut_impl(&mut self, ty: &mut Type) {
109        match ty {
110            #![cfg_attr(all(test, exhaustive), deny(non_exhaustive_omitted_patterns))]
111            Type::Array(ty) => {
112                self.visit_type_mut(&mut ty.elem);
113                self.visit_expr_mut(&mut ty.len);
114            }
115            Type::BareFn(ty) => {
116                for arg in &mut ty.inputs {
117                    self.visit_type_mut(&mut arg.ty);
118                }
119                self.visit_return_type_mut(&mut ty.output);
120            }
121            Type::Group(ty) => self.visit_type_mut(&mut ty.elem),
122            Type::ImplTrait(ty) => {
123                for bound in &mut ty.bounds {
124                    self.visit_type_param_bound_mut(bound);
125                }
126            }
127            Type::Macro(ty) => self.visit_macro_mut(&mut ty.mac),
128            Type::Paren(ty) => self.visit_type_mut(&mut ty.elem),
129            Type::Path(ty) => {
130                if let Some(qself) = &mut ty.qself {
131                    self.visit_type_mut(&mut qself.ty);
132                }
133                self.visit_path_mut(&mut ty.path);
134            }
135            Type::Ptr(ty) => self.visit_type_mut(&mut ty.elem),
136            Type::Reference(ty) => self.visit_type_mut(&mut ty.elem),
137            Type::Slice(ty) => self.visit_type_mut(&mut ty.elem),
138            Type::TraitObject(ty) => {
139                for bound in &mut ty.bounds {
140                    self.visit_type_param_bound_mut(bound);
141                }
142            }
143            Type::Tuple(ty) => {
144                for elem in &mut ty.elems {
145                    self.visit_type_mut(elem);
146                }
147            }
148
149            Type::Infer(_) | Type::Never(_) | Type::Verbatim(_) => {}
150
151            _ => {}
152        }
153    }
154
155    fn visit_type_path_mut_impl(&mut self, ty: &mut TypePath) {
156        if let Some(qself) = &mut ty.qself {
157            self.visit_type_mut(&mut qself.ty);
158        }
159        self.visit_path_mut(&mut ty.path);
160    }
161
162    fn visit_expr_path_mut_impl(&mut self, expr: &mut ExprPath) {
163        if let Some(qself) = &mut expr.qself {
164            self.visit_type_mut(&mut qself.ty);
165        }
166        self.visit_path_mut(&mut expr.path);
167    }
168
169    fn visit_path_mut(&mut self, path: &mut Path) {
170        for segment in &mut path.segments {
171            self.visit_path_arguments_mut(&mut segment.arguments);
172        }
173    }
174
175    fn visit_path_arguments_mut(&mut self, arguments: &mut PathArguments) {
176        match arguments {
177            PathArguments::None => {}
178            PathArguments::AngleBracketed(arguments) => {
179                for arg in &mut arguments.args {
180                    match arg {
181                        #![cfg_attr(all(test, exhaustive), deny(non_exhaustive_omitted_patterns))]
182                        GenericArgument::Type(arg) => self.visit_type_mut(arg),
183                        GenericArgument::AssocType(arg) => self.visit_type_mut(&mut arg.ty),
184                        GenericArgument::Lifetime(_)
185                        | GenericArgument::Const(_)
186                        | GenericArgument::AssocConst(_)
187                        | GenericArgument::Constraint(_) => {}
188                        _ => {}
189                    }
190                }
191            }
192            PathArguments::Parenthesized(arguments) => {
193                for argument in &mut arguments.inputs {
194                    self.visit_type_mut(argument);
195                }
196                self.visit_return_type_mut(&mut arguments.output);
197            }
198        }
199    }
200
201    fn visit_return_type_mut(&mut self, return_type: &mut ReturnType) {
202        match return_type {
203            ReturnType::Default => {}
204            ReturnType::Type(_, output) => self.visit_type_mut(output),
205        }
206    }
207
208    fn visit_type_param_bound_mut(&mut self, bound: &mut TypeParamBound) {
209        match bound {
210            #![cfg_attr(all(test, exhaustive), deny(non_exhaustive_omitted_patterns))]
211            TypeParamBound::Trait(bound) => self.visit_path_mut(&mut bound.path),
212            TypeParamBound::Lifetime(_) | TypeParamBound::Verbatim(_) => {}
213            _ => {}
214        }
215    }
216
217    fn visit_generics_mut(&mut self, generics: &mut Generics) {
218        for param in &mut generics.params {
219            match param {
220                GenericParam::Type(param) => {
221                    for bound in &mut param.bounds {
222                        self.visit_type_param_bound_mut(bound);
223                    }
224                }
225                GenericParam::Lifetime(_) | GenericParam::Const(_) => {}
226            }
227        }
228        if let Some(where_clause) = &mut generics.where_clause {
229            for predicate in &mut where_clause.predicates {
230                match predicate {
231                    #![cfg_attr(all(test, exhaustive), deny(non_exhaustive_omitted_patterns))]
232                    WherePredicate::Type(predicate) => {
233                        self.visit_type_mut(&mut predicate.bounded_ty);
234                        for bound in &mut predicate.bounds {
235                            self.visit_type_param_bound_mut(bound);
236                        }
237                    }
238                    WherePredicate::Lifetime(_) => {}
239                    _ => {}
240                }
241            }
242        }
243    }
244
245    fn visit_data_mut(&mut self, data: &mut Data) {
246        match data {
247            Data::Struct(data) => {
248                for field in &mut data.fields {
249                    self.visit_type_mut(&mut field.ty);
250                }
251            }
252            Data::Enum(data) => {
253                for variant in &mut data.variants {
254                    for field in &mut variant.fields {
255                        self.visit_type_mut(&mut field.ty);
256                    }
257                }
258            }
259            Data::Union(_) => {}
260        }
261    }
262
263    fn visit_expr_mut(&mut self, expr: &mut Expr) {
264        match expr {
265            Expr::Binary(expr) => {
266                self.visit_expr_mut(&mut expr.left);
267                self.visit_expr_mut(&mut expr.right);
268            }
269            Expr::Call(expr) => {
270                self.visit_expr_mut(&mut expr.func);
271                for arg in &mut expr.args {
272                    self.visit_expr_mut(arg);
273                }
274            }
275            Expr::Cast(expr) => {
276                self.visit_expr_mut(&mut expr.expr);
277                self.visit_type_mut(&mut expr.ty);
278            }
279            Expr::Field(expr) => self.visit_expr_mut(&mut expr.base),
280            Expr::Index(expr) => {
281                self.visit_expr_mut(&mut expr.expr);
282                self.visit_expr_mut(&mut expr.index);
283            }
284            Expr::Paren(expr) => self.visit_expr_mut(&mut expr.expr),
285            Expr::Path(expr) => self.visit_expr_path_mut(expr),
286            Expr::Unary(expr) => self.visit_expr_mut(&mut expr.expr),
287            _ => {}
288        }
289    }
290
291    fn visit_macro_mut(&mut self, _mac: &mut Macro) {}
292}
293