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.reindent_to(indent);
  65
  66             let expr_for_loop = make::expr_for_loop(param, receiver, block);
  67             builder.replace(range, expr_for_loop.to_string())
  68         },
  69     )
  70 }
  71
  72 fn validate_method_call_expr(
  73     ctx: &AssistContext,
  74     expr: ast::MethodCallExpr,
  75 ) -> Option<(ast::Expr, ast::Expr)> {
  76     let name_ref = expr.name_ref()?;
  77     if name_ref.syntax().text_range().intersect(ctx.frange.range).is_none() {
  78         cov_mark::hit!(test_for_each_not_applicable_invalid_cursor_pos);
  79         return None;
  80     }
  81     if name_ref.text() != "for_each" {
  82         return None;
  83     }
  84
  85     let sema = &ctx.sema;
  86
  87     let receiver = expr.receiver()?;
  88     let expr = ast::Expr::MethodCallExpr(expr);
  89
  90     let it_type = sema.type_of_expr(&receiver)?.adjusted();
  91     let module = sema.scope(receiver.syntax()).module()?;
  92     let krate = module.krate();
  93
  94     let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
  95     it_type.impls_trait(sema.db, iter_trait, &[]).then(|| (expr, receiver))
  96 }
  97
  98 #[cfg(test)]
  99 mod tests {
 100     use crate::tests::{check_assist, check_assist_not_applicable};
 101
 102     use super::*;
 103
 104     #[test]
 105     fn test_for_each_in_method_stmt() {
 106         check_assist(
 107             convert_iter_for_each_to_for,
 108             r#"
 109 //- minicore: iterators
 110 fn main() {
 111     let it = core::iter::repeat(92);
 112     it.$0for_each(|(x, y)| {
 113         println!("x: {}, y: {}", x, y);
 114     });
 115 }
 116 "#,
 117             r#"
 118 fn main() {
 119     let it = core::iter::repeat(92);
 120     for (x, y) in it {
 121         println!("x: {}, y: {}", x, y);
 122     }
 123 }
 124 "#,
 125         )
 126     }
 127
 128     #[test]
 129     fn test_for_each_in_method() {
 130         check_assist(
 131             convert_iter_for_each_to_for,
 132             r#"
 133 //- minicore: iterators
 134 fn main() {
 135     let it = core::iter::repeat(92);
 136     it.$0for_each(|(x, y)| {
 137         println!("x: {}, y: {}", x, y);
 138     })
 139 }
 140 "#,
 141             r#"
 142 fn main() {
 143     let it = core::iter::repeat(92);
 144     for (x, y) in it {
 145         println!("x: {}, y: {}", x, y);
 146     }
 147 }
 148 "#,
 149         )
 150     }
 151
 152     #[test]
 153     fn test_for_each_without_braces_stmt() {
 154         check_assist(
 155             convert_iter_for_each_to_for,
 156             r#"
 157 //- minicore: iterators
 158 fn main() {
 159     let it = core::iter::repeat(92);
 160     it.$0for_each(|(x, y)| println!("x: {}, y: {}", x, y));
 161 }
 162 "#,
 163             r#"
 164 fn main() {
 165     let it = core::iter::repeat(92);
 166     for (x, y) in it {
 167         println!("x: {}, y: {}", x, y)
 168     }
 169 }
 170 "#,
 171         )
 172     }
 173
 174     #[test]
 175     fn test_for_each_not_applicable() {
 176         check_assist_not_applicable(
 177             convert_iter_for_each_to_for,
 178             r#"
 179 //- minicore: iterators
 180 fn main() {
 181     ().$0for_each(|x| println!("{}", x));
 182 }"#,
 183         )
 184     }
 185
 186     #[test]
 187     fn test_for_each_not_applicable_invalid_cursor_pos() {
 188         cov_mark::check!(test_for_each_not_applicable_invalid_cursor_pos);
 189         check_assist_not_applicable(
 190             convert_iter_for_each_to_for,
 191             r#"
 192 //- minicore: iterators
 193 fn main() {
 194     core::iter::repeat(92).for_each(|(x, y)| $0println!("x: {}, y: {}", x, y));
 195 }"#,
 196         )
 197     }
 198 }