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