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