Merge pull request #22081 from A4-Tacks/migrate-loop-iter-for

internal: Migrate `convert_iter_for_each_to_for` assist to SyntaxEditor

Author: ShoyuVanilla

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

@@ -1,12 +1,11 @@
 use hir::{Name, sym};
 use ide_db::famous_defs::FamousDefs;
-use stdx::format_to;
 use syntax::{
     AstNode,
     ast::{self, HasArgList, HasLoopBody, edit::AstNodeEdit},
 };
 
-use crate::{AssistContext, AssistId, Assists};
+use crate::{AssistContext, AssistId, Assists, utils::wrap_paren};
 
 // Assist: convert_iter_for_each_to_for
 //
@@ -115,32 +114,40 @@ pub(crate) fn convert_for_loop_with_for_each(
         "Replace this for loop with `Iterator::for_each`",
         for_loop.syntax().text_range(),
         |builder| {
-            let mut buf = String::new();
+            let editor = builder.make_editor(for_loop.syntax());
+            let make = editor.make();
+
+            let mut receiver = iterable.clone();
 
-            if let Some((expr_behind_ref, method, krate)) =
+            let iter_method = if let Some((expr_behind_ref, method, krate)) =
                 is_ref_and_impls_iter_method(&ctx.sema, &iterable)
             {
+                receiver = expr_behind_ref;
                 // We have either "for x in &col" and col implements a method called iter
                 //             or "for x in &mut col" and col implements a method called iter_mut
-                format_to!(
-                    buf,
-                    "{expr_behind_ref}.{}()",
-                    method.display(ctx.db(), krate.edition(ctx.db()))
-                );
-            } else if let ast::Expr::RangeExpr(..) = iterable {
-                // range expressions need to be parenthesized for the syntax to be correct
-                format_to!(buf, "({iterable})");
-            } else if impls_core_iter(&ctx.sema, &iterable) {
-                format_to!(buf, "{iterable}");
-            } else if let ast::Expr::RefExpr(_) = iterable {
-                format_to!(buf, "({iterable}).into_iter()");
+                method.display(ctx.db(), krate.edition(ctx.db())).to_string()
             } else {
-                format_to!(buf, "{iterable}.into_iter()");
+                "into_iter".to_owned()
+            };
+
+            receiver = wrap_paren(receiver, make, ast::prec::ExprPrecedence::Postfix);
+
+            if !impls_core_iter(&ctx.sema, &iterable) {
+                receiver = make
+                    .expr_method_call(receiver, make.name_ref(&iter_method), make.arg_list([]))
+                    .into();
             }
 
-            format_to!(buf, ".for_each(|{pat}| {body});");
+            let loop_arg = make.expr_closure([make.untyped_param(pat)], body.into());
+            let for_each = make.expr_method_call(
+                receiver,
+                make.name_ref("for_each"),
+                make.arg_list([loop_arg.into()]),
+            );
+            let for_each = make.expr_stmt(for_each.into());
 
-            builder.replace(for_loop.syntax().text_range(), buf)
+            editor.replace(for_loop.syntax(), for_each.syntax());
+            builder.add_file_edits(ctx.vfs_file_id(), editor);
         },
     )
 }

crates/ide-db/src/source_change.rs

@@ -285,11 +285,11 @@ impl SourceChangeBuilder {
         SyntaxEditor::new(node.ancestors().last().unwrap_or_else(|| node.clone())).0
     }
 
-    pub fn add_file_edits(&mut self, file_id: impl Into<FileId>, edit: SyntaxEditor) {
+    pub fn add_file_edits(&mut self, file_id: impl Into<FileId>, editor: SyntaxEditor) {
         match self.file_editors.entry(file_id.into()) {
-            Entry::Occupied(mut entry) => entry.get_mut().merge(edit),
+            Entry::Occupied(mut entry) => entry.get_mut().merge(editor),
             Entry::Vacant(entry) => {
-                entry.insert(edit);
+                entry.insert(editor);
             }
         }
     }

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

@@ -193,6 +193,18 @@ impl SyntaxFactory {
         ast
     }
 
+    pub fn untyped_param(&self, pat: ast::Pat) -> ast::Param {
+        let ast = make::untyped_param(pat.clone()).clone_for_update();
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(pat.syntax().clone(), ast.pat().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+
     pub fn ty_fn_ptr<I: Iterator<Item = Param>>(
         &self,
         is_unsafe: bool,
@@ -1088,13 +1100,15 @@ impl SyntaxFactory {
         let ast = make::expr_closure(args, expr.clone()).clone_for_update();
 
         if let Some(mut mapping) = self.mappings() {
-            let mut builder = SyntaxMappingBuilder::new(ast.syntax.clone());
-            builder.map_children(
-                input,
-                ast.param_list().unwrap().params().map(|param| param.syntax().clone()),
-            );
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
             builder.map_node(expr.syntax().clone(), ast.body().unwrap().syntax().clone());
             builder.finish(&mut mapping);
+
+            let param_list = ast.param_list().unwrap();
+            let mut params_builder = SyntaxMappingBuilder::new(param_list.syntax().clone());
+            params_builder
+                .map_children(input, param_list.params().map(|param| param.syntax().clone()));
+            params_builder.finish(&mut mapping);
         }
 
         ast