Pull up with clone subtree

crates/ide-assists/src/handlers/generate_getter_or_setter.rs

@@ -294,7 +294,7 @@ fn generate_setter_from_info(info: &AssistInfo, record_field_info: &RecordFieldI
     let self_expr = make::ext::expr_self();
     let lhs = make::expr_field(self_expr, field_name);
     let rhs = make::expr_path(make::ext::ident_path(field_name));
-    let assign_stmt = make::expr_stmt(make::expr_assignment(lhs, rhs));
+    let assign_stmt = make::expr_stmt(make::expr_assignment(lhs, rhs).into());
     let body = make::block_expr([assign_stmt.into()], None);
 
     // Make the setter fn

crates/ide-assists/src/handlers/pull_assignment_up.rs

@@ -1,7 +1,8 @@
 use syntax::{
     AstNode,
-    ast::{self, make},
-    ted,
+    algo::find_node_at_range,
+    ast::{self, syntax_factory::SyntaxFactory},
+    syntax_editor::SyntaxEditor,
 };
 
 use crate::{
@@ -66,33 +67,52 @@ pub(crate) fn pull_assignment_up(acc: &mut Assists, ctx: &AssistContext<'_>) ->
             return None;
         }
     }
-
+    let target = tgt.syntax().text_range();
+
+    let edit_tgt = tgt.syntax().clone_subtree();
+    let assignments: Vec<_> = collector
+        .assignments
+        .into_iter()
+        .filter_map(|(stmt, rhs)| {
+            Some((
+                find_node_at_range::<ast::BinExpr>(
+                    &edit_tgt,
+                    stmt.syntax().text_range() - target.start(),
+                )?,
+                find_node_at_range::<ast::Expr>(
+                    &edit_tgt,
+                    rhs.syntax().text_range() - target.start(),
+                )?,
+            ))
+        })
+        .collect();
+
+    let mut editor = SyntaxEditor::new(edit_tgt);
+
+    for (stmt, rhs) in assignments {
+        let mut stmt = stmt.syntax().clone();
+        if let Some(parent) = stmt.parent() {
+            if ast::ExprStmt::cast(parent.clone()).is_some() {
+                stmt = parent.clone();
+            }
+        }
+        editor.replace(stmt, rhs.syntax());
+    }
+    let new_tgt_root = editor.finish().new_root().clone();
+    let new_tgt = ast::Expr::cast(new_tgt_root)?;
     acc.add(
         AssistId::refactor_extract("pull_assignment_up"),
         "Pull assignment up",
-        tgt.syntax().text_range(),
+        target,
         move |edit| {
-            let assignments: Vec<_> = collector
-                .assignments
-                .into_iter()
-                .map(|(stmt, rhs)| (edit.make_mut(stmt), rhs.clone_for_update()))
-                .collect();
-
-            let tgt = edit.make_mut(tgt);
-
-            for (stmt, rhs) in assignments {
-                let mut stmt = stmt.syntax().clone();
-                if let Some(parent) = stmt.parent() {
-                    if ast::ExprStmt::cast(parent.clone()).is_some() {
-                        stmt = parent.clone();
-                    }
-                }
-                ted::replace(stmt, 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());
+            let make = SyntaxFactory::with_mappings();
+            let mut editor = edit.make_editor(tgt.syntax());
+            let assign_expr = make.expr_assignment(collector.common_lhs, new_tgt.clone());
+            let assign_stmt = make.expr_stmt(assign_expr.into());
+
+            editor.replace(tgt.syntax(), assign_stmt.syntax());
+            editor.add_mappings(make.finish_with_mappings());
+            edit.add_file_edits(ctx.vfs_file_id(), editor);
         },
     )
 }

crates/syntax/src/ast/make.rs

@@ -680,7 +680,7 @@ pub fn expr_tuple(elements: impl IntoIterator<Item = ast::Expr>) -> ast::TupleEx
     let expr = elements.into_iter().format(", ");
     expr_from_text(&format!("({expr})"))
 }
-pub fn expr_assignment(lhs: ast::Expr, rhs: ast::Expr) -> ast::Expr {
+pub fn expr_assignment(lhs: ast::Expr, rhs: ast::Expr) -> ast::BinExpr {
     expr_from_text(&format!("{lhs} = {rhs}"))
 }
 fn expr_from_text<E: Into<ast::Expr> + AstNode>(text: &str) -> E {

crates/syntax/src/ast/syntax_factory/constructors.rs

@@ -440,6 +440,19 @@ impl SyntaxFactory {
         ast
     }
 
+    pub fn expr_assignment(&self, lhs: ast::Expr, rhs: ast::Expr) -> ast::BinExpr {
+        let ast = make::expr_assignment(lhs.clone(), rhs.clone()).clone_for_update();
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(lhs.syntax().clone(), ast.lhs().unwrap().syntax().clone());
+            builder.map_node(rhs.syntax().clone(), ast.rhs().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+
     pub fn expr_bin(&self, lhs: ast::Expr, op: ast::BinaryOp, rhs: ast::Expr) -> ast::BinExpr {
         let ast::Expr::BinExpr(ast) =
             make::expr_bin_op(lhs.clone(), op, rhs.clone()).clone_for_update()