Unnamed repository; edit this file 'description' to name the repository.
| -rw-r--r-- | crates/ide-assists/src/handlers/convert_iter_for_each_to_for.rs | 45 | ||||
| -rw-r--r-- | crates/syntax/src/ast/syntax_factory/constructors.rs | 12 |
2 files changed, 38 insertions, 19 deletions
diff --git a/crates/ide-assists/src/handlers/convert_iter_for_each_to_for.rs b/crates/ide-assists/src/handlers/convert_iter_for_each_to_for.rs index ecfbc3a07d..cc5cc490f1 100644 --- a/crates/ide-assists/src/handlers/convert_iter_for_each_to_for.rs +++ b/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); }, ) } diff --git a/crates/syntax/src/ast/syntax_factory/constructors.rs b/crates/syntax/src/ast/syntax_factory/constructors.rs index 5688b8b1f5..0f3b3d301c 100644 --- a/crates/syntax/src/ast/syntax_factory/constructors.rs +++ b/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, |