Skip to content

Commit dc14c43

Browse files
bors[bot]mattyhall
andauthored
Merge #7741
7741: Add convert_for_to_iter_for_each assist r=mattyhall a=mattyhall Implements one direction of #7681 I wonder if this tries to guess too much at the right thing here. A common pattern is: ```rust let col = vec![1, 2, 3]; for v in &mut col { *v *= 2; } // equivalent to: col.iter_mut().for_each(|v| *v *= 2); ``` I've tried to detect this case by checking if the expression after the `in` is a (mutable) reference and if not inserting iter()/iter_mut(). This is just a convention used in the stdlib however, so could sometimes be wrong. I'd be happy to make an improvement for this, but not sure what would be best. A few options spring to mind: 1. Only allow this for types that are known to have iter/iter_mut (ie stdlib types) 2. Try to check if iter/iter_mut exists and they return the right iterator type 3. Don't try to do this and just add `.into_iter()` to whatever is after `in` Co-authored-by: Matt Hall <matthew@quickbeam.me.uk>
2 parents 0537510 + a28e862 commit dc14c43

File tree

4 files changed

+327
-0
lines changed

4 files changed

+327
-0
lines changed

crates/hir_expand/src/name.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,7 @@ pub mod known {
189189
// Components of known path (function name)
190190
filter_map,
191191
next,
192+
iter_mut,
192193
// Builtin macros
193194
file,
194195
column,
Lines changed: 301 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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+
}

crates/ide_assists/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ mod handlers {
114114
mod apply_demorgan;
115115
mod auto_import;
116116
mod change_visibility;
117+
mod convert_for_to_iter_for_each;
117118
mod convert_integer_literal;
118119
mod early_return;
119120
mod expand_glob_import;
@@ -175,6 +176,7 @@ mod handlers {
175176
apply_demorgan::apply_demorgan,
176177
auto_import::auto_import,
177178
change_visibility::change_visibility,
179+
convert_for_to_iter_for_each::convert_for_to_iter_for_each,
178180
convert_integer_literal::convert_integer_literal,
179181
early_return::convert_to_guarded_return,
180182
expand_glob_import::expand_glob_import,

crates/ide_assists/src/tests/generated.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,29 @@ pub(crate) fn frobnicate() {}
192192
)
193193
}
194194

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+
195218
#[test]
196219
fn doctest_convert_integer_literal() {
197220
check_doc_test(

0 commit comments

Comments
 (0)