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