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.rs225
1 files changed, 225 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..d858474c6
--- /dev/null
+++ b/crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs
@@ -0,0 +1,225 @@
1use ast::LoopBodyOwner;
2use ide_db::helpers::FamousDefs;
3use stdx::format_to;
4use syntax::{ast, AstNode};
5
6use crate::{AssistContext, AssistId, AssistKind, Assists};
7
8// Assist: convert_for_to_iter_for_each
9//
10// Converts a for loop into a for_each loop on the Iterator.
11//
12// ```
13// fn main() {
14// let x = vec![1, 2, 3];
15// for $0v in x {
16// let y = v * 2;
17// }
18// }
19// ```
20// ->
21// ```
22// fn main() {
23// let x = vec![1, 2, 3];
24// x.into_iter().for_each(|v| {
25// let y = v * 2;
26// });
27// }
28// ```
29pub(crate) fn convert_for_to_iter_for_each(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
30 let for_loop = ctx.find_node_at_offset::<ast::ForExpr>()?;
31 let iterable = for_loop.iterable()?;
32 let pat = for_loop.pat()?;
33 let body = for_loop.loop_body()?;
34
35 let mut buf = String::new();
36
37 if impls_core_iter(&ctx.sema, &iterable) {
38 buf += &iterable.to_string();
39 } else {
40 match iterable {
41 ast::Expr::RefExpr(r) => {
42 if r.mut_token().is_some() {
43 format_to!(buf, "{}.iter_mut()", r.expr()?);
44 } else {
45 format_to!(buf, "{}.iter()", r.expr()?);
46 }
47 }
48 _ => format_to!(buf, "{}.into_iter()", iterable),
49 }
50 }
51
52 format_to!(buf, ".for_each(|{}| {});", pat, body);
53
54 acc.add(
55 AssistId("convert_for_to_iter_for_each", AssistKind::RefactorRewrite),
56 "Convert a for loop into an Iterator::for_each",
57 for_loop.syntax().text_range(),
58 |builder| builder.replace(for_loop.syntax().text_range(), buf),
59 )
60}
61
62fn impls_core_iter(sema: &hir::Semantics<ide_db::RootDatabase>, iterable: &ast::Expr) -> bool {
63 let it_typ = if let Some(i) = sema.type_of_expr(iterable) {
64 i
65 } else {
66 return false;
67 };
68 let module = if let Some(m) = sema.scope(iterable.syntax()).module() {
69 m
70 } else {
71 return false;
72 };
73 let krate = module.krate();
74 if let Some(iter_trait) = FamousDefs(sema, Some(krate)).core_iter_Iterator() {
75 return it_typ.impls_trait(sema.db, iter_trait, &[]);
76 }
77 false
78}
79
80#[cfg(test)]
81mod tests {
82 use crate::tests::{check_assist, check_assist_not_applicable};
83
84 use super::*;
85
86 const EMPTY_ITER_FIXTURE: &'static str = r"
87//- /lib.rs deps:core crate:empty_iter
88pub struct EmptyIter;
89impl Iterator for EmptyIter {
90 type Item = usize;
91 fn next(&mut self) -> Option<Self::Item> { None }
92}
93
94pub struct Empty;
95impl Empty {
96 pub fn iter(&self) -> EmptyIter { EmptyIter }
97}
98";
99
100 #[test]
101 fn test_not_for() {
102 check_assist_not_applicable(
103 convert_for_to_iter_for_each,
104 r"
105let mut x = vec![1, 2, 3];
106x.iter_mut().$0for_each(|v| *v *= 2);
107 ",
108 )
109 }
110
111 #[test]
112 fn test_simple_for() {
113 check_assist(
114 convert_for_to_iter_for_each,
115 r"
116fn main() {
117 let x = vec![1, 2, 3];
118 for $0v in x {
119 v *= 2;
120 }
121}",
122 r"
123fn main() {
124 let x = vec![1, 2, 3];
125 x.into_iter().for_each(|v| {
126 v *= 2;
127 });
128}",
129 )
130 }
131
132 #[test]
133 fn test_for_borrowed() {
134 check_assist(
135 convert_for_to_iter_for_each,
136 r"
137fn main() {
138 let x = vec![1, 2, 3];
139 for $0v in &x {
140 let a = v * 2;
141 }
142}",
143 r"
144fn main() {
145 let x = vec![1, 2, 3];
146 x.iter().for_each(|v| {
147 let a = v * 2;
148 });
149}",
150 )
151 }
152
153 #[test]
154 fn test_for_borrowed_mut() {
155 check_assist(
156 convert_for_to_iter_for_each,
157 r"
158fn main() {
159 let x = vec![1, 2, 3];
160 for $0v in &mut x {
161 *v *= 2;
162 }
163}",
164 r"
165fn main() {
166 let x = vec![1, 2, 3];
167 x.iter_mut().for_each(|v| {
168 *v *= 2;
169 });
170}",
171 )
172 }
173
174 #[test]
175 fn test_for_borrowed_mut_behind_var() {
176 check_assist(
177 convert_for_to_iter_for_each,
178 r"
179fn main() {
180 let x = vec![1, 2, 3];
181 let y = &mut x;
182 for $0v in y {
183 *v *= 2;
184 }
185}",
186 r"
187fn main() {
188 let x = vec![1, 2, 3];
189 let y = &mut x;
190 y.into_iter().for_each(|v| {
191 *v *= 2;
192 });
193}",
194 )
195 }
196
197 #[test]
198 fn test_take() {
199 let before = r#"
200use empty_iter::*;
201fn main() {
202 let x = Empty;
203 for$0 a in x.iter().take(1) {
204 println!("{}", a);
205 }
206}
207"#;
208 let after = r#"
209use empty_iter::*;
210fn main() {
211 let x = Empty;
212 x.iter().take(1).for_each(|a| {
213 println!("{}", a);
214 });
215}
216"#;
217 let before = &format!(
218 "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
219 before,
220 FamousDefs::FIXTURE,
221 EMPTY_ITER_FIXTURE
222 );
223 check_assist(convert_for_to_iter_for_each, before, after);
224 }
225}