crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs

   1 use ide_db::helpers::FamousDefs;
   2 use stdx::format_to;
   3 use syntax::{AstNode, ast::{self, ArgListOwner}};
   4
   5 use crate::{AssistContext, AssistId, AssistKind, Assists};
   6
   7 /// Assist: convert_iter_for_each_to_for
   8 //
   9 /// Converts an Iterator::for_each function into a for loop.
  10 ///
  11 /// ```rust
  12 /// fn main() {
  13 ///     let vec = vec![(1, 2), (2, 3), (3, 4)];
  14 ///     x.iter().for_each(|(x, y)| {
  15 ///         println!("x: {}, y: {}", x, y);
  16 ///    })
  17 /// }
  18 /// ```
  19 /// ->
  20 /// ```rust
  21 /// fn main() {
  22 ///     let vec = vec![(1, 2), (2, 3), (3, 4)];
  23 ///     for (x, y) in x.iter() {
  24 ///         println!("x: {}, y: {}", x, y);
  25 ///     });
  26 /// }
  27 /// ```
  28 pub(crate) fn convert_iter_for_each_to_for(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  29     let closure;
  30
  31     let total_expr = match ctx.find_node_at_offset::<ast::Expr>()? {
  32         ast::Expr::MethodCallExpr(expr) => {
  33             closure = match expr.arg_list()?.args().next()? {
  34                 ast::Expr::ClosureExpr(expr) => expr,
  35                 _ => { return None; }
  36             };
  37
  38             expr
  39         },
  40         ast::Expr::ClosureExpr(expr) => {
  41             closure = expr;
  42             ast::MethodCallExpr::cast(closure.syntax().ancestors().nth(2)?)?
  43         },
  44         _ => { return None; }
  45     };
  46
  47     let (total_expr, parent) = validate_method_call_expr(&ctx.sema, total_expr)?;
  48
  49     let param_list = closure.param_list()?;
  50     let param = param_list.params().next()?;
  51     let body = closure.body()?;
  52
  53     acc.add(
  54         AssistId("convert_iter_for_each_to_for", AssistKind::RefactorRewrite),
  55         "Replace this `Iterator::for_each` with a for loop",
  56         total_expr.syntax().text_range(),
  57         |builder| {
  58             let mut buf = String::new();
  59
  60             format_to!(buf, "for {} in {} ", param, parent);
  61
  62             match body {
  63                 ast::Expr::BlockExpr(body) => format_to!(buf, "{}", body),
  64                 _ => format_to!(buf, "{{\n{}\n}}", body)
  65             }
  66
  67             builder.replace(total_expr.syntax().text_range(), buf)
  68         },
  69     )
  70 }
  71
  72 fn validate_method_call_expr(
  73     sema: &hir::Semantics<ide_db::RootDatabase>,
  74     expr: ast::MethodCallExpr,
  75 ) -> Option<(ast::Expr, ast::Expr)> {
  76     if expr.name_ref()?.text() != "for_each" {
  77         return None;
  78     }
  79
  80     let expr = ast::Expr::MethodCallExpr(expr);
  81     let parent = ast::Expr::cast(expr.syntax().first_child()?)?;
  82
  83     let it_type = sema.type_of_expr(&parent)?;
  84     let module = sema.scope(parent.syntax()).module()?;
  85     let krate = module.krate();
  86
  87     let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
  88     it_type.impls_trait(sema.db, iter_trait, &[]).then(|| (expr, parent))
  89 }
  90
  91 #[cfg(test)]
  92 mod tests {
  93     use crate::tests::check_assist;
  94
  95     use super::*;
  96
  97     #[test]
  98     fn test_for_each_in_method() {
  99         check_assist(
 100             convert_iter_for_each_to_for,
 101             r"
 102 fn main() {
 103     let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
 104     x.iter().$0for_each(|(x, y)| {
 105         dbg!(x, y)
 106     });
 107 }",
 108             r"
 109 fn main() {
 110     let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
 111     for (x, y) in x.iter() {
 112         dbg!(x, y)
 113     };
 114 }",
 115         )
 116     }
 117
 118     #[test]
 119     fn test_for_each_in_closure() {
 120         check_assist(
 121             convert_iter_for_each_to_for,
 122             r"
 123 fn main() {
 124     let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
 125     x.iter().for_each($0|(x, y)| {
 126         dbg!(x, y)
 127     });
 128 }",
 129             r"
 130 fn main() {
 131     let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
 132     for (x, y) in x.iter() {
 133         dbg!(x, y)
 134     };
 135 }",
 136         )
 137     }
 138 }