2 files changed, 41 insertions, 49 deletions

diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs
index 217f692a4..368d48a71 100644
--- a/crates/ra_assists/src/assist_context.rs
+++ b/crates/ra_assists/src/assist_context.rs
@@ -5,14 +5,13 @@ use std::mem;
 use algo::find_covering_element;
 use hir::Semantics;
 use ra_db::{FileId, FileRange};
-use ra_fmt::{leading_indent, reindent};
 use ra_ide_db::{
     source_change::{SourceChange, SourceFileEdit},
     RootDatabase,
 };
 use syntax::{
     algo::{self, find_node_at_offset, SyntaxRewriter},
-    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextSize,
+    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxToken, TextRange, TextSize,
     TokenAtOffset,
 };
 use text_edit::{TextEdit, TextEditBuilder};
@@ -269,20 +268,6 @@ impl AssistBuilder {
     pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
         algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
     }
-    /// Replaces specified `node` of text with a given string, reindenting the
-    /// string to maintain `node`'s existing indent.
-    // FIXME: remove in favor of syntax::edit::IndentLevel::increase_indent
-    pub(crate) fn replace_node_and_indent(
-        &mut self,
-        node: &SyntaxNode,
-        replace_with: impl Into<String>,
-    ) {
-        let mut replace_with = replace_with.into();
-        if let Some(indent) = leading_indent(node) {
-            replace_with = reindent(&replace_with, &indent)
-        }
-        self.replace(node.text_range(), replace_with)
-    }
     pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
         let node = rewriter.rewrite_root().unwrap();
         let new = rewriter.rewrite(&node);
diff --git a/crates/ra_assists/src/handlers/change_return_type_to_result.rs b/crates/ra_assists/src/handlers/change_return_type_to_result.rs
index d5a68a24c..be480943c 100644
--- a/crates/ra_assists/src/handlers/change_return_type_to_result.rs
+++ b/crates/ra_assists/src/handlers/change_return_type_to_result.rs
@@ -1,10 +1,12 @@
+use std::iter;
+
 use syntax::{
-    ast::{self, BlockExpr, Expr, LoopBodyOwner},
+    ast::{self, make, BlockExpr, Expr, LoopBodyOwner},
     AstNode, SyntaxNode,
 };
+use test_utils::mark;
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
-use test_utils::mark;
 
 // Assist: change_return_type_to_result
 //
@@ -44,7 +46,13 @@ pub(crate) fn change_return_type_to_result(acc: &mut Assists, ctx: &AssistContex
             tail_return_expr_collector.collect_tail_exprs(block_expr);
 
             for ret_expr_arg in tail_return_expr_collector.exprs_to_wrap {
-                builder.replace_node_and_indent(&ret_expr_arg, format!("Ok({})", ret_expr_arg));
+                let ok_wrapped = make::expr_call(
+                    make::expr_path(make::path_unqualified(make::path_segment(make::name_ref(
+                        "Ok",
+                    )))),
+                    make::arg_list(iter::once(ret_expr_arg.clone())),
+                );
+                builder.replace_ast(ret_expr_arg, ok_wrapped);
             }
 
             match ctx.config.snippet_cap {
@@ -60,7 +68,7 @@ pub(crate) fn change_return_type_to_result(acc: &mut Assists, ctx: &AssistContex
 }
 
 struct TailReturnCollector {
-    exprs_to_wrap: Vec<SyntaxNode>,
+    exprs_to_wrap: Vec<ast::Expr>,
 }
 
 impl TailReturnCollector {
@@ -86,7 +94,8 @@ impl TailReturnCollector {
             if let Some(last_exprs) = get_tail_expr_from_block(&expr) {
                 for last_expr in last_exprs {
                     let last_expr = match last_expr {
-                        NodeType::Node(expr) | NodeType::Leaf(expr) => expr,
+                        NodeType::Node(expr) => expr,
+                        NodeType::Leaf(expr) => expr.syntax().clone(),
                     };
 
                     if let Some(last_expr) = Expr::cast(last_expr.clone()) {
@@ -113,12 +122,12 @@ impl TailReturnCollector {
             }
             Expr::ReturnExpr(ret_expr) => {
                 if let Some(ret_expr_arg) = &ret_expr.expr() {
-                    self.exprs_to_wrap.push(ret_expr_arg.syntax().clone());
+                    self.exprs_to_wrap.push(ret_expr_arg.clone());
                 }
             }
             Expr::BreakExpr(break_expr) if collect_break => {
                 if let Some(break_expr_arg) = &break_expr.expr() {
-                    self.exprs_to_wrap.push(break_expr_arg.syntax().clone());
+                    self.exprs_to_wrap.push(break_expr_arg.clone());
                 }
             }
             Expr::IfExpr(if_expr) => {
@@ -166,14 +175,11 @@ impl TailReturnCollector {
                     NodeType::Leaf(expr) => {
                         self.exprs_to_wrap.push(expr.clone());
                     }
-                    NodeType::Node(expr) => match &Expr::cast(expr.clone()) {
-                        Some(last_expr) => {
-                            self.fetch_tail_exprs(last_expr);
-                        }
-                        None => {
-                            self.exprs_to_wrap.push(expr.clone());
+                    NodeType::Node(expr) => {
+                        if let Some(last_expr) = Expr::cast(expr.clone()) {
+                            self.fetch_tail_exprs(&last_expr);
                         }
-                    },
+                    }
                 }
             }
         }
@@ -182,7 +188,7 @@ impl TailReturnCollector {
 
 #[derive(Debug)]
 enum NodeType {
-    Leaf(SyntaxNode),
+    Leaf(ast::Expr),
     Node(SyntaxNode),
 }
 
@@ -233,25 +239,26 @@ fn get_tail_expr_from_block(expr: &Expr) -> Option<Vec<NodeType>> {
 
             Some(arms)
         }
-        Expr::BreakExpr(expr) => expr.expr().map(|e| vec![NodeType::Leaf(e.syntax().clone())]),
+        Expr::BreakExpr(expr) => expr.expr().map(|e| vec![NodeType::Leaf(e)]),
         Expr::ReturnExpr(ret_expr) => Some(vec![NodeType::Node(ret_expr.syntax().clone())]),
-        Expr::CallExpr(call_expr) => Some(vec![NodeType::Leaf(call_expr.syntax().clone())]),
-        Expr::Literal(lit_expr) => Some(vec![NodeType::Leaf(lit_expr.syntax().clone())]),
-        Expr::TupleExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::ArrayExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::ParenExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::PathExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::RecordExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::IndexExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::MethodCallExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::AwaitExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::CastExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::RefExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::PrefixExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::RangeExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::BinExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::MacroCall(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
-        Expr::BoxExpr(expr) => Some(vec![NodeType::Leaf(expr.syntax().clone())]),
+
+        Expr::CallExpr(_)
+        | Expr::Literal(_)
+        | Expr::TupleExpr(_)
+        | Expr::ArrayExpr(_)
+        | Expr::ParenExpr(_)
+        | Expr::PathExpr(_)
+        | Expr::RecordExpr(_)
+        | Expr::IndexExpr(_)
+        | Expr::MethodCallExpr(_)
+        | Expr::AwaitExpr(_)
+        | Expr::CastExpr(_)
+        | Expr::RefExpr(_)
+        | Expr::PrefixExpr(_)
+        | Expr::RangeExpr(_)
+        | Expr::BinExpr(_)
+        | Expr::MacroCall(_)
+        | Expr::BoxExpr(_) => Some(vec![NodeType::Leaf(expr.clone())]),
         _ => None,
     }
 }