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