@@ -48,36 +48,42 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
4848 let impl_def = ctx. find_node_at_offset :: < ast:: Impl > ( ) ?. clone_for_update ( ) ;
4949 let indent = impl_def. indent_level ( ) ;
5050
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+
5157 let trait_ = impl_def. trait_ ( ) ?;
5258 if let ast:: Type :: PathType ( trait_path) = trait_ {
5359 let trait_type = ctx. sema . resolve_trait ( & trait_path. path ( ) ?) ?;
5460 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) ? {
5663 return None ;
5764 }
5865 }
5966
6067 // Index -> IndexMut
61- let index_trait = impl_def
62- . syntax ( )
63- . descendants ( )
64- . filter_map ( ast:: NameRef :: cast)
65- . find ( |it| it. text ( ) == "Index" ) ?;
6668 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 ( ) ,
6971 ) ;
7072
7173 // index -> index_mut
72- let trait_method_name = impl_def
74+ let ( trait_method_name, new_trait_method_name ) = impl_def
7375 . syntax ( )
7476 . descendants ( )
7577 . 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+ ) ;
7883
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+ }
8187
8288 // &self -> &mut self
8389 let mut_self_param = make:: mut_self_param ( ) ;
@@ -87,10 +93,8 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
8793
8894 // &Self::Output -> &mut Self::Output
8995 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 ( ) ) ;
9498
9599 let fn_ = impl_def. assoc_item_list ( ) ?. assoc_items ( ) . find_map ( |it| match it {
96100 ast:: AssocItem :: Fn ( f) => Some ( f) ,
@@ -104,14 +108,51 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
104108 let target = impl_def. syntax ( ) . text_range ( ) ;
105109 acc. add (
106110 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" ) ,
108112 target,
109113 |edit| {
110114 edit. insert ( target. start ( ) , format ! ( "$0{impl_def}\n \n {indent}" ) ) ;
111115 } ,
112116 )
113117}
114118
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+
115156#[ cfg( test) ]
116157mod tests {
117158 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 {
186227 var_name
187228 }
188229}
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+ }
189259"# ,
190260 ) ;
191261 }
@@ -285,6 +355,14 @@ mod foo {
285355pub trait Index<Idx: ?Sized> {}
286356
287357impl<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] {}
288366"# ,
289367 ) ;
290368 }
0 commit comments