crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs

   1 use ide_db::helpers::FamousDefs;
   2 use syntax::{
   3     ast::{self, edit::AstNodeEdit, make, ArgListOwner},
   4     AstNode,
   5 };
   6
   7 use crate::{AssistContext, AssistId, AssistKind, Assists};
   8
   9 // Assist: convert_iter_for_each_to_for
  10 //
  11 // Converts an Iterator::for_each function into a for loop.
  12 //
  13 // ```
  14 // # //- /lib.rs crate:core
  15 // # pub mod iter { pub mod traits { pub mod iterator { pub trait Iterator {} } } }
  16 // # pub struct SomeIter;
  17 // # impl self::iter::traits::iterator::Iterator for SomeIter {}
  18 // # //- /lib.rs crate:main deps:core
  19 // # use core::SomeIter;
  20 // fn main() {
  21 //     let iter = SomeIter;
  22 //     iter.for_each$0(|(x, y)| {
  23 //         println!("x: {}, y: {}", x, y);
  24 //     });
  25 // }
  26 // ```
  27 // ->
  28 // ```
  29 // # use core::SomeIter;
  30 // fn main() {
  31 //     let iter = SomeIter;
  32 //     for (x, y) in iter {
  33 //         println!("x: {}, y: {}", x, y);
  34 //     }
  35 // }
  36 // ```
  37
  38 pub(crate) fn convert_iter_for_each_to_for(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  39     let method = ctx.find_node_at_offset::<ast::MethodCallExpr>()?;
  40
  41     let closure = match method.arg_list()?.args().next()? {
  42         ast::Expr::ClosureExpr(expr) => expr,
  43         _ => return None,
  44     };
  45
  46     let (method, receiver) = validate_method_call_expr(ctx, method)?;
  47
  48     let param_list = closure.param_list()?;
  49     let param = param_list.params().next()?.pat()?;
  50     let body = closure.body()?;
  51
  52     let stmt = method.syntax().parent().and_then(ast::ExprStmt::cast);
  53     let syntax = stmt.as_ref().map_or(method.syntax(), |stmt| stmt.syntax());
  54
  55     acc.add(
  56         AssistId("convert_iter_for_each_to_for", AssistKind::RefactorRewrite),
  57         "Replace this `Iterator::for_each` with a for loop",
  58         syntax.text_range(),
  59         |builder| {
  60             let indent = stmt.as_ref().map_or(method.indent_level(), |stmt| stmt.indent_level());
  61
  62             let block = match body {
  63                 ast::Expr::BlockExpr(block) => block,
  64                 _ => make::block_expr(Vec::new(), Some(body)),
  65             }
  66             .reset_indent()
  67             .indent(indent);
  68
  69             let expr_for_loop = make::expr_for_loop(param, receiver, block);
  70             builder.replace(syntax.text_range(), expr_for_loop.syntax().text())
  71         },
  72     )
  73 }
  74
  75 fn validate_method_call_expr(
  76     ctx: &AssistContext,
  77     expr: ast::MethodCallExpr,
  78 ) -> Option<(ast::Expr, ast::Expr)> {
  79     let name_ref = expr.name_ref()?;
  80     if name_ref.syntax().text_range().intersect(ctx.frange.range).is_none() {
  81         cov_mark::hit!(test_for_each_not_applicable_invalid_cursor_pos);
  82         return None;
  83     }
  84     if name_ref.text() != "for_each" {
  85         return None;
  86     }
  87
  88     let sema = &ctx.sema;
  89
  90     let receiver = expr.receiver()?;
  91     let expr = ast::Expr::MethodCallExpr(expr);
  92
  93     let it_type = sema.type_of_expr(&receiver)?;
  94     let module = sema.scope(receiver.syntax()).module()?;
  95     let krate = module.krate();
  96
  97     let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
  98     it_type.impls_trait(sema.db, iter_trait, &[]).then(|| (expr, receiver))
  99 }
 100
 101 #[cfg(test)]
 102 mod tests {
 103     use crate::tests::{check_assist, check_assist_not_applicable};
 104
 105     use super::*;
 106
 107     #[test]
 108     fn test_for_each_in_method_stmt() {
 109         check_assist(
 110             convert_iter_for_each_to_for,
 111             r#"
 112 //- minicore: iterators
 113 fn main() {
 114     let it = core::iter::repeat(92);
 115     it.$0for_each(|(x, y)| {
 116         println!("x: {}, y: {}", x, y);
 117     });
 118 }
 119 "#,
 120             r#"
 121 fn main() {
 122     let it = core::iter::repeat(92);
 123     for (x, y) in it {
 124         println!("x: {}, y: {}", x, y);
 125     }
 126 }
 127 "#,
 128         )
 129     }
 130
 131     #[test]
 132     fn test_for_each_in_method() {
 133         check_assist(
 134             convert_iter_for_each_to_for,
 135             r#"
 136 //- minicore: iterators
 137 fn main() {
 138     let it = core::iter::repeat(92);
 139     it.$0for_each(|(x, y)| {
 140         println!("x: {}, y: {}", x, y);
 141     })
 142 }
 143 "#,
 144             r#"
 145 fn main() {
 146     let it = core::iter::repeat(92);
 147     for (x, y) in it {
 148         println!("x: {}, y: {}", x, y);
 149     }
 150 }
 151 "#,
 152         )
 153     }
 154
 155     #[test]
 156     fn test_for_each_without_braces_stmt() {
 157         check_assist(
 158             convert_iter_for_each_to_for,
 159             r#"
 160 //- minicore: iterators
 161 fn main() {
 162     let it = core::iter::repeat(92);
 163     it.$0for_each(|(x, y)| println!("x: {}, y: {}", x, y));
 164 }
 165 "#,
 166             r#"
 167 fn main() {
 168     let it = core::iter::repeat(92);
 169     for (x, y) in it {
 170         println!("x: {}, y: {}", x, y)
 171     }
 172 }
 173 "#,
 174         )
 175     }
 176
 177     #[test]
 178     fn test_for_each_not_applicable() {
 179         check_assist_not_applicable(
 180             convert_iter_for_each_to_for,
 181             r#"
 182 //- minicore: iterators
 183 fn main() {
 184     ().$0for_each(|x| println!("{}", x));
 185 }"#,
 186         )
 187     }
 188
 189     #[test]
 190     fn test_for_each_not_applicable_invalid_cursor_pos() {
 191         cov_mark::check!(test_for_each_not_applicable_invalid_cursor_pos);
 192         check_assist_not_applicable(
 193             convert_iter_for_each_to_for,
 194             r#"
 195 //- minicore: iterators
 196 fn main() {
 197     core::iter::repeat(92).for_each(|(x, y)| $0println!("x: {}, y: {}", x, y));
 198 }"#,
 199         )
 200     }
 201 }