aboutsummaryrefslogtreecommitdiff
path: root/crates/ide_assists/src
diff options
context:
space:
mode:
Diffstat (limited to 'crates/ide_assists/src')
-rw-r--r--crates/ide_assists/src/handlers/convert_for_to_iter_for_each.rs158
1 files changed, 125 insertions, 33 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
index d858474c6..f329d435f 100644
--- 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
@@ -32,33 +32,68 @@ pub(crate) fn convert_for_to_iter_for_each(acc: &mut Assists, ctx: &AssistContex
32 let pat = for_loop.pat()?; 32 let pat = for_loop.pat()?;
33 let body = for_loop.loop_body()?; 33 let body = for_loop.loop_body()?;
34 34
35 let mut buf = String::new(); 35 acc.add(
36 AssistId("convert_for_to_iter_for_each", AssistKind::RefactorRewrite),
37 "Convert a for loop into an Iterator::for_each",
38 for_loop.syntax().text_range(),
39 |builder| {
40 let mut buf = String::new();
36 41
37 if impls_core_iter(&ctx.sema, &iterable) { 42 if let Some((expr_behind_ref, method)) =
38 buf += &iterable.to_string(); 43 is_ref_and_impls_iter_method(&ctx.sema, &iterable)
39 } else { 44 {
40 match iterable { 45 // We have either "for x in &col" and col implements a method called iter
41 ast::Expr::RefExpr(r) => { 46 // or "for x in &mut col" and col implements a method called iter_mut
42 if r.mut_token().is_some() { 47 format_to!(buf, "{}.{}()", expr_behind_ref, method);
43 format_to!(buf, "{}.iter_mut()", r.expr()?); 48 } else if impls_core_iter(&ctx.sema, &iterable) {
49 format_to!(buf, "{}", iterable);
50 } else {
51 if let ast::Expr::RefExpr(_) = iterable {
52 format_to!(buf, "({}).into_iter()", iterable);
44 } else { 53 } else {
45 format_to!(buf, "{}.iter()", r.expr()?); 54 format_to!(buf, "{}.into_iter()", iterable);
46 } 55 }
47 } 56 }
48 _ => format_to!(buf, "{}.into_iter()", iterable),
49 }
50 }
51 57
52 format_to!(buf, ".for_each(|{}| {});", pat, body); 58 format_to!(buf, ".for_each(|{}| {});", pat, body);
53 59
54 acc.add( 60 builder.replace(for_loop.syntax().text_range(), buf)
55 AssistId("convert_for_to_iter_for_each", AssistKind::RefactorRewrite), 61 },
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 ) 62 )
60} 63}
61 64
65/// If iterable is a reference where the expression behind the reference implements a method
66/// returning an Iterator called iter or iter_mut (depending on the type of reference) then return
67/// the expression behind the reference and the method name
68fn is_ref_and_impls_iter_method(
69 sema: &hir::Semantics<ide_db::RootDatabase>,
70 iterable: &ast::Expr,
71) -> Option<(ast::Expr, &'static str)> {
72 let ref_expr = match iterable {
73 ast::Expr::RefExpr(r) => r,
74 _ => return None,
75 };
76 let wanted_method = if ref_expr.mut_token().is_some() { "iter_mut" } else { "iter" };
77 let expr_behind_ref = ref_expr.expr()?;
78 let typ = sema.type_of_expr(&expr_behind_ref)?;
79 let scope = sema.scope(iterable.syntax());
80 let krate = scope.module()?.krate();
81 let traits_in_scope = scope.traits_in_scope();
82 let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
83 let has_wanted_method =
84 typ.iterate_method_candidates(sema.db, krate, &traits_in_scope, None, |_, func| {
85 if func.name(sema.db).to_string() != wanted_method {
86 return None;
87 }
88 if func.ret_type(sema.db).impls_trait(sema.db, iter_trait, &[]) {
89 return Some(());
90 }
91 None
92 });
93 has_wanted_method.and(Some((expr_behind_ref, wanted_method)))
94}
95
96/// Whether iterable implements core::Iterator
62fn impls_core_iter(sema: &hir::Semantics<ide_db::RootDatabase>, iterable: &ast::Expr) -> bool { 97fn 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) { 98 let it_typ = if let Some(i) = sema.type_of_expr(iterable) {
64 i 99 i
@@ -94,7 +129,10 @@ impl Iterator for EmptyIter {
94pub struct Empty; 129pub struct Empty;
95impl Empty { 130impl Empty {
96 pub fn iter(&self) -> EmptyIter { EmptyIter } 131 pub fn iter(&self) -> EmptyIter { EmptyIter }
132 pub fn iter_mut(&self) -> EmptyIter { EmptyIter }
97} 133}
134
135pub struct NoIterMethod;
98"; 136";
99 137
100 #[test] 138 #[test]
@@ -131,43 +169,97 @@ fn main() {
131 169
132 #[test] 170 #[test]
133 fn test_for_borrowed() { 171 fn test_for_borrowed() {
134 check_assist( 172 let before = r"
135 convert_for_to_iter_for_each, 173use empty_iter::*;
136 r"
137fn main() { 174fn main() {
138 let x = vec![1, 2, 3]; 175 let x = Empty;
139 for $0v in &x { 176 for $0v in &x {
140 let a = v * 2; 177 let a = v * 2;
141 } 178 }
142}", 179}
180";
181 let before = &format!(
182 "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
183 before,
184 FamousDefs::FIXTURE,
185 EMPTY_ITER_FIXTURE
186 );
187 check_assist(
188 convert_for_to_iter_for_each,
189 before,
143 r" 190 r"
191use empty_iter::*;
144fn main() { 192fn main() {
145 let x = vec![1, 2, 3]; 193 let x = Empty;
146 x.iter().for_each(|v| { 194 x.iter().for_each(|v| {
147 let a = v * 2; 195 let a = v * 2;
148 }); 196 });
149}", 197}
198",
150 ) 199 )
151 } 200 }
152 201
153 #[test] 202 #[test]
154 fn test_for_borrowed_mut() { 203 fn test_for_borrowed_no_iter_method() {
204 let before = r"
205use empty_iter::*;
206fn main() {
207 let x = NoIterMethod;
208 for $0v in &x {
209 let a = v * 2;
210 }
211}
212";
213 let before = &format!(
214 "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
215 before,
216 FamousDefs::FIXTURE,
217 EMPTY_ITER_FIXTURE
218 );
155 check_assist( 219 check_assist(
156 convert_for_to_iter_for_each, 220 convert_for_to_iter_for_each,
221 before,
157 r" 222 r"
223use empty_iter::*;
158fn main() { 224fn main() {
159 let x = vec![1, 2, 3]; 225 let x = NoIterMethod;
226 (&x).into_iter().for_each(|v| {
227 let a = v * 2;
228 });
229}
230",
231 )
232 }
233
234 #[test]
235 fn test_for_borrowed_mut() {
236 let before = r"
237use empty_iter::*;
238fn main() {
239 let x = Empty;
160 for $0v in &mut x { 240 for $0v in &mut x {
161 *v *= 2; 241 let a = v * 2;
162 } 242 }
163}", 243}
244";
245 let before = &format!(
246 "//- /main.rs crate:main deps:core,empty_iter{}{}{}",
247 before,
248 FamousDefs::FIXTURE,
249 EMPTY_ITER_FIXTURE
250 );
251 check_assist(
252 convert_for_to_iter_for_each,
253 before,
164 r" 254 r"
255use empty_iter::*;
165fn main() { 256fn main() {
166 let x = vec![1, 2, 3]; 257 let x = Empty;
167 x.iter_mut().for_each(|v| { 258 x.iter_mut().for_each(|v| {
168 *v *= 2; 259 let a = v * 2;
169 }); 260 });
170}", 261}
262",
171 ) 263 )
172 } 264 }
173 265
@@ -195,7 +287,7 @@ fn main() {
195 } 287 }
196 288
197 #[test] 289 #[test]
198 fn test_take() { 290 fn test_already_impls_iterator() {
199 let before = r#" 291 let before = r#"
200use empty_iter::*; 292use empty_iter::*;
201fn main() { 293fn main() {