@@ -17,7 +17,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
17
17
// fn foo() -> i32 { 42i32 }
18
18
// ```
19
19
pub ( crate ) fn infer_function_return_type ( acc : & mut Assists , ctx : & AssistContext ) -> Option < ( ) > {
20
- let ( tail_expr , builder_edit_pos , wrap_expr ) = extract_tail ( ctx) ?;
20
+ let ( fn_type , tail_expr , builder_edit_pos ) = extract_tail ( ctx) ?;
21
21
let module = ctx. sema . scope ( tail_expr. syntax ( ) ) . module ( ) ?;
22
22
let ty = ctx. sema . type_of_expr ( & tail_expr) ?;
23
23
if ty. is_unit ( ) {
@@ -27,7 +27,10 @@ pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext)
27
27
28
28
acc. add (
29
29
AssistId ( "infer_function_return_type" , AssistKind :: RefactorRewrite ) ,
30
- "Add this function's return type" ,
30
+ match fn_type {
31
+ FnType :: Function => "Add this function's return type" ,
32
+ FnType :: Closure { .. } => "Add this closure's return type" ,
33
+ } ,
31
34
tail_expr. syntax ( ) . text_range ( ) ,
32
35
|builder| {
33
36
match builder_edit_pos {
@@ -38,7 +41,7 @@ pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext)
38
41
builder. replace ( text_range, & format ! ( "-> {}" , ty) )
39
42
}
40
43
}
41
- if wrap_expr {
44
+ if let FnType :: Closure { wrap_expr : true } = fn_type {
42
45
mark:: hit!( wrap_closure_non_block_expr) ;
43
46
// `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
44
47
builder. replace ( tail_expr. syntax ( ) . text_range ( ) , & format ! ( "{{{}}}" , tail_expr) ) ;
@@ -72,8 +75,13 @@ fn ret_ty_to_action(ret_ty: Option<ast::RetType>, insert_pos: TextSize) -> Optio
72
75
}
73
76
}
74
77
75
- fn extract_tail ( ctx : & AssistContext ) -> Option < ( ast:: Expr , InsertOrReplace , bool ) > {
76
- let ( tail_expr, return_type_range, action, wrap_expr) =
78
+ enum FnType {
79
+ Function ,
80
+ Closure { wrap_expr : bool } ,
81
+ }
82
+
83
+ fn extract_tail ( ctx : & AssistContext ) -> Option < ( FnType , ast:: Expr , InsertOrReplace ) > {
84
+ let ( fn_type, tail_expr, return_type_range, action) =
77
85
if let Some ( closure) = ctx. find_node_at_offset :: < ast:: ClosureExpr > ( ) {
78
86
let rpipe_pos = closure. param_list ( ) ?. syntax ( ) . last_token ( ) ?. text_range ( ) . end ( ) ;
79
87
let action = ret_ty_to_action ( closure. ret_type ( ) , rpipe_pos) ?;
@@ -86,7 +94,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(ast::Expr, InsertOrReplace, bool
86
94
} ;
87
95
88
96
let ret_range = TextRange :: new ( rpipe_pos, body_start) ;
89
- ( tail_expr, ret_range, action, wrap_expr )
97
+ ( FnType :: Closure { wrap_expr } , tail_expr, ret_range, action)
90
98
} else {
91
99
let func = ctx. find_node_at_offset :: < ast:: Fn > ( ) ?;
92
100
let rparen_pos = func. param_list ( ) ?. r_paren_token ( ) ?. text_range ( ) . end ( ) ;
@@ -97,7 +105,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(ast::Expr, InsertOrReplace, bool
97
105
98
106
let ret_range_end = body. l_curly_token ( ) ?. text_range ( ) . start ( ) ;
99
107
let ret_range = TextRange :: new ( rparen_pos, ret_range_end) ;
100
- ( tail_expr, ret_range, action, false )
108
+ ( FnType :: Function , tail_expr, ret_range, action)
101
109
} ;
102
110
let frange = ctx. frange . range ;
103
111
if return_type_range. contains_range ( frange) {
@@ -109,7 +117,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(ast::Expr, InsertOrReplace, bool
109
117
} else {
110
118
return None ;
111
119
}
112
- Some ( ( tail_expr , action , wrap_expr ) )
120
+ Some ( ( fn_type , tail_expr , action ) )
113
121
}
114
122
115
123
#[ cfg( test) ]
0 commit comments