@@ -3,8 +3,8 @@ use ra_syntax::{
33 SyntaxKind , SyntaxNode , TextUnit ,
44} ;
55
6- use crate :: { Assist , AssistCtx , AssistId } ;
7- use ast:: { edit:: IndentLevel , ArgListOwner , CallExpr , Expr } ;
6+ use crate :: { Assist , AssistCtx , AssistFile , AssistId } ;
7+ use ast:: { edit:: IndentLevel , ArgListOwner , ModuleItemOwner } ;
88use hir:: HirDisplay ;
99use rustc_hash:: { FxHashMap , FxHashSet } ;
1010
@@ -16,7 +16,7 @@ use rustc_hash::{FxHashMap, FxHashSet};
1616// struct Baz;
1717// fn baz() -> Baz { Baz }
1818// fn foo() {
19- // bar<|>("", baz());
19+ // bar<|>("", baz());
2020// }
2121//
2222// ```
@@ -25,7 +25,7 @@ use rustc_hash::{FxHashMap, FxHashSet};
2525// struct Baz;
2626// fn baz() -> Baz { Baz }
2727// fn foo() {
28- // bar("", baz());
28+ // bar("", baz());
2929// }
3030//
3131// fn bar(arg: &str, baz: Baz) {
@@ -38,21 +38,30 @@ pub(crate) fn add_function(ctx: AssistCtx) -> Option<Assist> {
3838 let call = path_expr. syntax ( ) . parent ( ) . and_then ( ast:: CallExpr :: cast) ?;
3939 let path = path_expr. path ( ) ?;
4040
41- if path. qualifier ( ) . is_some ( ) {
42- return None ;
43- }
44-
4541 if ctx. sema . resolve_path ( & path) . is_some ( ) {
4642 // The function call already resolves, no need to add a function
4743 return None ;
4844 }
4945
50- let function_builder = FunctionBuilder :: from_call ( & ctx, & call) ?;
46+ let target_module = if let Some ( qualifier) = path. qualifier ( ) {
47+ if let Some ( hir:: PathResolution :: Def ( hir:: ModuleDef :: Module ( module) ) ) =
48+ ctx. sema . resolve_path ( & qualifier)
49+ {
50+ Some ( module. definition_source ( ctx. sema . db ) )
51+ } else {
52+ return None ;
53+ }
54+ } else {
55+ None
56+ } ;
57+
58+ let function_builder = FunctionBuilder :: from_call ( & ctx, & call, & path, target_module) ?;
5159
5260 ctx. add_assist ( AssistId ( "add_function" ) , "Add function" , |edit| {
5361 edit. target ( call. syntax ( ) . text_range ( ) ) ;
5462
5563 if let Some ( function_template) = function_builder. render ( ) {
64+ edit. set_file ( function_template. file ) ;
5665 edit. set_cursor ( function_template. cursor_offset ) ;
5766 edit. insert ( function_template. insert_offset , function_template. fn_def . to_string ( ) ) ;
5867 }
@@ -63,29 +72,67 @@ struct FunctionTemplate {
6372 insert_offset : TextUnit ,
6473 cursor_offset : TextUnit ,
6574 fn_def : ast:: SourceFile ,
75+ file : AssistFile ,
6676}
6777
6878struct FunctionBuilder {
69- append_fn_at : SyntaxNode ,
79+ target : GeneratedFunctionTarget ,
7080 fn_name : ast:: Name ,
7181 type_params : Option < ast:: TypeParamList > ,
7282 params : ast:: ParamList ,
83+ file : AssistFile ,
84+ needs_pub : bool ,
7385}
7486
7587impl FunctionBuilder {
76- fn from_call ( ctx : & AssistCtx , call : & ast:: CallExpr ) -> Option < Self > {
77- let append_fn_at = next_space_for_fn ( & call) ?;
78- let fn_name = fn_name ( & call) ?;
88+ /// Prepares a generated function that matches `call` in `generate_in`
89+ /// (or as close to `call` as possible, if `generate_in` is `None`)
90+ fn from_call (
91+ ctx : & AssistCtx ,
92+ call : & ast:: CallExpr ,
93+ path : & ast:: Path ,
94+ target_module : Option < hir:: InFile < hir:: ModuleSource > > ,
95+ ) -> Option < Self > {
96+ let needs_pub = target_module. is_some ( ) ;
97+ let mut file = AssistFile :: default ( ) ;
98+ let target = if let Some ( target_module) = target_module {
99+ let ( in_file, target) = next_space_for_fn_in_module ( ctx. sema . db , target_module) ?;
100+ file = in_file;
101+ target
102+ } else {
103+ next_space_for_fn_after_call_site ( & call) ?
104+ } ;
105+ let fn_name = fn_name ( & path) ?;
79106 let ( type_params, params) = fn_args ( ctx, & call) ?;
80- Some ( Self { append_fn_at , fn_name, type_params, params } )
107+ Some ( Self { target , fn_name, type_params, params, file , needs_pub } )
81108 }
109+
82110 fn render ( self ) -> Option < FunctionTemplate > {
83111 let placeholder_expr = ast:: make:: expr_todo ( ) ;
84112 let fn_body = ast:: make:: block_expr ( vec ! [ ] , Some ( placeholder_expr) ) ;
85- let fn_def = ast:: make:: fn_def ( self . fn_name , self . type_params , self . params , fn_body) ;
86- let fn_def = ast:: make:: add_newlines ( 2 , fn_def) ;
87- let fn_def = IndentLevel :: from_node ( & self . append_fn_at ) . increase_indent ( fn_def) ;
88- let insert_offset = self . append_fn_at . text_range ( ) . end ( ) ;
113+ let mut fn_def = ast:: make:: fn_def ( self . fn_name , self . type_params , self . params , fn_body) ;
114+ if self . needs_pub {
115+ fn_def = ast:: make:: add_pub_crate_modifier ( fn_def) ;
116+ }
117+
118+ let ( fn_def, insert_offset) = match self . target {
119+ GeneratedFunctionTarget :: BehindItem ( it) => {
120+ let with_leading_blank_line = ast:: make:: add_leading_newlines ( 2 , fn_def) ;
121+ let indented = IndentLevel :: from_node ( & it) . increase_indent ( with_leading_blank_line) ;
122+ ( indented, it. text_range ( ) . end ( ) )
123+ }
124+ GeneratedFunctionTarget :: InEmptyItemList ( it) => {
125+ let indent_once = IndentLevel ( 1 ) ;
126+ let indent = IndentLevel :: from_node ( it. syntax ( ) ) ;
127+
128+ let fn_def = ast:: make:: add_leading_newlines ( 1 , fn_def) ;
129+ let fn_def = indent_once. increase_indent ( fn_def) ;
130+ let fn_def = ast:: make:: add_trailing_newlines ( 1 , fn_def) ;
131+ let fn_def = indent. increase_indent ( fn_def) ;
132+ ( fn_def, it. syntax ( ) . text_range ( ) . start ( ) + TextUnit :: from_usize ( 1 ) )
133+ }
134+ } ;
135+
89136 let cursor_offset_from_fn_start = fn_def
90137 . syntax ( )
91138 . descendants ( )
@@ -94,19 +141,24 @@ impl FunctionBuilder {
94141 . text_range ( )
95142 . start ( ) ;
96143 let cursor_offset = insert_offset + cursor_offset_from_fn_start;
97- Some ( FunctionTemplate { insert_offset, cursor_offset, fn_def } )
144+ Some ( FunctionTemplate { insert_offset, cursor_offset, fn_def, file : self . file } )
98145 }
99146}
100147
101- fn fn_name ( call : & CallExpr ) -> Option < ast:: Name > {
102- let name = call. expr ( ) ?. syntax ( ) . to_string ( ) ;
148+ enum GeneratedFunctionTarget {
149+ BehindItem ( SyntaxNode ) ,
150+ InEmptyItemList ( ast:: ItemList ) ,
151+ }
152+
153+ fn fn_name ( call : & ast:: Path ) -> Option < ast:: Name > {
154+ let name = call. segment ( ) ?. syntax ( ) . to_string ( ) ;
103155 Some ( ast:: make:: name ( & name) )
104156}
105157
106158/// Computes the type variables and arguments required for the generated function
107159fn fn_args (
108160 ctx : & AssistCtx ,
109- call : & CallExpr ,
161+ call : & ast :: CallExpr ,
110162) -> Option < ( Option < ast:: TypeParamList > , ast:: ParamList ) > {
111163 let mut arg_names = Vec :: new ( ) ;
112164 let mut arg_types = Vec :: new ( ) ;
@@ -158,9 +210,9 @@ fn deduplicate_arg_names(arg_names: &mut Vec<String>) {
158210 }
159211}
160212
161- fn fn_arg_name ( fn_arg : & Expr ) -> Option < String > {
213+ fn fn_arg_name ( fn_arg : & ast :: Expr ) -> Option < String > {
162214 match fn_arg {
163- Expr :: CastExpr ( cast_expr) => fn_arg_name ( & cast_expr. expr ( ) ?) ,
215+ ast :: Expr :: CastExpr ( cast_expr) => fn_arg_name ( & cast_expr. expr ( ) ?) ,
164216 _ => Some (
165217 fn_arg
166218 . syntax ( )
@@ -172,7 +224,7 @@ fn fn_arg_name(fn_arg: &Expr) -> Option<String> {
172224 }
173225}
174226
175- fn fn_arg_type ( ctx : & AssistCtx , fn_arg : & Expr ) -> Option < String > {
227+ fn fn_arg_type ( ctx : & AssistCtx , fn_arg : & ast :: Expr ) -> Option < String > {
176228 let ty = ctx. sema . type_of_expr ( fn_arg) ?;
177229 if ty. is_unknown ( ) {
178230 return None ;
@@ -184,7 +236,7 @@ fn fn_arg_type(ctx: &AssistCtx, fn_arg: &Expr) -> Option<String> {
184236/// directly after the current block
185237/// We want to write the generated function directly after
186238/// fns, impls or macro calls, but inside mods
187- fn next_space_for_fn ( expr : & CallExpr ) -> Option < SyntaxNode > {
239+ fn next_space_for_fn_after_call_site ( expr : & ast :: CallExpr ) -> Option < GeneratedFunctionTarget > {
188240 let mut ancestors = expr. syntax ( ) . ancestors ( ) . peekable ( ) ;
189241 let mut last_ancestor: Option < SyntaxNode > = None ;
190242 while let Some ( next_ancestor) = ancestors. next ( ) {
@@ -201,7 +253,32 @@ fn next_space_for_fn(expr: &CallExpr) -> Option<SyntaxNode> {
201253 }
202254 last_ancestor = Some ( next_ancestor) ;
203255 }
204- last_ancestor
256+ last_ancestor. map ( GeneratedFunctionTarget :: BehindItem )
257+ }
258+
259+ fn next_space_for_fn_in_module (
260+ db : & dyn hir:: db:: AstDatabase ,
261+ module : hir:: InFile < hir:: ModuleSource > ,
262+ ) -> Option < ( AssistFile , GeneratedFunctionTarget ) > {
263+ let file = module. file_id . original_file ( db) ;
264+ let assist_file = AssistFile :: TargetFile ( file) ;
265+ let assist_item = match module. value {
266+ hir:: ModuleSource :: SourceFile ( it) => {
267+ if let Some ( last_item) = it. items ( ) . last ( ) {
268+ GeneratedFunctionTarget :: BehindItem ( last_item. syntax ( ) . clone ( ) )
269+ } else {
270+ GeneratedFunctionTarget :: BehindItem ( it. syntax ( ) . clone ( ) )
271+ }
272+ }
273+ hir:: ModuleSource :: Module ( it) => {
274+ if let Some ( last_item) = it. item_list ( ) . and_then ( |it| it. items ( ) . last ( ) ) {
275+ GeneratedFunctionTarget :: BehindItem ( last_item. syntax ( ) . clone ( ) )
276+ } else {
277+ GeneratedFunctionTarget :: InEmptyItemList ( it. item_list ( ) ?)
278+ }
279+ }
280+ } ;
281+ Some ( ( assist_file, assist_item) )
205282}
206283
207284#[ cfg( test) ]
@@ -713,6 +790,111 @@ fn bar(baz_1: Baz, baz_2: Baz, arg_1: &str, arg_2: &str) {
713790 )
714791 }
715792
793+ #[ test]
794+ fn add_function_in_module ( ) {
795+ check_assist (
796+ add_function,
797+ r"
798+ mod bar {}
799+
800+ fn foo() {
801+ bar::my_fn<|>()
802+ }
803+ " ,
804+ r"
805+ mod bar {
806+ pub(crate) fn my_fn() {
807+ <|>todo!()
808+ }
809+ }
810+
811+ fn foo() {
812+ bar::my_fn()
813+ }
814+ " ,
815+ )
816+ }
817+
818+ #[ test]
819+ fn add_function_in_module_containing_other_items ( ) {
820+ check_assist (
821+ add_function,
822+ r"
823+ mod bar {
824+ fn something_else() {}
825+ }
826+
827+ fn foo() {
828+ bar::my_fn<|>()
829+ }
830+ " ,
831+ r"
832+ mod bar {
833+ fn something_else() {}
834+
835+ pub(crate) fn my_fn() {
836+ <|>todo!()
837+ }
838+ }
839+
840+ fn foo() {
841+ bar::my_fn()
842+ }
843+ " ,
844+ )
845+ }
846+
847+ #[ test]
848+ fn add_function_in_nested_module ( ) {
849+ check_assist (
850+ add_function,
851+ r"
852+ mod bar {
853+ mod baz {}
854+ }
855+
856+ fn foo() {
857+ bar::baz::my_fn<|>()
858+ }
859+ " ,
860+ r"
861+ mod bar {
862+ mod baz {
863+ pub(crate) fn my_fn() {
864+ <|>todo!()
865+ }
866+ }
867+ }
868+
869+ fn foo() {
870+ bar::baz::my_fn()
871+ }
872+ " ,
873+ )
874+ }
875+
876+ #[ test]
877+ fn add_function_in_another_file ( ) {
878+ check_assist (
879+ add_function,
880+ r"
881+ //- /main.rs
882+ mod foo;
883+
884+ fn main() {
885+ foo::bar<|>()
886+ }
887+ //- /foo.rs
888+ " ,
889+ r"
890+
891+
892+ pub(crate) fn bar() {
893+ <|>todo!()
894+ }" ,
895+ )
896+ }
897+
716898 #[ test]
717899 fn add_function_not_applicable_if_function_already_exists ( ) {
718900 check_assist_not_applicable (
0 commit comments