Skip to content

Commit eea2149

Browse files
committed
feat: add assist to conver for_each into for loops
1 parent 21913d0 commit eea2149

File tree

2 files changed

+140
-0
lines changed

2 files changed

+140
-0
lines changed
Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
use ide_db::helpers::FamousDefs;
2+
use stdx::format_to;
3+
use syntax::{AstNode, ast::{self, ArgListOwner}};
4+
5+
use crate::{AssistContext, AssistId, AssistKind, Assists};
6+
7+
/// Assist: convert_iter_for_each_to_for
8+
//
9+
/// Converts an Iterator::for_each function into a for loop.
10+
///
11+
/// ```rust
12+
/// fn main() {
13+
/// let vec = vec![(1, 2), (2, 3), (3, 4)];
14+
/// x.iter().for_each(|(x, y)| {
15+
/// println!("x: {}, y: {}", x, y);
16+
/// })
17+
/// }
18+
/// ```
19+
/// ->
20+
/// ```rust
21+
/// fn main() {
22+
/// let vec = vec![(1, 2), (2, 3), (3, 4)];
23+
/// for (x, y) in x.iter() {
24+
/// println!("x: {}, y: {}", x, y);
25+
/// });
26+
/// }
27+
/// ```
28+
pub(crate) fn convert_iter_for_each_to_for(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
29+
let closure;
30+
31+
let total_expr = match ctx.find_node_at_offset::<ast::Expr>()? {
32+
ast::Expr::MethodCallExpr(expr) => {
33+
closure = match expr.arg_list()?.args().next()? {
34+
ast::Expr::ClosureExpr(expr) => expr,
35+
_ => { return None; }
36+
};
37+
38+
expr
39+
},
40+
ast::Expr::ClosureExpr(expr) => {
41+
closure = expr;
42+
ast::MethodCallExpr::cast(closure.syntax().ancestors().nth(2)?)?
43+
},
44+
_ => { return None; }
45+
};
46+
47+
let (total_expr, parent) = validate_method_call_expr(&ctx.sema, total_expr)?;
48+
49+
let param_list = closure.param_list()?;
50+
let param = param_list.params().next()?;
51+
let body = closure.body()?;
52+
53+
acc.add(
54+
AssistId("convert_iter_for_each_to_for", AssistKind::RefactorRewrite),
55+
"Replace this `Iterator::for_each` with a for loop",
56+
total_expr.syntax().text_range(),
57+
|builder| {
58+
let mut buf = String::new();
59+
60+
format_to!(buf, "for {} in {} ", param, parent);
61+
62+
match body {
63+
ast::Expr::BlockExpr(body) => format_to!(buf, "{}", body),
64+
_ => format_to!(buf, "{{\n{}\n}}", body)
65+
}
66+
67+
builder.replace(total_expr.syntax().text_range(), buf)
68+
},
69+
)
70+
}
71+
72+
fn validate_method_call_expr(
73+
sema: &hir::Semantics<ide_db::RootDatabase>,
74+
expr: ast::MethodCallExpr,
75+
) -> Option<(ast::Expr, ast::Expr)> {
76+
if expr.name_ref()?.text() != "for_each" {
77+
return None;
78+
}
79+
80+
let expr = ast::Expr::MethodCallExpr(expr);
81+
let parent = ast::Expr::cast(expr.syntax().first_child()?)?;
82+
83+
let it_type = sema.type_of_expr(&parent)?;
84+
let module = sema.scope(parent.syntax()).module()?;
85+
let krate = module.krate();
86+
87+
let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
88+
it_type.impls_trait(sema.db, iter_trait, &[]).then(|| (expr, parent))
89+
}
90+
91+
#[cfg(test)]
92+
mod tests {
93+
use crate::tests::check_assist;
94+
95+
use super::*;
96+
97+
#[test]
98+
fn test_for_each_in_method() {
99+
check_assist(
100+
convert_iter_for_each_to_for,
101+
r"
102+
fn main() {
103+
let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
104+
x.iter().$0for_each(|(x, y)| {
105+
dbg!(x, y)
106+
});
107+
}",
108+
r"
109+
fn main() {
110+
let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
111+
for (x, y) in x.iter() {
112+
dbg!(x, y)
113+
};
114+
}",
115+
)
116+
}
117+
118+
#[test]
119+
fn test_for_each_in_closure() {
120+
check_assist(
121+
convert_iter_for_each_to_for,
122+
r"
123+
fn main() {
124+
let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
125+
x.iter().for_each($0|(x, y)| {
126+
dbg!(x, y)
127+
});
128+
}",
129+
r"
130+
fn main() {
131+
let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
132+
for (x, y) in x.iter() {
133+
dbg!(x, y)
134+
};
135+
}",
136+
)
137+
}
138+
}

crates/ide_assists/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ mod handlers {
116116
mod change_visibility;
117117
mod convert_integer_literal;
118118
mod convert_comment_block;
119+
mod convert_iter_for_each_to_for;
119120
mod early_return;
120121
mod expand_glob_import;
121122
mod extract_function;
@@ -181,6 +182,7 @@ mod handlers {
181182
change_visibility::change_visibility,
182183
convert_integer_literal::convert_integer_literal,
183184
convert_comment_block::convert_comment_block,
185+
convert_iter_for_each_to_for::convert_iter_for_each_to_for,
184186
early_return::convert_to_guarded_return,
185187
expand_glob_import::expand_glob_import,
186188
extract_struct_from_enum_variant::extract_struct_from_enum_variant,

0 commit comments

Comments
 (0)