3 files changed, 171 insertions, 131 deletions

diff --git a/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs b/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs
index b5b5ada5e..70949ca35 100644
--- a/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs
+++ b/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs
@@ -107,7 +107,7 @@ fn edit_struct_references(
     names: &[ast::Name],
 ) {
     let strukt_def = Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Struct(strukt)));
-    let usages = strukt_def.usages(&ctx.sema).include_self_kw_refs(true).all();
+    let usages = strukt_def.usages(&ctx.sema).include_self_refs().all();
 
     let edit_node = |edit: &mut AssistBuilder, node: SyntaxNode| -> Option<()> {
         match_ast! {
diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index 93b28370c..4116985ae 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/crates/ide_assists/src/handlers/extract_function.rs
@@ -10,7 +10,6 @@ use ide_db::{
 use itertools::Itertools;
 use stdx::format_to;
 use syntax::{
-    algo::SyntaxRewriter,
     ast::{
         self,
         edit::{AstNodeEdit, IndentLevel},
@@ -1362,7 +1361,8 @@ fn rewrite_body_segment(
     syntax: &SyntaxNode,
 ) -> SyntaxNode {
     let syntax = fix_param_usages(ctx, params, syntax);
-    update_external_control_flow(handler, &syntax)
+    update_external_control_flow(handler, &syntax);
+    syntax
 }
 
 /// change all usages to account for added `&`/`&mut` for some params
@@ -1415,75 +1415,65 @@ fn fix_param_usages(ctx: &AssistContext, params: &[Param], syntax: &SyntaxNode)
     res
 }
 
-fn update_external_control_flow(handler: &FlowHandler, syntax: &SyntaxNode) -> SyntaxNode {
-    let mut rewriter = SyntaxRewriter::default();
-
+fn update_external_control_flow(handler: &FlowHandler, syntax: &SyntaxNode) {
     let mut nested_loop = None;
     let mut nested_scope = None;
     for event in syntax.preorder() {
-        let node = match event {
-            WalkEvent::Enter(e) => {
-                match e.kind() {
-                    SyntaxKind::LOOP_EXPR | SyntaxKind::WHILE_EXPR | SyntaxKind::FOR_EXPR => {
-                        if nested_loop.is_none() {
-                            nested_loop = Some(e.clone());
-                        }
+        match event {
+            WalkEvent::Enter(e) => match e.kind() {
+                SyntaxKind::LOOP_EXPR | SyntaxKind::WHILE_EXPR | SyntaxKind::FOR_EXPR => {
+                    if nested_loop.is_none() {
+                        nested_loop = Some(e.clone());
                     }
-                    SyntaxKind::FN
-                    | SyntaxKind::CONST
-                    | SyntaxKind::STATIC
-                    | SyntaxKind::IMPL
-                    | SyntaxKind::MODULE => {
-                        if nested_scope.is_none() {
-                            nested_scope = Some(e.clone());
-                        }
+                }
+                SyntaxKind::FN
+                | SyntaxKind::CONST
+                | SyntaxKind::STATIC
+                | SyntaxKind::IMPL
+                | SyntaxKind::MODULE => {
+                    if nested_scope.is_none() {
+                        nested_scope = Some(e.clone());
                     }
-                    _ => {}
                 }
-                e
-            }
+                _ => {}
+            },
             WalkEvent::Leave(e) => {
+                if nested_scope.is_none() {
+                    if let Some(expr) = ast::Expr::cast(e.clone()) {
+                        match expr {
+                            ast::Expr::ReturnExpr(return_expr) if nested_scope.is_none() => {
+                                let expr = return_expr.expr();
+                                if let Some(replacement) = make_rewritten_flow(handler, expr) {
+                                    ted::replace(return_expr.syntax(), replacement.syntax())
+                                }
+                            }
+                            ast::Expr::BreakExpr(break_expr) if nested_loop.is_none() => {
+                                let expr = break_expr.expr();
+                                if let Some(replacement) = make_rewritten_flow(handler, expr) {
+                                    ted::replace(break_expr.syntax(), replacement.syntax())
+                                }
+                            }
+                            ast::Expr::ContinueExpr(continue_expr) if nested_loop.is_none() => {
+                                if let Some(replacement) = make_rewritten_flow(handler, None) {
+                                    ted::replace(continue_expr.syntax(), replacement.syntax())
+                                }
+                            }
+                            _ => {
+                                // do nothing
+                            }
+                        }
+                    }
+                }
+
                 if nested_loop.as_ref() == Some(&e) {
                     nested_loop = None;
                 }
                 if nested_scope.as_ref() == Some(&e) {
                     nested_scope = None;
                 }
-                continue;
             }
         };
-        if nested_scope.is_some() {
-            continue;
-        }
-        let expr = match ast::Expr::cast(node) {
-            Some(e) => e,
-            None => continue,
-        };
-        match expr {
-            ast::Expr::ReturnExpr(return_expr) if nested_scope.is_none() => {
-                let expr = return_expr.expr();
-                if let Some(replacement) = make_rewritten_flow(handler, expr) {
-                    rewriter.replace_ast(&return_expr.into(), &replacement);
-                }
-            }
-            ast::Expr::BreakExpr(break_expr) if nested_loop.is_none() => {
-                let expr = break_expr.expr();
-                if let Some(replacement) = make_rewritten_flow(handler, expr) {
-                    rewriter.replace_ast(&break_expr.into(), &replacement);
-                }
-            }
-            ast::Expr::ContinueExpr(continue_expr) if nested_loop.is_none() => {
-                if let Some(replacement) = make_rewritten_flow(handler, None) {
-                    rewriter.replace_ast(&continue_expr.into(), &replacement);
-                }
-            }
-            _ => {
-                // do nothing
-            }
-        }
     }
-
-    rewriter.rewrite(syntax)
 }
 
 fn make_rewritten_flow(handler: &FlowHandler, arg_expr: Option<ast::Expr>) -> Option<ast::Expr> {
@@ -1502,7 +1492,7 @@ fn make_rewritten_flow(handler: &FlowHandler, arg_expr: Option<ast::Expr>) -> Op
             make::expr_call(make::expr_path(make_path_from_text("Err")), args)
         }
     };
-    Some(make::expr_return(Some(value)))
+    Some(make::expr_return(Some(value)).clone_for_update())
 }
 
 #[cfg(test)]
diff --git a/crates/ide_assists/src/handlers/pull_assignment_up.rs b/crates/ide_assists/src/handlers/pull_assignment_up.rs
index 04bae4e58..28d14b9c3 100644
--- a/crates/ide_assists/src/handlers/pull_assignment_up.rs
+++ b/crates/ide_assists/src/handlers/pull_assignment_up.rs
@@ -1,6 +1,6 @@
 use syntax::{
-    ast::{self, edit::AstNodeEdit, make},
-    AstNode,
+    ast::{self, make},
+    ted, AstNode,
 };
 
 use crate::{
@@ -37,103 +37,101 @@ use crate::{
 // ```
 pub(crate) fn pull_assignment_up(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let assign_expr = ctx.find_node_at_offset::<ast::BinExpr>()?;
-    let name_expr = if assign_expr.op_kind()? == ast::BinOp::Assignment {
-        assign_expr.lhs()?
-    } else {
+
+    let op_kind = assign_expr.op_kind()?;
+    if op_kind != ast::BinOp::Assignment {
+        cov_mark::hit!(test_cant_pull_non_assignments);
         return None;
+    }
+
+    let mut collector = AssignmentsCollector {
+        sema: &ctx.sema,
+        common_lhs: assign_expr.lhs()?,
+        assignments: Vec::new(),
     };
 
-    let (old_stmt, new_stmt) = if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
-        (
-            ast::Expr::cast(if_expr.syntax().to_owned())?,
-            exprify_if(&if_expr, &ctx.sema, &name_expr)?.indent(if_expr.indent_level()),
-        )
+    let tgt: ast::Expr = if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
+        collector.collect_if(&if_expr)?;
+        if_expr.into()
     } else if let Some(match_expr) = ctx.find_node_at_offset::<ast::MatchExpr>() {
-        (
-            ast::Expr::cast(match_expr.syntax().to_owned())?,
-            exprify_match(&match_expr, &ctx.sema, &name_expr)?,
-        )
+        collector.collect_match(&match_expr)?;
+        match_expr.into()
     } else {
         return None;
     };
 
-    let expr_stmt = make::expr_stmt(new_stmt);
-
     acc.add(
         AssistId("pull_assignment_up", AssistKind::RefactorExtract),
         "Pull assignment up",
-        old_stmt.syntax().text_range(),
+        tgt.syntax().text_range(),
         move |edit| {
-            edit.replace(old_stmt.syntax().text_range(), format!("{} = {};", name_expr, expr_stmt));
+            let assignments: Vec<_> = collector
+                .assignments
+                .into_iter()
+                .map(|(stmt, rhs)| (edit.make_ast_mut(stmt), rhs.clone_for_update()))
+                .collect();
+
+            let tgt = edit.make_ast_mut(tgt);
+
+            for (stmt, rhs) in assignments {
+                ted::replace(stmt.syntax(), rhs.syntax());
+            }
+            let assign_expr = make::expr_assignment(collector.common_lhs, tgt.clone());
+            let assign_stmt = make::expr_stmt(assign_expr);
+
+            ted::replace(tgt.syntax(), assign_stmt.syntax().clone_for_update());
         },
     )
 }
 
-fn exprify_match(
-    match_expr: &ast::MatchExpr,
-    sema: &hir::Semantics<ide_db::RootDatabase>,
-    name: &ast::Expr,
-) -> Option<ast::Expr> {
-    let new_arm_list = match_expr
-        .match_arm_list()?
-        .arms()
-        .map(|arm| {
-            if let ast::Expr::BlockExpr(block) = arm.expr()? {
-                let new_block = exprify_block(&block, sema, name)?.indent(block.indent_level());
-                Some(arm.replace_descendant(block, new_block))
-            } else {
-                None
-            }
-        })
-        .collect::<Option<Vec<_>>>()?;
-    let new_arm_list = match_expr
-        .match_arm_list()?
-        .replace_descendants(match_expr.match_arm_list()?.arms().zip(new_arm_list));
-    Some(make::expr_match(match_expr.expr()?, new_arm_list))
+struct AssignmentsCollector<'a> {
+    sema: &'a hir::Semantics<'a, ide_db::RootDatabase>,
+    common_lhs: ast::Expr,
+    assignments: Vec<(ast::ExprStmt, ast::Expr)>,
 }
 
-fn exprify_if(
-    statement: &ast::IfExpr,
-    sema: &hir::Semantics<ide_db::RootDatabase>,
-    name: &ast::Expr,
-) -> Option<ast::Expr> {
-    let then_branch = exprify_block(&statement.then_branch()?, sema, name)?;
-    let else_branch = match statement.else_branch()? {
-        ast::ElseBranch::Block(ref block) => {
-            ast::ElseBranch::Block(exprify_block(block, sema, name)?)
-        }
-        ast::ElseBranch::IfExpr(expr) => {
-            cov_mark::hit!(test_pull_assignment_up_chained_if);
-            ast::ElseBranch::IfExpr(ast::IfExpr::cast(
-                exprify_if(&expr, sema, name)?.syntax().to_owned(),
-            )?)
+impl<'a> AssignmentsCollector<'a> {
+    fn collect_match(&mut self, match_expr: &ast::MatchExpr) -> Option<()> {
+        for arm in match_expr.match_arm_list()?.arms() {
+            match arm.expr()? {
+                ast::Expr::BlockExpr(block) => self.collect_block(&block)?,
+                _ => return None,
+            }
         }
-    };
-    Some(make::expr_if(statement.condition()?, then_branch, Some(else_branch)))
-}
 
-fn exprify_block(
-    block: &ast::BlockExpr,
-    sema: &hir::Semantics<ide_db::RootDatabase>,
-    name: &ast::Expr,
-) -> Option<ast::BlockExpr> {
-    if block.tail_expr().is_some() {
-        return None;
+        Some(())
     }
+    fn collect_if(&mut self, if_expr: &ast::IfExpr) -> Option<()> {
+        let then_branch = if_expr.then_branch()?;
+        self.collect_block(&then_branch)?;
+
+        match if_expr.else_branch()? {
+            ast::ElseBranch::Block(block) => self.collect_block(&block),
+            ast::ElseBranch::IfExpr(expr) => {
+                cov_mark::hit!(test_pull_assignment_up_chained_if);
+                self.collect_if(&expr)
+            }
+        }
+    }
+    fn collect_block(&mut self, block: &ast::BlockExpr) -> Option<()> {
+        if block.tail_expr().is_some() {
+            return None;
+        }
 
-    let mut stmts: Vec<_> = block.statements().collect();
-    let stmt = stmts.pop()?;
-
-    if let ast::Stmt::ExprStmt(stmt) = stmt {
-        if let ast::Expr::BinExpr(expr) = stmt.expr()? {
-            if expr.op_kind()? == ast::BinOp::Assignment && is_equivalent(sema, &expr.lhs()?, name)
-            {
-                // The last statement in the block is an assignment to the name we want
-                return Some(make::block_expr(stmts, Some(expr.rhs()?)));
+        let last_stmt = block.statements().last()?;
+        if let ast::Stmt::ExprStmt(stmt) = last_stmt {
+            if let ast::Expr::BinExpr(expr) = stmt.expr()? {
+                if expr.op_kind()? == ast::BinOp::Assignment
+                    && is_equivalent(self.sema, &expr.lhs()?, &self.common_lhs)
+                {
+                    self.assignments.push((stmt, expr.rhs()?));
+                    return Some(());
+                }
             }
         }
+
+        None
     }
-    None
 }
 
 fn is_equivalent(
@@ -243,6 +241,38 @@ fn foo() {
     }
 
     #[test]
+    #[ignore]
+    fn test_pull_assignment_up_assignment_expressions() {
+        check_assist(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    match 1 {
+        1 => { $0a = 2; },
+        2 => a = 3,
+        3 => {
+            a = 4
+        }
+    }
+}"#,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    a = match 1 {
+        1 => { 2 },
+        2 => 3,
+        3 => {
+            4
+        }
+    };
+}"#,
+        );
+    }
+
+    #[test]
     fn test_pull_assignment_up_not_last_not_applicable() {
         check_assist_not_applicable(
             pull_assignment_up,
@@ -439,4 +469,24 @@ fn foo() {
 "#,
         )
     }
+
+    #[test]
+    fn test_cant_pull_non_assignments() {
+        cov_mark::check!(test_cant_pull_non_assignments);
+        check_assist_not_applicable(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+    let b = &mut a;
+
+    if true {
+        $0*b + 2;
+    } else {
+        *b + 3;
+    }
+}
+"#,
+        )
+    }
 }