@@ -17,7 +17,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
1717// fn foo() -> i32 { 42i32 }
1818// ```
1919pub ( 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) ?;
2121 let module = ctx. sema . scope ( tail_expr. syntax ( ) ) . module ( ) ?;
2222 let ty = ctx. sema . type_of_expr ( & tail_expr) ?;
2323 if ty. is_unit ( ) {
@@ -27,7 +27,10 @@ pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext)
2727
2828 acc. add (
2929 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+ } ,
3134 tail_expr. syntax ( ) . text_range ( ) ,
3235 |builder| {
3336 match builder_edit_pos {
@@ -38,7 +41,7 @@ pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext)
3841 builder. replace ( text_range, & format ! ( "-> {}" , ty) )
3942 }
4043 }
41- if wrap_expr {
44+ if let FnType :: Closure { wrap_expr : true } = fn_type {
4245 mark:: hit!( wrap_closure_non_block_expr) ;
4346 // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
4447 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
7275 }
7376}
7477
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) =
7785 if let Some ( closure) = ctx. find_node_at_offset :: < ast:: ClosureExpr > ( ) {
7886 let rpipe_pos = closure. param_list ( ) ?. syntax ( ) . last_token ( ) ?. text_range ( ) . end ( ) ;
7987 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
8694 } ;
8795
8896 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)
9098 } else {
9199 let func = ctx. find_node_at_offset :: < ast:: Fn > ( ) ?;
92100 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
97105
98106 let ret_range_end = body. l_curly_token ( ) ?. text_range ( ) . start ( ) ;
99107 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)
101109 } ;
102110 let frange = ctx. frange . range ;
103111 if return_type_range. contains_range ( frange) {
@@ -109,7 +117,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(ast::Expr, InsertOrReplace, bool
109117 } else {
110118 return None ;
111119 }
112- Some ( ( tail_expr , action , wrap_expr ) )
120+ Some ( ( fn_type , tail_expr , action ) )
113121}
114122
115123#[ cfg( test) ]
0 commit comments