diff options
-rw-r--r-- | crates/hir_expand/src/name.rs | 1 | ||||
-rw-r--r-- | crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs | 301 | ||||
-rw-r--r-- | crates/ide_assists/src/lib.rs | 2 | ||||
-rw-r--r-- | crates/ide_assists/src/tests/generated.rs | 23 |
4 files changed, 327 insertions, 0 deletions
diff --git a/crates/hir_expand/src/name.rs b/crates/hir_expand/src/name.rs index c7609e90d..c94fb580a 100644 --- a/crates/hir_expand/src/name.rs +++ b/crates/hir_expand/src/name.rs | |||
@@ -189,6 +189,7 @@ pub mod known { | |||
189 | // Components of known path (function name) | 189 | // Components of known path (function name) |
190 | filter_map, | 190 | filter_map, |
191 | next, | 191 | next, |
192 | iter_mut, | ||
192 | // Builtin macros | 193 | // Builtin macros |
193 | file, | 194 | file, |
194 | column, | 195 | column, |
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 @@ | |||
1 | use ast::LoopBodyOwner; | ||
2 | use hir::known; | ||
3 | use ide_db::helpers::FamousDefs; | ||
4 | use stdx::format_to; | ||
5 | use syntax::{ast, AstNode}; | ||
6 | |||
7 | use 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 | // ``` | ||
30 | pub(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 | ||
69 | fn 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 | ||
100 | fn 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)] | ||
119 | mod 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 | ||
126 | pub struct EmptyIter; | ||
127 | impl Iterator for EmptyIter { | ||
128 | type Item = usize; | ||
129 | fn next(&mut self) -> Option<Self::Item> { None } | ||
130 | } | ||
131 | |||
132 | pub struct Empty; | ||
133 | impl Empty { | ||
134 | pub fn iter(&self) -> EmptyIter { EmptyIter } | ||
135 | pub fn iter_mut(&self) -> EmptyIter { EmptyIter } | ||
136 | } | ||
137 | |||
138 | pub 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" | ||
156 | let mut x = vec![1, 2, 3]; | ||
157 | x.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" | ||
167 | fn main() { | ||
168 | let x = vec![1, 2, 3]; | ||
169 | for $0v in x { | ||
170 | v *= 2; | ||
171 | } | ||
172 | }", | ||
173 | r" | ||
174 | fn 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" | ||
187 | use empty_iter::*; | ||
188 | fn main() { | ||
189 | let x = Empty; | ||
190 | for $0v in &x { | ||
191 | let a = v * 2; | ||
192 | } | ||
193 | } | ||
194 | ", | ||
195 | r" | ||
196 | use empty_iter::*; | ||
197 | fn 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" | ||
211 | use empty_iter::*; | ||
212 | fn main() { | ||
213 | let x = NoIterMethod; | ||
214 | for $0v in &x { | ||
215 | let a = v * 2; | ||
216 | } | ||
217 | } | ||
218 | ", | ||
219 | r" | ||
220 | use empty_iter::*; | ||
221 | fn 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" | ||
235 | use empty_iter::*; | ||
236 | fn main() { | ||
237 | let x = Empty; | ||
238 | for $0v in &mut x { | ||
239 | let a = v * 2; | ||
240 | } | ||
241 | } | ||
242 | ", | ||
243 | r" | ||
244 | use empty_iter::*; | ||
245 | fn 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" | ||
260 | fn 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" | ||
268 | fn 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#" | ||
282 | use empty_iter::*; | ||
283 | fn main() { | ||
284 | let x = Empty; | ||
285 | for$0 a in x.iter().take(1) { | ||
286 | println!("{}", a); | ||
287 | } | ||
288 | } | ||
289 | "#, | ||
290 | r#" | ||
291 | use empty_iter::*; | ||
292 | fn main() { | ||
293 | let x = Empty; | ||
294 | x.iter().take(1).for_each(|a| { | ||
295 | println!("{}", a); | ||
296 | }); | ||
297 | } | ||
298 | "#, | ||
299 | ); | ||
300 | } | ||
301 | } | ||
diff --git a/crates/ide_assists/src/lib.rs b/crates/ide_assists/src/lib.rs index 7067cf8b6..f4c7e6fbf 100644 --- a/crates/ide_assists/src/lib.rs +++ b/crates/ide_assists/src/lib.rs | |||
@@ -114,6 +114,7 @@ mod handlers { | |||
114 | mod apply_demorgan; | 114 | mod apply_demorgan; |
115 | mod auto_import; | 115 | mod auto_import; |
116 | mod change_visibility; | 116 | mod change_visibility; |
117 | mod convert_for_to_iter_for_each; | ||
117 | mod convert_integer_literal; | 118 | mod convert_integer_literal; |
118 | mod early_return; | 119 | mod early_return; |
119 | mod expand_glob_import; | 120 | mod expand_glob_import; |
@@ -175,6 +176,7 @@ mod handlers { | |||
175 | apply_demorgan::apply_demorgan, | 176 | apply_demorgan::apply_demorgan, |
176 | auto_import::auto_import, | 177 | auto_import::auto_import, |
177 | change_visibility::change_visibility, | 178 | change_visibility::change_visibility, |
179 | convert_for_to_iter_for_each::convert_for_to_iter_for_each, | ||
178 | convert_integer_literal::convert_integer_literal, | 180 | convert_integer_literal::convert_integer_literal, |
179 | early_return::convert_to_guarded_return, | 181 | early_return::convert_to_guarded_return, |
180 | expand_glob_import::expand_glob_import, | 182 | expand_glob_import::expand_glob_import, |
diff --git a/crates/ide_assists/src/tests/generated.rs b/crates/ide_assists/src/tests/generated.rs index 701091a6b..d42875822 100644 --- a/crates/ide_assists/src/tests/generated.rs +++ b/crates/ide_assists/src/tests/generated.rs | |||
@@ -193,6 +193,29 @@ pub(crate) fn frobnicate() {} | |||
193 | } | 193 | } |
194 | 194 | ||
195 | #[test] | 195 | #[test] |
196 | fn doctest_convert_for_to_iter_for_each() { | ||
197 | check_doc_test( | ||
198 | "convert_for_to_iter_for_each", | ||
199 | r#####" | ||
200 | fn main() { | ||
201 | let x = vec![1, 2, 3]; | ||
202 | for $0v in x { | ||
203 | let y = v * 2; | ||
204 | } | ||
205 | } | ||
206 | "#####, | ||
207 | r#####" | ||
208 | fn main() { | ||
209 | let x = vec![1, 2, 3]; | ||
210 | x.into_iter().for_each(|v| { | ||
211 | let y = v * 2; | ||
212 | }); | ||
213 | } | ||
214 | "#####, | ||
215 | ) | ||
216 | } | ||
217 | |||
218 | #[test] | ||
196 | fn doctest_convert_integer_literal() { | 219 | fn doctest_convert_integer_literal() { |
197 | check_doc_test( | 220 | check_doc_test( |
198 | "convert_integer_literal", | 221 | "convert_integer_literal", |