@@ -48,36 +48,42 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
48
48
let impl_def = ctx. find_node_at_offset :: < ast:: Impl > ( ) ?. clone_for_update ( ) ;
49
49
let indent = impl_def. indent_level ( ) ;
50
50
51
+ let ( apply_trait, new_apply_trait) = impl_def
52
+ . syntax ( )
53
+ . descendants ( )
54
+ . filter_map ( ast:: NameRef :: cast)
55
+ . find_map ( process_trait_name) ?;
56
+
51
57
let trait_ = impl_def. trait_ ( ) ?;
52
58
if let ast:: Type :: PathType ( trait_path) = trait_ {
53
59
let trait_type = ctx. sema . resolve_trait ( & trait_path. path ( ) ?) ?;
54
60
let scope = ctx. sema . scope ( trait_path. syntax ( ) ) ?;
55
- if trait_type != FamousDefs ( & ctx. sema , scope. krate ( ) ) . core_convert_Index ( ) ? {
61
+ let famous_defs = FamousDefs ( & ctx. sema , scope. krate ( ) ) ;
62
+ if trait_type != get_famous ( & apply_trait. text ( ) , famous_defs) ? {
56
63
return None ;
57
64
}
58
65
}
59
66
60
67
// Index -> IndexMut
61
- let index_trait = impl_def
62
- . syntax ( )
63
- . descendants ( )
64
- . filter_map ( ast:: NameRef :: cast)
65
- . find ( |it| it. text ( ) == "Index" ) ?;
66
68
ted:: replace (
67
- index_trait . syntax ( ) ,
68
- make:: path_segment ( make:: name_ref ( "IndexMut" ) ) . clone_for_update ( ) . syntax ( ) ,
69
+ apply_trait . syntax ( ) ,
70
+ make:: path_segment ( make:: name_ref ( new_apply_trait ) ) . clone_for_update ( ) . syntax ( ) ,
69
71
) ;
70
72
71
73
// index -> index_mut
72
- let trait_method_name = impl_def
74
+ let ( trait_method_name, new_trait_method_name ) = impl_def
73
75
. syntax ( )
74
76
. descendants ( )
75
77
. filter_map ( ast:: Name :: cast)
76
- . find ( |it| it. text ( ) == "index" ) ?;
77
- ted:: replace ( trait_method_name. syntax ( ) , make:: name ( "index_mut" ) . clone_for_update ( ) . syntax ( ) ) ;
78
+ . find_map ( process_method_name) ?;
79
+ ted:: replace (
80
+ trait_method_name. syntax ( ) ,
81
+ make:: name ( new_trait_method_name) . clone_for_update ( ) . syntax ( ) ,
82
+ ) ;
78
83
79
- let type_alias = impl_def. syntax ( ) . descendants ( ) . find_map ( ast:: TypeAlias :: cast) ?;
80
- ted:: remove ( type_alias. syntax ( ) ) ;
84
+ if let Some ( type_alias) = impl_def. syntax ( ) . descendants ( ) . find_map ( ast:: TypeAlias :: cast) {
85
+ ted:: remove ( type_alias. syntax ( ) ) ;
86
+ }
81
87
82
88
// &self -> &mut self
83
89
let mut_self_param = make:: mut_self_param ( ) ;
@@ -87,10 +93,8 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
87
93
88
94
// &Self::Output -> &mut Self::Output
89
95
let ret_type = impl_def. syntax ( ) . descendants ( ) . find_map ( ast:: RetType :: cast) ?;
90
- ted:: replace (
91
- ret_type. syntax ( ) ,
92
- make:: ret_type ( make:: ty ( "&mut Self::Output" ) ) . clone_for_update ( ) . syntax ( ) ,
93
- ) ;
96
+ let new_ret_type = process_ret_type ( & ret_type) ?;
97
+ ted:: replace ( ret_type. syntax ( ) , make:: ret_type ( new_ret_type) . clone_for_update ( ) . syntax ( ) ) ;
94
98
95
99
let fn_ = impl_def. assoc_item_list ( ) ?. assoc_items ( ) . find_map ( |it| match it {
96
100
ast:: AssocItem :: Fn ( f) => Some ( f) ,
@@ -104,14 +108,51 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
104
108
let target = impl_def. syntax ( ) . text_range ( ) ;
105
109
acc. add (
106
110
AssistId :: generate ( "generate_mut_trait_impl" ) ,
107
- "Generate `IndexMut ` impl from this `Index ` trait" ,
111
+ format ! ( "Generate `{new_apply_trait} ` impl from this `{apply_trait} ` trait" ) ,
108
112
target,
109
113
|edit| {
110
114
edit. insert ( target. start ( ) , format ! ( "$0{impl_def}\n \n {indent}" ) ) ;
111
115
} ,
112
116
)
113
117
}
114
118
119
+ fn get_famous ( apply_trait : & str , famous : FamousDefs < ' _ , ' _ > ) -> Option < hir:: Trait > {
120
+ match apply_trait {
121
+ "Index" => famous. core_convert_Index ( ) ,
122
+ "AsRef" => famous. core_convert_AsRef ( ) ,
123
+ "Borrow" => famous. core_borrow_Borrow ( ) ,
124
+ _ => None ,
125
+ }
126
+ }
127
+
128
+ fn process_trait_name ( name : ast:: NameRef ) -> Option < ( ast:: NameRef , & ' static str ) > {
129
+ let new_name = match & * name. text ( ) {
130
+ "Index" => "IndexMut" ,
131
+ "AsRef" => "AsMut" ,
132
+ "Borrow" => "BorrowMut" ,
133
+ _ => return None ,
134
+ } ;
135
+ Some ( ( name, new_name) )
136
+ }
137
+
138
+ fn process_method_name ( name : ast:: Name ) -> Option < ( ast:: Name , & ' static str ) > {
139
+ let new_name = match & * name. text ( ) {
140
+ "index" => "index_mut" ,
141
+ "as_ref" => "as_mut" ,
142
+ "borrow" => "borrow_mut" ,
143
+ _ => return None ,
144
+ } ;
145
+ Some ( ( name, new_name) )
146
+ }
147
+
148
+ fn process_ret_type ( ref_ty : & ast:: RetType ) -> Option < ast:: Type > {
149
+ let ty = ref_ty. ty ( ) ?;
150
+ let ast:: Type :: RefType ( ref_type) = ty else {
151
+ return None ;
152
+ } ;
153
+ Some ( make:: ty_ref ( ref_type. ty ( ) ?, true ) )
154
+ }
155
+
115
156
#[ cfg( test) ]
116
157
mod tests {
117
158
use crate :: tests:: { check_assist, check_assist_not_applicable} ;
@@ -186,6 +227,35 @@ impl<T> core::ops::Index<Axis> for [T; 3] where T: Copy {
186
227
var_name
187
228
}
188
229
}
230
+ "# ,
231
+ ) ;
232
+
233
+ check_assist (
234
+ generate_mut_trait_impl,
235
+ r#"
236
+ //- minicore: as_ref
237
+ struct Foo(i32);
238
+
239
+ impl<T> core::convert::AsRef$0<i32> for Foo {
240
+ fn as_ref(&self) -> &i32 {
241
+ &self.0
242
+ }
243
+ }
244
+ "# ,
245
+ r#"
246
+ struct Foo(i32);
247
+
248
+ $0impl<T> core::convert::AsMut<i32> for Foo {
249
+ fn as_mut(&mut self) -> &mut i32 {
250
+ &self.0
251
+ }
252
+ }
253
+
254
+ impl<T> core::convert::AsRef<i32> for Foo {
255
+ fn as_ref(&self) -> &i32 {
256
+ &self.0
257
+ }
258
+ }
189
259
"# ,
190
260
) ;
191
261
}
@@ -285,6 +355,14 @@ mod foo {
285
355
pub trait Index<Idx: ?Sized> {}
286
356
287
357
impl<T> Index$0<i32> for [T; 3] {}
358
+ "# ,
359
+ ) ;
360
+ check_assist_not_applicable (
361
+ generate_mut_trait_impl,
362
+ r#"
363
+ pub trait AsRef<T: ?Sized> {}
364
+
365
+ impl<T> AsRef$0<i32> for [T; 3] {}
288
366
"# ,
289
367
) ;
290
368
}
0 commit comments