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 // # //- /lib.rs crate:core
  15 // # pub mod iter { pub mod traits { pub mod iterator { pub trait Iterator {} } } }
  16 // # pub struct SomeIter;
  17 // # impl self::iter::traits::iterator::Iterator for SomeIter {}
  18 // # //- /lib.rs crate:main deps:core
  19 // # use core::SomeIter;
  20 // fn main() {
  21 //     let iter = SomeIter;
  22 //     iter.for_each$0(|(x, y)| {
  23 //         println!("x: {}, y: {}", x, y);
  24 //     });
  25 // }
  26 // ```
  27 // ->
  28 // ```
  29 // # use core::SomeIter;
  30 // fn main() {
  31 //     let iter = SomeIter;
  32 //     for (x, y) in iter {
  33 //         println!("x: {}, y: {}", x, y);
  34 //     }
  35 // }
  36 // ```
  37
  38 pub(crate) fn convert_iter_for_each_to_for(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  39     let method = ctx.find_node_at_offset::<ast::MethodCallExpr>()?;
  40
  41     let closure = match method.arg_list()?.args().next()? {
  42         ast::Expr::ClosureExpr(expr) => expr,
  43         _ => return None,
  44     };
  45
  46     let (method, receiver) = validate_method_call_expr(ctx, method)?;
  47
  48     let param_list = closure.param_list()?;
  49     let param = param_list.params().next()?.pat()?;
  50     let body = closure.body()?;
  51
  52     let stmt = method.syntax().parent().and_then(ast::ExprStmt::cast);
  53     let syntax = stmt.as_ref().map_or(method.syntax(), |stmt| stmt.syntax());
  54
  55     acc.add(
  56         AssistId("convert_iter_for_each_to_for", AssistKind::RefactorRewrite),
  57         "Replace this `Iterator::for_each` with a for loop",
  58         syntax.text_range(),
  59         |builder| {
  60             let indent = stmt.as_ref().map_or(method.indent_level(), |stmt| stmt.indent_level());
  61
  62             let block = match body {
  63                 ast::Expr::BlockExpr(block) => block,
  64                 _ => make::block_expr(Vec::new(), Some(body)),
  65             }
  66             .reset_indent()
  67             .indent(indent);
  68
  69             let expr_for_loop = make::expr_for_loop(param, receiver, block);
  70             builder.replace(syntax.text_range(), expr_for_loop.syntax().text())
  71         },
  72     )
  73 }
  74
  75 fn validate_method_call_expr(
  76     ctx: &AssistContext,
  77     expr: ast::MethodCallExpr,
  78 ) -> Option<(ast::Expr, ast::Expr)> {
  79     let name_ref = expr.name_ref()?;
  80     if name_ref.syntax().text_range().intersect(ctx.frange.range).is_none()
  81         || name_ref.text() != "for_each"
  82     {
  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)?;
  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::{self, check_assist};
 102
 103     use super::*;
 104
 105     const EMPTY_ITER_FIXTURE: &'static str = r"
 106 //- /lib.rs deps:core crate:empty_iter
 107 pub struct EmptyIter;
 108 impl Iterator for EmptyIter {
 109     type Item = usize;
 110     fn next(&mut self) -> Option<Self::Item> { None }
 111 }
 112 pub struct Empty;
 113 impl Empty {
 114     pub fn iter(&self) -> EmptyIter { EmptyIter }
 115 }
 116 ";
 117
 118     fn check_assist_with_fixtures(before: &str, after: &str) {
 119         let before = &format!(
 120             "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
 121             before,
 122             EMPTY_ITER_FIXTURE,
 123             FamousDefs::FIXTURE,
 124         );
 125         check_assist(convert_iter_for_each_to_for, before, after);
 126     }
 127
 128     fn check_assist_not_applicable(before: &str) {
 129         let before = &format!(
 130             "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
 131             before,
 132             EMPTY_ITER_FIXTURE,
 133             FamousDefs::FIXTURE,
 134         );
 135         tests::check_assist_not_applicable(convert_iter_for_each_to_for, before);
 136     }
 137
 138     #[test]
 139     fn test_for_each_in_method_stmt() {
 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_method() {
 163         check_assist_with_fixtures(
 164             r#"
 165 use empty_iter::*;
 166 fn main() {
 167     let x = Empty;
 168     x.iter().$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;
 176     for (x, y) in x.iter() {
 177         println!("x: {}, y: {}", x, y);
 178     }
 179 }
 180 "#,
 181         )
 182     }
 183
 184     #[test]
 185     fn test_for_each_in_iter_stmt() {
 186         check_assist_with_fixtures(
 187             r#"
 188 use empty_iter::*;
 189 fn main() {
 190     let x = Empty.iter();
 191     x.$0for_each(|(x, y)| {
 192         println!("x: {}, y: {}", x, y);
 193     });
 194 }"#,
 195             r#"
 196 use empty_iter::*;
 197 fn main() {
 198     let x = Empty.iter();
 199     for (x, y) in x {
 200         println!("x: {}, y: {}", x, y);
 201     }
 202 }
 203 "#,
 204         )
 205     }
 206
 207     #[test]
 208     fn test_for_each_without_braces_stmt() {
 209         check_assist_with_fixtures(
 210             r#"
 211 use empty_iter::*;
 212 fn main() {
 213     let x = Empty;
 214     x.iter().$0for_each(|(x, y)| println!("x: {}, y: {}", x, y));
 215 }"#,
 216             r#"
 217 use empty_iter::*;
 218 fn main() {
 219     let x = Empty;
 220     for (x, y) in x.iter() {
 221         println!("x: {}, y: {}", x, y)
 222     }
 223 }
 224 "#,
 225         )
 226     }
 227
 228     #[test]
 229     fn test_for_each_not_applicable() {
 230         check_assist_not_applicable(
 231             r#"
 232 fn main() {
 233     ().$0for_each(|x| println!("{}", x));
 234 }"#,
 235         )
 236     }
 237
 238     #[test]
 239     fn test_for_each_not_applicable_invalid_cursor_pos() {
 240         check_assist_not_applicable(
 241             r#"
 242 use empty_iter::*;
 243 fn main() {
 244     Empty.iter().for_each(|(x, y)| $0println!("x: {}, y: {}", x, y));
 245 }"#,
 246         )
 247     }
 248 }