@@ -3,8 +3,8 @@ use ra_syntax::{
3
3
SyntaxKind , SyntaxNode , TextUnit ,
4
4
} ;
5
5
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 } ;
8
8
use hir:: HirDisplay ;
9
9
use rustc_hash:: { FxHashMap , FxHashSet } ;
10
10
@@ -16,7 +16,7 @@ use rustc_hash::{FxHashMap, FxHashSet};
16
16
// struct Baz;
17
17
// fn baz() -> Baz { Baz }
18
18
// fn foo() {
19
- // bar<|>("", baz());
19
+ // bar<|>("", baz());
20
20
// }
21
21
//
22
22
// ```
@@ -25,7 +25,7 @@ use rustc_hash::{FxHashMap, FxHashSet};
25
25
// struct Baz;
26
26
// fn baz() -> Baz { Baz }
27
27
// fn foo() {
28
- // bar("", baz());
28
+ // bar("", baz());
29
29
// }
30
30
//
31
31
// fn bar(arg: &str, baz: Baz) {
@@ -38,21 +38,30 @@ pub(crate) fn add_function(ctx: AssistCtx) -> Option<Assist> {
38
38
let call = path_expr. syntax ( ) . parent ( ) . and_then ( ast:: CallExpr :: cast) ?;
39
39
let path = path_expr. path ( ) ?;
40
40
41
- if path. qualifier ( ) . is_some ( ) {
42
- return None ;
43
- }
44
-
45
41
if ctx. sema . resolve_path ( & path) . is_some ( ) {
46
42
// The function call already resolves, no need to add a function
47
43
return None ;
48
44
}
49
45
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) ?;
51
59
52
60
ctx. add_assist ( AssistId ( "add_function" ) , "Add function" , |edit| {
53
61
edit. target ( call. syntax ( ) . text_range ( ) ) ;
54
62
55
63
if let Some ( function_template) = function_builder. render ( ) {
64
+ edit. set_file ( function_template. file ) ;
56
65
edit. set_cursor ( function_template. cursor_offset ) ;
57
66
edit. insert ( function_template. insert_offset , function_template. fn_def . to_string ( ) ) ;
58
67
}
@@ -63,29 +72,67 @@ struct FunctionTemplate {
63
72
insert_offset : TextUnit ,
64
73
cursor_offset : TextUnit ,
65
74
fn_def : ast:: SourceFile ,
75
+ file : AssistFile ,
66
76
}
67
77
68
78
struct FunctionBuilder {
69
- append_fn_at : SyntaxNode ,
79
+ target : GeneratedFunctionTarget ,
70
80
fn_name : ast:: Name ,
71
81
type_params : Option < ast:: TypeParamList > ,
72
82
params : ast:: ParamList ,
83
+ file : AssistFile ,
84
+ needs_pub : bool ,
73
85
}
74
86
75
87
impl 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) ?;
79
106
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 } )
81
108
}
109
+
82
110
fn render ( self ) -> Option < FunctionTemplate > {
83
111
let placeholder_expr = ast:: make:: expr_todo ( ) ;
84
112
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
+
89
136
let cursor_offset_from_fn_start = fn_def
90
137
. syntax ( )
91
138
. descendants ( )
@@ -94,19 +141,24 @@ impl FunctionBuilder {
94
141
. text_range ( )
95
142
. start ( ) ;
96
143
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 } )
98
145
}
99
146
}
100
147
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 ( ) ;
103
155
Some ( ast:: make:: name ( & name) )
104
156
}
105
157
106
158
/// Computes the type variables and arguments required for the generated function
107
159
fn fn_args (
108
160
ctx : & AssistCtx ,
109
- call : & CallExpr ,
161
+ call : & ast :: CallExpr ,
110
162
) -> Option < ( Option < ast:: TypeParamList > , ast:: ParamList ) > {
111
163
let mut arg_names = Vec :: new ( ) ;
112
164
let mut arg_types = Vec :: new ( ) ;
@@ -158,9 +210,9 @@ fn deduplicate_arg_names(arg_names: &mut Vec<String>) {
158
210
}
159
211
}
160
212
161
- fn fn_arg_name ( fn_arg : & Expr ) -> Option < String > {
213
+ fn fn_arg_name ( fn_arg : & ast :: Expr ) -> Option < String > {
162
214
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 ( ) ?) ,
164
216
_ => Some (
165
217
fn_arg
166
218
. syntax ( )
@@ -172,7 +224,7 @@ fn fn_arg_name(fn_arg: &Expr) -> Option<String> {
172
224
}
173
225
}
174
226
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 > {
176
228
let ty = ctx. sema . type_of_expr ( fn_arg) ?;
177
229
if ty. is_unknown ( ) {
178
230
return None ;
@@ -184,7 +236,7 @@ fn fn_arg_type(ctx: &AssistCtx, fn_arg: &Expr) -> Option<String> {
184
236
/// directly after the current block
185
237
/// We want to write the generated function directly after
186
238
/// 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 > {
188
240
let mut ancestors = expr. syntax ( ) . ancestors ( ) . peekable ( ) ;
189
241
let mut last_ancestor: Option < SyntaxNode > = None ;
190
242
while let Some ( next_ancestor) = ancestors. next ( ) {
@@ -201,7 +253,32 @@ fn next_space_for_fn(expr: &CallExpr) -> Option<SyntaxNode> {
201
253
}
202
254
last_ancestor = Some ( next_ancestor) ;
203
255
}
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) )
205
282
}
206
283
207
284
#[ cfg( test) ]
@@ -713,6 +790,111 @@ fn bar(baz_1: Baz, baz_2: Baz, arg_1: &str, arg_2: &str) {
713
790
)
714
791
}
715
792
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
+
716
898
#[ test]
717
899
fn add_function_not_applicable_if_function_already_exists ( ) {
718
900
check_assist_not_applicable (
0 commit comments