From 506293ca43f4cae7520c9e14b6b36b42b1af4ce1 Mon Sep 17 00:00:00 2001 From: Matt Hall Date: Sun, 21 Feb 2021 17:36:37 +0000 Subject: Add convert_for_to_iter_for_each assist --- .../src/handlers/convert_for_to_iter_for_each.rs | 225 +++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs (limited to 'crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs') diff --git a/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs new file mode 100644 index 000000000..d858474c6 --- /dev/null +++ b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs @@ -0,0 +1,225 @@ +use ast::LoopBodyOwner; +use ide_db::helpers::FamousDefs; +use stdx::format_to; +use syntax::{ast, AstNode}; + +use crate::{AssistContext, AssistId, AssistKind, Assists}; + +// Assist: convert_for_to_iter_for_each +// +// Converts a for loop into a for_each loop on the Iterator. +// +// ``` +// fn main() { +// let x = vec![1, 2, 3]; +// for $0v in x { +// let y = v * 2; +// } +// } +// ``` +// -> +// ``` +// fn main() { +// let x = vec![1, 2, 3]; +// x.into_iter().for_each(|v| { +// let y = v * 2; +// }); +// } +// ``` +pub(crate) fn convert_for_to_iter_for_each(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { + let for_loop = ctx.find_node_at_offset::()?; + let iterable = for_loop.iterable()?; + let pat = for_loop.pat()?; + let body = for_loop.loop_body()?; + + let mut buf = String::new(); + + if impls_core_iter(&ctx.sema, &iterable) { + buf += &iterable.to_string(); + } else { + match iterable { + ast::Expr::RefExpr(r) => { + if r.mut_token().is_some() { + format_to!(buf, "{}.iter_mut()", r.expr()?); + } else { + format_to!(buf, "{}.iter()", r.expr()?); + } + } + _ => format_to!(buf, "{}.into_iter()", iterable), + } + } + + format_to!(buf, ".for_each(|{}| {});", pat, body); + + acc.add( + AssistId("convert_for_to_iter_for_each", AssistKind::RefactorRewrite), + "Convert a for loop into an Iterator::for_each", + for_loop.syntax().text_range(), + |builder| builder.replace(for_loop.syntax().text_range(), buf), + ) +} + +fn impls_core_iter(sema: &hir::Semantics, iterable: &ast::Expr) -> bool { + let it_typ = if let Some(i) = sema.type_of_expr(iterable) { + i + } else { + return false; + }; + let module = if let Some(m) = sema.scope(iterable.syntax()).module() { + m + } else { + return false; + }; + let krate = module.krate(); + if let Some(iter_trait) = FamousDefs(sema, Some(krate)).core_iter_Iterator() { + return it_typ.impls_trait(sema.db, iter_trait, &[]); + } + false +} + +#[cfg(test)] +mod tests { + use crate::tests::{check_assist, check_assist_not_applicable}; + + use super::*; + + const EMPTY_ITER_FIXTURE: &'static str = r" +//- /lib.rs deps:core crate:empty_iter +pub struct EmptyIter; +impl Iterator for EmptyIter { + type Item = usize; + fn next(&mut self) -> Option { None } +} + +pub struct Empty; +impl Empty { + pub fn iter(&self) -> EmptyIter { EmptyIter } +} +"; + + #[test] + fn test_not_for() { + check_assist_not_applicable( + convert_for_to_iter_for_each, + r" +let mut x = vec![1, 2, 3]; +x.iter_mut().$0for_each(|v| *v *= 2); + ", + ) + } + + #[test] + fn test_simple_for() { + check_assist( + convert_for_to_iter_for_each, + r" +fn main() { + let x = vec![1, 2, 3]; + for $0v in x { + v *= 2; + } +}", + r" +fn main() { + let x = vec![1, 2, 3]; + x.into_iter().for_each(|v| { + v *= 2; + }); +}", + ) + } + + #[test] + fn test_for_borrowed() { + check_assist( + convert_for_to_iter_for_each, + r" +fn main() { + let x = vec![1, 2, 3]; + for $0v in &x { + let a = v * 2; + } +}", + r" +fn main() { + let x = vec![1, 2, 3]; + x.iter().for_each(|v| { + let a = v * 2; + }); +}", + ) + } + + #[test] + fn test_for_borrowed_mut() { + check_assist( + convert_for_to_iter_for_each, + r" +fn main() { + let x = vec![1, 2, 3]; + for $0v in &mut x { + *v *= 2; + } +}", + r" +fn main() { + let x = vec![1, 2, 3]; + x.iter_mut().for_each(|v| { + *v *= 2; + }); +}", + ) + } + + #[test] + fn test_for_borrowed_mut_behind_var() { + check_assist( + convert_for_to_iter_for_each, + r" +fn main() { + let x = vec![1, 2, 3]; + let y = &mut x; + for $0v in y { + *v *= 2; + } +}", + r" +fn main() { + let x = vec![1, 2, 3]; + let y = &mut x; + y.into_iter().for_each(|v| { + *v *= 2; + }); +}", + ) + } + + #[test] + fn test_take() { + let before = r#" +use empty_iter::*; +fn main() { + let x = Empty; + for$0 a in x.iter().take(1) { + println!("{}", a); + } +} +"#; + let after = r#" +use empty_iter::*; +fn main() { + let x = Empty; + x.iter().take(1).for_each(|a| { + println!("{}", a); + }); +} +"#; + let before = &format!( + "//- /main.rs crate:main deps:core,empty_iter{}{}{}", + before, + FamousDefs::FIXTURE, + EMPTY_ITER_FIXTURE + ); + check_assist(convert_for_to_iter_for_each, before, after); + } +} -- cgit v1.2.3 From 98a626450d43515f8fe91db9f7918c6e69804693 Mon Sep 17 00:00:00 2001 From: Matt Hall Date: Tue, 23 Feb 2021 19:19:48 +0000 Subject: Address review comments * Move code to build replacement into closure * Look for iter/iter_mut methods on types behind reference --- .../src/handlers/convert_for_to_iter_for_each.rs | 158 ++++++++++++++++----- 1 file changed, 125 insertions(+), 33 deletions(-) (limited to 'crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs') diff --git a/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs index d858474c6..f329d435f 100644 --- a/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs +++ b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs @@ -32,33 +32,68 @@ pub(crate) fn convert_for_to_iter_for_each(acc: &mut Assists, ctx: &AssistContex let pat = for_loop.pat()?; let body = for_loop.loop_body()?; - let mut buf = String::new(); + acc.add( + AssistId("convert_for_to_iter_for_each", AssistKind::RefactorRewrite), + "Convert a for loop into an Iterator::for_each", + for_loop.syntax().text_range(), + |builder| { + let mut buf = String::new(); - if impls_core_iter(&ctx.sema, &iterable) { - buf += &iterable.to_string(); - } else { - match iterable { - ast::Expr::RefExpr(r) => { - if r.mut_token().is_some() { - format_to!(buf, "{}.iter_mut()", r.expr()?); + if let Some((expr_behind_ref, method)) = + is_ref_and_impls_iter_method(&ctx.sema, &iterable) + { + // 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); + } else if impls_core_iter(&ctx.sema, &iterable) { + format_to!(buf, "{}", iterable); + } else { + if let ast::Expr::RefExpr(_) = iterable { + format_to!(buf, "({}).into_iter()", iterable); } else { - format_to!(buf, "{}.iter()", r.expr()?); + format_to!(buf, "{}.into_iter()", iterable); } } - _ => format_to!(buf, "{}.into_iter()", iterable), - } - } - format_to!(buf, ".for_each(|{}| {});", pat, body); + format_to!(buf, ".for_each(|{}| {});", pat, body); - acc.add( - AssistId("convert_for_to_iter_for_each", AssistKind::RefactorRewrite), - "Convert a for loop into an Iterator::for_each", - for_loop.syntax().text_range(), - |builder| builder.replace(for_loop.syntax().text_range(), buf), + builder.replace(for_loop.syntax().text_range(), buf) + }, ) } +/// If iterable is a reference where the expression behind the reference implements a method +/// returning an Iterator called iter or iter_mut (depending on the type of reference) then return +/// the expression behind the reference and the method name +fn is_ref_and_impls_iter_method( + sema: &hir::Semantics, + iterable: &ast::Expr, +) -> Option<(ast::Expr, &'static str)> { + let ref_expr = match iterable { + ast::Expr::RefExpr(r) => r, + _ => return None, + }; + let wanted_method = if ref_expr.mut_token().is_some() { "iter_mut" } else { "iter" }; + let expr_behind_ref = ref_expr.expr()?; + let typ = sema.type_of_expr(&expr_behind_ref)?; + let scope = sema.scope(iterable.syntax()); + let krate = scope.module()?.krate(); + let traits_in_scope = scope.traits_in_scope(); + let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?; + let has_wanted_method = + typ.iterate_method_candidates(sema.db, krate, &traits_in_scope, None, |_, func| { + if func.name(sema.db).to_string() != wanted_method { + return None; + } + if func.ret_type(sema.db).impls_trait(sema.db, iter_trait, &[]) { + return Some(()); + } + None + }); + has_wanted_method.and(Some((expr_behind_ref, wanted_method))) +} + +/// Whether iterable implements core::Iterator fn impls_core_iter(sema: &hir::Semantics, iterable: &ast::Expr) -> bool { let it_typ = if let Some(i) = sema.type_of_expr(iterable) { i @@ -94,7 +129,10 @@ impl Iterator for EmptyIter { pub struct Empty; impl Empty { pub fn iter(&self) -> EmptyIter { EmptyIter } + pub fn iter_mut(&self) -> EmptyIter { EmptyIter } } + +pub struct NoIterMethod; "; #[test] @@ -131,43 +169,97 @@ fn main() { #[test] fn test_for_borrowed() { - check_assist( - convert_for_to_iter_for_each, - r" + let before = r" +use empty_iter::*; fn main() { - let x = vec![1, 2, 3]; + let x = Empty; for $0v in &x { let a = v * 2; } -}", +} +"; + let before = &format!( + "//- /main.rs crate:main deps:core,empty_iter{}{}{}", + before, + FamousDefs::FIXTURE, + EMPTY_ITER_FIXTURE + ); + check_assist( + convert_for_to_iter_for_each, + before, r" +use empty_iter::*; fn main() { - let x = vec![1, 2, 3]; + let x = Empty; x.iter().for_each(|v| { let a = v * 2; }); -}", +} +", ) } #[test] - fn test_for_borrowed_mut() { + fn test_for_borrowed_no_iter_method() { + let before = r" +use empty_iter::*; +fn main() { + let x = NoIterMethod; + for $0v in &x { + let a = v * 2; + } +} +"; + let before = &format!( + "//- /main.rs crate:main deps:core,empty_iter{}{}{}", + before, + FamousDefs::FIXTURE, + EMPTY_ITER_FIXTURE + ); check_assist( convert_for_to_iter_for_each, + before, r" +use empty_iter::*; fn main() { - let x = vec![1, 2, 3]; + let x = NoIterMethod; + (&x).into_iter().for_each(|v| { + let a = v * 2; + }); +} +", + ) + } + + #[test] + fn test_for_borrowed_mut() { + let before = r" +use empty_iter::*; +fn main() { + let x = Empty; for $0v in &mut x { - *v *= 2; + let a = v * 2; } -}", +} +"; + let before = &format!( + "//- /main.rs crate:main deps:core,empty_iter{}{}{}", + before, + FamousDefs::FIXTURE, + EMPTY_ITER_FIXTURE + ); + check_assist( + convert_for_to_iter_for_each, + before, r" +use empty_iter::*; fn main() { - let x = vec![1, 2, 3]; + let x = Empty; x.iter_mut().for_each(|v| { - *v *= 2; + let a = v * 2; }); -}", +} +", ) } @@ -195,7 +287,7 @@ fn main() { } #[test] - fn test_take() { + fn test_already_impls_iterator() { let before = r#" use empty_iter::*; fn main() { -- cgit v1.2.3 From a28e8628255198aa36bcde1f380763ef257beabd Mon Sep 17 00:00:00 2001 From: Matt Hall Date: Wed, 24 Feb 2021 19:23:12 +0000 Subject: Address further review comments * Use known names for iter/iter_mut method (simplifies checking if the method exists * Extract code to check assist with fixtures to function --- .../src/handlers/convert_for_to_iter_for_each.rs | 86 +++++++++------------- 1 file changed, 35 insertions(+), 51 deletions(-) (limited to 'crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs') diff --git a/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs index f329d435f..9fddf889c 100644 --- a/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs +++ b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs @@ -1,4 +1,5 @@ use ast::LoopBodyOwner; +use hir::known; use ide_db::helpers::FamousDefs; use stdx::format_to; use syntax::{ast, AstNode}; @@ -68,28 +69,30 @@ pub(crate) fn convert_for_to_iter_for_each(acc: &mut Assists, ctx: &AssistContex fn is_ref_and_impls_iter_method( sema: &hir::Semantics, iterable: &ast::Expr, -) -> Option<(ast::Expr, &'static str)> { +) -> Option<(ast::Expr, hir::Name)> { let ref_expr = match iterable { ast::Expr::RefExpr(r) => r, _ => return None, }; - let wanted_method = if ref_expr.mut_token().is_some() { "iter_mut" } else { "iter" }; + let wanted_method = if ref_expr.mut_token().is_some() { known::iter_mut } else { known::iter }; let expr_behind_ref = ref_expr.expr()?; let typ = sema.type_of_expr(&expr_behind_ref)?; let scope = sema.scope(iterable.syntax()); let krate = scope.module()?.krate(); let traits_in_scope = scope.traits_in_scope(); let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?; - let has_wanted_method = - typ.iterate_method_candidates(sema.db, krate, &traits_in_scope, None, |_, func| { - if func.name(sema.db).to_string() != wanted_method { - return None; - } + let has_wanted_method = typ.iterate_method_candidates( + sema.db, + krate, + &traits_in_scope, + Some(&wanted_method), + |_, func| { if func.ret_type(sema.db).impls_trait(sema.db, iter_trait, &[]) { return Some(()); } None - }); + }, + ); has_wanted_method.and(Some((expr_behind_ref, wanted_method))) } @@ -135,6 +138,16 @@ impl Empty { pub struct NoIterMethod; "; + fn check_assist_with_fixtures(before: &str, after: &str) { + let before = &format!( + "//- /main.rs crate:main deps:core,empty_iter{}{}{}", + before, + FamousDefs::FIXTURE, + EMPTY_ITER_FIXTURE + ); + check_assist(convert_for_to_iter_for_each, before, after); + } + #[test] fn test_not_for() { check_assist_not_applicable( @@ -169,7 +182,8 @@ fn main() { #[test] fn test_for_borrowed() { - let before = r" + check_assist_with_fixtures( + r" use empty_iter::*; fn main() { let x = Empty; @@ -177,16 +191,7 @@ fn main() { let a = v * 2; } } -"; - let before = &format!( - "//- /main.rs crate:main deps:core,empty_iter{}{}{}", - before, - FamousDefs::FIXTURE, - EMPTY_ITER_FIXTURE - ); - check_assist( - convert_for_to_iter_for_each, - before, +", r" use empty_iter::*; fn main() { @@ -201,7 +206,8 @@ fn main() { #[test] fn test_for_borrowed_no_iter_method() { - let before = r" + check_assist_with_fixtures( + r" use empty_iter::*; fn main() { let x = NoIterMethod; @@ -209,16 +215,7 @@ fn main() { let a = v * 2; } } -"; - let before = &format!( - "//- /main.rs crate:main deps:core,empty_iter{}{}{}", - before, - FamousDefs::FIXTURE, - EMPTY_ITER_FIXTURE - ); - check_assist( - convert_for_to_iter_for_each, - before, +", r" use empty_iter::*; fn main() { @@ -233,7 +230,8 @@ fn main() { #[test] fn test_for_borrowed_mut() { - let before = r" + check_assist_with_fixtures( + r" use empty_iter::*; fn main() { let x = Empty; @@ -241,16 +239,7 @@ fn main() { let a = v * 2; } } -"; - let before = &format!( - "//- /main.rs crate:main deps:core,empty_iter{}{}{}", - before, - FamousDefs::FIXTURE, - EMPTY_ITER_FIXTURE - ); - check_assist( - convert_for_to_iter_for_each, - before, +", r" use empty_iter::*; fn main() { @@ -288,7 +277,8 @@ fn main() { #[test] fn test_already_impls_iterator() { - let before = r#" + check_assist_with_fixtures( + r#" use empty_iter::*; fn main() { let x = Empty; @@ -296,8 +286,8 @@ fn main() { println!("{}", a); } } -"#; - let after = r#" +"#, + r#" use empty_iter::*; fn main() { let x = Empty; @@ -305,13 +295,7 @@ fn main() { println!("{}", a); }); } -"#; - let before = &format!( - "//- /main.rs crate:main deps:core,empty_iter{}{}{}", - before, - FamousDefs::FIXTURE, - EMPTY_ITER_FIXTURE +"#, ); - check_assist(convert_for_to_iter_for_each, before, after); } } -- cgit v1.2.3