@@ -9,9 +9,10 @@ use syntax::{
99 ast:: {
1010 self ,
1111 edit:: { AstNodeEdit , IndentLevel } ,
12- make, BlockExpr , Expr , ExprStmt , HasArgList ,
12+ syntax_factory:: SyntaxFactory ,
13+ BlockExpr , Expr , ExprStmt , HasArgList ,
1314 } ,
14- ted , AstNode , AstPtr , TextSize ,
15+ AstNode , AstPtr , TextSize ,
1516} ;
1617
1718use crate :: { adjusted_display_range, fix, Assist , Diagnostic , DiagnosticCode , DiagnosticsContext } ;
@@ -223,8 +224,9 @@ fn remove_unnecessary_wrapper(
223224
224225 let inner_arg = call_expr. arg_list ( ) ?. args ( ) . next ( ) ?;
225226
226- let mut builder = SourceChangeBuilder :: new ( expr_ptr. file_id . original_file ( ctx. sema . db ) ) ;
227-
227+ let file_id = expr_ptr. file_id . original_file ( db) ;
228+ let mut builder = SourceChangeBuilder :: new ( file_id) ;
229+ let mut editor;
228230 match inner_arg {
229231 // We're returning `()`
230232 Expr :: TupleExpr ( tup) if tup. fields ( ) . next ( ) . is_none ( ) => {
@@ -233,35 +235,33 @@ fn remove_unnecessary_wrapper(
233235 . parent ( )
234236 . and_then ( Either :: < ast:: ReturnExpr , ast:: StmtList > :: cast) ?;
235237
238+ editor = builder. make_editor ( parent. syntax ( ) ) ;
239+ let make = SyntaxFactory :: new ( ) ;
240+
236241 match parent {
237242 Either :: Left ( ret_expr) => {
238- let old = builder. make_mut ( ret_expr) ;
239- let new = make:: expr_return ( None ) . clone_for_update ( ) ;
240-
241- ted:: replace ( old. syntax ( ) , new. syntax ( ) ) ;
243+ editor. replace ( ret_expr. syntax ( ) , make. expr_return ( None ) . syntax ( ) ) ;
242244 }
243245 Either :: Right ( stmt_list) => {
244- if stmt_list. statements ( ) . count ( ) == 0 {
245- let block = stmt_list. syntax ( ) . parent ( ) . and_then ( ast:: BlockExpr :: cast) ?;
246- let old = builder. make_mut ( block) ;
247- let new = make:: expr_empty_block ( ) . clone_for_update ( ) ;
248-
249- ted:: replace ( old. syntax ( ) , new. syntax ( ) ) ;
246+ let new_block = if stmt_list. statements ( ) . next ( ) . is_none ( ) {
247+ make. expr_empty_block ( )
250248 } else {
251- let old = builder . make_syntax_mut ( stmt_list. syntax ( ) . parent ( ) ? ) ;
252- let new = make :: block_expr ( stmt_list . statements ( ) , None ) . clone_for_update ( ) ;
249+ make . block_expr ( stmt_list. statements ( ) , None )
250+ } ;
253251
254- ted:: replace ( old, new. syntax ( ) ) ;
255- }
252+ editor. replace ( stmt_list. syntax ( ) . parent ( ) ?, new_block. syntax ( ) ) ;
256253 }
257254 }
255+
256+ editor. add_mappings ( make. finish_with_mappings ( ) ) ;
258257 }
259258 _ => {
260- let call_mut = builder. make_mut ( call_expr. clone ( ) ) ;
261- ted :: replace ( call_mut . syntax ( ) , inner_arg. clone_for_update ( ) . syntax ( ) ) ;
259+ editor = builder. make_editor ( call_expr. syntax ( ) ) ;
260+ editor . replace ( call_expr . syntax ( ) , inner_arg. syntax ( ) ) ;
262261 }
263262 }
264263
264+ builder. add_file_edits ( file_id, editor) ;
265265 let name = format ! ( "Remove unnecessary {}() wrapper" , variant. name( db) . as_str( ) ) ;
266266 acc. push ( fix (
267267 "remove_unnecessary_wrapper" ,
@@ -848,6 +848,29 @@ fn div(x: i32, y: i32) -> i32 {
848848 ) ;
849849 }
850850
851+ #[ test]
852+ fn unwrap_return_type_option_tail_unit ( ) {
853+ check_fix (
854+ r#"
855+ //- minicore: option, result
856+ fn div(x: i32, y: i32) {
857+ if y == 0 {
858+ panic!();
859+ }
860+
861+ Ok(())$0
862+ }
863+ "# ,
864+ r#"
865+ fn div(x: i32, y: i32) {
866+ if y == 0 {
867+ panic!();
868+ }
869+ }
870+ "# ,
871+ ) ;
872+ }
873+
851874 #[ test]
852875 fn unwrap_return_type_handles_generic_functions ( ) {
853876 check_fix (
0 commit comments