aboutsummaryrefslogtreecommitdiff
path: root/crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs
diff options
context:
space:
mode:
authorAleksey Kladov <[email protected]>2021-02-28 18:11:41 +0000
committerAleksey Kladov <[email protected]>2021-02-28 18:14:34 +0000
commit406d96c7d4c3a78d42b58a91ba633333ec37d487 (patch)
treeab17d9e641f19bcf514673c9ed5e9c35646fe269 /crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs
parentaa04e3bbb2fab8ee7f7aa8eb406943d314976a0d (diff)
Use consistent naming for assist
Diffstat (limited to 'crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs')
-rw-r--r--crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs321
1 files changed, 321 insertions, 0 deletions
diff --git a/crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs b/crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs
new file mode 100644
index 000000000..27da28bc0
--- /dev/null
+++ b/crates/ide_assists/src/handlers/replace_for_loop_with_for_each.rs
@@ -0,0 +1,321 @@
1use ast::LoopBodyOwner;
2use hir::known;
3use ide_db::helpers::FamousDefs;
4use stdx::format_to;
5use syntax::{ast, AstNode};
6use test_utils::mark;
7
8use crate::{AssistContext, AssistId, AssistKind, Assists};
9
10// Assist: replace_for_loop_with_for_each
11//
12// Converts a for loop into a for_each loop on the Iterator.
13//
14// ```
15// fn main() {
16// let x = vec![1, 2, 3];
17// for$0 v in x {
18// let y = v * 2;
19// }
20// }
21// ```
22// ->
23// ```
24// fn main() {
25// let x = vec![1, 2, 3];
26// x.into_iter().for_each(|v| {
27// let y = v * 2;
28// });
29// }
30// ```
31pub(crate) fn replace_for_loop_with_for_each(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
32 let for_loop = ctx.find_node_at_offset::<ast::ForExpr>()?;
33 let iterable = for_loop.iterable()?;
34 let pat = for_loop.pat()?;
35 let body = for_loop.loop_body()?;
36 if body.syntax().text_range().start() < ctx.offset() {
37 mark::hit!(not_available_in_body);
38 return None;
39 }
40
41 acc.add(
42 AssistId("replace_for_loop_with_for_each", AssistKind::RefactorRewrite),
43 "Replace this for loop with `Iterator::for_each`",
44 for_loop.syntax().text_range(),
45 |builder| {
46 let mut buf = String::new();
47
48 if let Some((expr_behind_ref, method)) =
49 is_ref_and_impls_iter_method(&ctx.sema, &iterable)
50 {
51 // We have either "for x in &col" and col implements a method called iter
52 // or "for x in &mut col" and col implements a method called iter_mut
53 format_to!(buf, "{}.{}()", expr_behind_ref, method);
54 } else if impls_core_iter(&ctx.sema, &iterable) {
55 format_to!(buf, "{}", iterable);
56 } else {
57 if let ast::Expr::RefExpr(_) = iterable {
58 format_to!(buf, "({}).into_iter()", iterable);
59 } else {
60 format_to!(buf, "{}.into_iter()", iterable);
61 }
62 }
63
64 format_to!(buf, ".for_each(|{}| {});", pat, body);
65
66 builder.replace(for_loop.syntax().text_range(), buf)
67 },
68 )
69}
70
71/// If iterable is a reference where the expression behind the reference implements a method
72/// returning an Iterator called iter or iter_mut (depending on the type of reference) then return
73/// the expression behind the reference and the method name
74fn is_ref_and_impls_iter_method(
75 sema: &hir::Semantics<ide_db::RootDatabase>,
76 iterable: &ast::Expr,
77) -> Option<(ast::Expr, hir::Name)> {
78 let ref_expr = match iterable {
79 ast::Expr::RefExpr(r) => r,
80 _ => return None,
81 };
82 let wanted_method = if ref_expr.mut_token().is_some() { known::iter_mut } else { known::iter };
83 let expr_behind_ref = ref_expr.expr()?;
84 let typ = sema.type_of_expr(&expr_behind_ref)?;
85 let scope = sema.scope(iterable.syntax());
86 let krate = scope.module()?.krate();
87 let traits_in_scope = scope.traits_in_scope();
88 let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
89 let has_wanted_method = typ.iterate_method_candidates(
90 sema.db,
91 krate,
92 &traits_in_scope,
93 Some(&wanted_method),
94 |_, func| {
95 if func.ret_type(sema.db).impls_trait(sema.db, iter_trait, &[]) {
96 return Some(());
97 }
98 None
99 },
100 );
101 has_wanted_method.and(Some((expr_behind_ref, wanted_method)))
102}
103
104/// Whether iterable implements core::Iterator
105fn impls_core_iter(sema: &hir::Semantics<ide_db::RootDatabase>, iterable: &ast::Expr) -> bool {
106 let it_typ = if let Some(i) = sema.type_of_expr(iterable) {
107 i
108 } else {
109 return false;
110 };
111 let module = if let Some(m) = sema.scope(iterable.syntax()).module() {
112 m
113 } else {
114 return false;
115 };
116 let krate = module.krate();
117 if let Some(iter_trait) = FamousDefs(sema, Some(krate)).core_iter_Iterator() {
118 return it_typ.impls_trait(sema.db, iter_trait, &[]);
119 }
120 false
121}
122
123#[cfg(test)]
124mod tests {
125 use crate::tests::{check_assist, check_assist_not_applicable};
126
127 use super::*;
128
129 const EMPTY_ITER_FIXTURE: &'static str = r"
130//- /lib.rs deps:core crate:empty_iter
131pub struct EmptyIter;
132impl Iterator for EmptyIter {
133 type Item = usize;
134 fn next(&mut self) -> Option<Self::Item> { None }
135}
136
137pub struct Empty;
138impl Empty {
139 pub fn iter(&self) -> EmptyIter { EmptyIter }
140 pub fn iter_mut(&self) -> EmptyIter { EmptyIter }
141}
142
143pub struct NoIterMethod;
144";
145
146 fn check_assist_with_fixtures(before: &str, after: &str) {
147 let before = &format!(
148 "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
149 before,
150 FamousDefs::FIXTURE,
151 EMPTY_ITER_FIXTURE
152 );
153 check_assist(replace_for_loop_with_for_each, before, after);
154 }
155
156 #[test]
157 fn test_not_for() {
158 check_assist_not_applicable(
159 replace_for_loop_with_for_each,
160 r"
161let mut x = vec![1, 2, 3];
162x.iter_mut().$0for_each(|v| *v *= 2);
163 ",
164 )
165 }
166
167 #[test]
168 fn test_simple_for() {
169 check_assist(
170 replace_for_loop_with_for_each,
171 r"
172fn main() {
173 let x = vec![1, 2, 3];
174 for $0v in x {
175 v *= 2;
176 }
177}",
178 r"
179fn main() {
180 let x = vec![1, 2, 3];
181 x.into_iter().for_each(|v| {
182 v *= 2;
183 });
184}",
185 )
186 }
187
188 #[test]
189 fn not_available_in_body() {
190 mark::check!(not_available_in_body);
191 check_assist_not_applicable(
192 replace_for_loop_with_for_each,
193 r"
194fn main() {
195 let x = vec![1, 2, 3];
196 for v in x {
197 $0v *= 2;
198 }
199}",
200 )
201 }
202
203 #[test]
204 fn test_for_borrowed() {
205 check_assist_with_fixtures(
206 r"
207use empty_iter::*;
208fn main() {
209 let x = Empty;
210 for $0v in &x {
211 let a = v * 2;
212 }
213}
214",
215 r"
216use empty_iter::*;
217fn main() {
218 let x = Empty;
219 x.iter().for_each(|v| {
220 let a = v * 2;
221 });
222}
223",
224 )
225 }
226
227 #[test]
228 fn test_for_borrowed_no_iter_method() {
229 check_assist_with_fixtures(
230 r"
231use empty_iter::*;
232fn main() {
233 let x = NoIterMethod;
234 for $0v in &x {
235 let a = v * 2;
236 }
237}
238",
239 r"
240use empty_iter::*;
241fn main() {
242 let x = NoIterMethod;
243 (&x).into_iter().for_each(|v| {
244 let a = v * 2;
245 });
246}
247",
248 )
249 }
250
251 #[test]
252 fn test_for_borrowed_mut() {
253 check_assist_with_fixtures(
254 r"
255use empty_iter::*;
256fn main() {
257 let x = Empty;
258 for $0v in &mut x {
259 let a = v * 2;
260 }
261}
262",
263 r"
264use empty_iter::*;
265fn main() {
266 let x = Empty;
267 x.iter_mut().for_each(|v| {
268 let a = v * 2;
269 });
270}
271",
272 )
273 }
274
275 #[test]
276 fn test_for_borrowed_mut_behind_var() {
277 check_assist(
278 replace_for_loop_with_for_each,
279 r"
280fn main() {
281 let x = vec![1, 2, 3];
282 let y = &mut x;
283 for $0v in y {
284 *v *= 2;
285 }
286}",
287 r"
288fn main() {
289 let x = vec![1, 2, 3];
290 let y = &mut x;
291 y.into_iter().for_each(|v| {
292 *v *= 2;
293 });
294}",
295 )
296 }
297
298 #[test]
299 fn test_already_impls_iterator() {
300 check_assist_with_fixtures(
301 r#"
302use empty_iter::*;
303fn main() {
304 let x = Empty;
305 for$0 a in x.iter().take(1) {
306 println!("{}", a);
307 }
308}
309"#,
310 r#"
311use empty_iter::*;
312fn main() {
313 let x = Empty;
314 x.iter().take(1).for_each(|a| {
315 println!("{}", a);
316 });
317}
318"#,
319 );
320 }
321}