1+ use ide_db:: imports_locator;
12use itertools:: Itertools ;
23use syntax:: {
3- ast:: { self , AstNode } ,
4+ ast:: { self , make , AstNode } ,
45 Direction , SmolStr ,
56 SyntaxKind :: { IDENT , WHITESPACE } ,
67 TextRange , TextSize ,
78} ;
89
910use crate :: {
10- assist_context:: { AssistContext , Assists } ,
11+ assist_config:: SnippetCap ,
12+ assist_context:: { AssistBuilder , AssistContext , Assists } ,
13+ utils:: mod_path_to_ast,
1114 AssistId , AssistKind ,
1215} ;
1316
@@ -30,78 +33,151 @@ use crate::{
3033// ```
3134pub ( crate ) fn add_custom_impl ( acc : & mut Assists , ctx : & AssistContext ) -> Option < ( ) > {
3235 let attr = ctx. find_node_at_offset :: < ast:: Attr > ( ) ?;
33- let input = attr. token_tree ( ) ?;
3436
3537 let attr_name = attr
3638 . syntax ( )
3739 . descendants_with_tokens ( )
3840 . filter ( |t| t. kind ( ) == IDENT )
39- . find_map ( |i| i . into_token ( ) )
40- . filter ( |t| * t. text ( ) == "derive" ) ?
41+ . find_map ( syntax :: NodeOrToken :: into_token)
42+ . filter ( |t| t. text ( ) == "derive" ) ?
4143 . text ( )
4244 . clone ( ) ;
4345
4446 let trait_token =
4547 ctx. token_at_offset ( ) . find ( |t| t. kind ( ) == IDENT && * t. text ( ) != attr_name) ?;
48+ let trait_path = make:: path_unqualified ( make:: path_segment ( make:: name_ref ( trait_token. text ( ) ) ) ) ;
4649
4750 let annotated = attr. syntax ( ) . siblings ( Direction :: Next ) . find_map ( ast:: Name :: cast) ?;
4851 let annotated_name = annotated. syntax ( ) . text ( ) . to_string ( ) ;
49- let start_offset = annotated. syntax ( ) . parent ( ) ?. text_range ( ) . end ( ) ;
52+ let insert_pos = annotated. syntax ( ) . parent ( ) ?. text_range ( ) . end ( ) ;
53+
54+ let current_module = ctx. sema . scope ( annotated. syntax ( ) ) . module ( ) ?;
55+ let current_crate = current_module. krate ( ) ;
56+
57+ let found_traits = imports_locator:: find_imports ( & ctx. sema , current_crate, trait_token. text ( ) )
58+ . into_iter ( )
59+ . filter_map ( |candidate : either:: Either < hir:: ModuleDef , hir:: MacroDef > | match candidate {
60+ either:: Either :: Left ( hir:: ModuleDef :: Trait ( trait_) ) => Some ( trait_) ,
61+ _ => None ,
62+ } )
63+ . flat_map ( |trait_| {
64+ current_module
65+ . find_use_path ( ctx. sema . db , hir:: ModuleDef :: Trait ( trait_) )
66+ . as_ref ( )
67+ . map ( mod_path_to_ast)
68+ . zip ( Some ( trait_) )
69+ } ) ;
5070
51- let label =
52- format ! ( "Add custom impl `{}` for `{}`" , trait_token. text( ) . as_str( ) , annotated_name) ;
71+ let mut no_traits_found = true ;
72+ for ( trait_path, _trait) in found_traits. inspect ( |_| no_traits_found = false ) {
73+ add_assist ( acc, ctx. config . snippet_cap , & attr, & trait_path, & annotated_name, insert_pos) ?;
74+ }
75+ if no_traits_found {
76+ add_assist ( acc, ctx. config . snippet_cap , & attr, & trait_path, & annotated_name, insert_pos) ?;
77+ }
78+ Some ( ( ) )
79+ }
5380
81+ fn add_assist (
82+ acc : & mut Assists ,
83+ snippet_cap : Option < SnippetCap > ,
84+ attr : & ast:: Attr ,
85+ trait_path : & ast:: Path ,
86+ annotated_name : & str ,
87+ insert_pos : TextSize ,
88+ ) -> Option < ( ) > {
5489 let target = attr. syntax ( ) . text_range ( ) ;
55- acc. add ( AssistId ( "add_custom_impl" , AssistKind :: Refactor ) , label, target, |builder| {
56- let new_attr_input = input
57- . syntax ( )
58- . descendants_with_tokens ( )
59- . filter ( |t| t. kind ( ) == IDENT )
60- . filter_map ( |t| t. into_token ( ) . map ( |t| t. text ( ) . clone ( ) ) )
61- . filter ( |t| t != trait_token. text ( ) )
62- . collect :: < Vec < SmolStr > > ( ) ;
63- let has_more_derives = !new_attr_input. is_empty ( ) ;
64-
65- if has_more_derives {
66- let new_attr_input = format ! ( "({})" , new_attr_input. iter( ) . format( ", " ) ) ;
67- builder. replace ( input. syntax ( ) . text_range ( ) , new_attr_input) ;
68- } else {
69- let attr_range = attr. syntax ( ) . text_range ( ) ;
70- builder. delete ( attr_range) ;
71-
72- let line_break_range = attr
73- . syntax ( )
74- . next_sibling_or_token ( )
75- . filter ( |t| t. kind ( ) == WHITESPACE )
76- . map ( |t| t. text_range ( ) )
77- . unwrap_or_else ( || TextRange :: new ( TextSize :: from ( 0 ) , TextSize :: from ( 0 ) ) ) ;
78- builder. delete ( line_break_range) ;
79- }
90+ let input = attr. token_tree ( ) ?;
91+ let label = format ! ( "Add custom impl `{}` for `{}`" , trait_path, annotated_name) ;
92+ let trait_name = trait_path. segment ( ) . and_then ( |seg| seg. name_ref ( ) ) ?;
8093
81- match ctx. config . snippet_cap {
94+ acc. add ( AssistId ( "add_custom_impl" , AssistKind :: Refactor ) , label, target, |builder| {
95+ update_attribute ( builder, & input, & trait_name, & attr) ;
96+ match snippet_cap {
8297 Some ( cap) => {
8398 builder. insert_snippet (
8499 cap,
85- start_offset ,
86- format ! ( "\n \n impl {} for {} {{\n $0\n }}" , trait_token , annotated_name) ,
100+ insert_pos ,
101+ format ! ( "\n \n impl {} for {} {{\n $0\n }}" , trait_path , annotated_name) ,
87102 ) ;
88103 }
89104 None => {
90105 builder. insert (
91- start_offset ,
92- format ! ( "\n \n impl {} for {} {{\n \n }}" , trait_token , annotated_name) ,
106+ insert_pos ,
107+ format ! ( "\n \n impl {} for {} {{\n \n }}" , trait_path , annotated_name) ,
93108 ) ;
94109 }
95110 }
96111 } )
97112}
98113
114+ fn update_attribute (
115+ builder : & mut AssistBuilder ,
116+ input : & ast:: TokenTree ,
117+ trait_name : & ast:: NameRef ,
118+ attr : & ast:: Attr ,
119+ ) {
120+ let new_attr_input = input
121+ . syntax ( )
122+ . descendants_with_tokens ( )
123+ . filter ( |t| t. kind ( ) == IDENT )
124+ . filter_map ( |t| t. into_token ( ) . map ( |t| t. text ( ) . clone ( ) ) )
125+ . filter ( |t| t != trait_name. text ( ) )
126+ . collect :: < Vec < SmolStr > > ( ) ;
127+ let has_more_derives = !new_attr_input. is_empty ( ) ;
128+
129+ if has_more_derives {
130+ let new_attr_input = format ! ( "({})" , new_attr_input. iter( ) . format( ", " ) ) ;
131+ builder. replace ( input. syntax ( ) . text_range ( ) , new_attr_input) ;
132+ } else {
133+ let attr_range = attr. syntax ( ) . text_range ( ) ;
134+ builder. delete ( attr_range) ;
135+
136+ let line_break_range = attr
137+ . syntax ( )
138+ . next_sibling_or_token ( )
139+ . filter ( |t| t. kind ( ) == WHITESPACE )
140+ . map ( |t| t. text_range ( ) )
141+ . unwrap_or_else ( || TextRange :: new ( TextSize :: from ( 0 ) , TextSize :: from ( 0 ) ) ) ;
142+ builder. delete ( line_break_range) ;
143+ }
144+ }
145+
99146#[ cfg( test) ]
100147mod tests {
101148 use crate :: tests:: { check_assist, check_assist_not_applicable} ;
102149
103150 use super :: * ;
104151
152+ #[ test]
153+ fn add_custom_impl_qualified ( ) {
154+ check_assist (
155+ add_custom_impl,
156+ "
157+ mod fmt {
158+ pub trait Debug {}
159+ }
160+
161+ #[derive(Debu<|>g)]
162+ struct Foo {
163+ bar: String,
164+ }
165+ " ,
166+ "
167+ mod fmt {
168+ pub trait Debug {}
169+ }
170+
171+ struct Foo {
172+ bar: String,
173+ }
174+
175+ impl fmt::Debug for Foo {
176+ $0
177+ }
178+ " ,
179+ )
180+ }
105181 #[ test]
106182 fn add_custom_impl_for_unique_input ( ) {
107183 check_assist (
0 commit comments