Skip to content

Commit 9e9ea26

Browse files
committed
Add ide-assist: convert_range_for_to_while
Convert for each range into while loop. ```rust fn foo() { $0for i in 3..7 { foo(i); } } ``` -> ```rust fn foo() { let mut i = 3; while i < 7 { foo(i); i += 1; } } ```
1 parent 1625b1e commit 9e9ea26

File tree

5 files changed

+294
-1
lines changed

5 files changed

+294
-1
lines changed
Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,254 @@
1+
use ide_db::assists::AssistId;
2+
use itertools::Itertools;
3+
use syntax::{
4+
AstNode, T,
5+
algo::previous_non_trivia_token,
6+
ast::{
7+
self, HasArgList, HasLoopBody, HasName, RangeItem, edit::AstNodeEdit, make,
8+
syntax_factory::SyntaxFactory,
9+
},
10+
syntax_editor::{Element, Position},
11+
};
12+
13+
use crate::assist_context::{AssistContext, Assists};
14+
15+
// Assist: convert_range_for_to_while
16+
//
17+
// Convert for each range into while loop.
18+
//
19+
// ```
20+
// fn foo() {
21+
// $0for i in 3..7 {
22+
// foo(i);
23+
// }
24+
// }
25+
// ```
26+
// ->
27+
// ```
28+
// fn foo() {
29+
// let mut i = 3;
30+
// while i < 7 {
31+
// foo(i);
32+
// i += 1;
33+
// }
34+
// }
35+
// ```
36+
pub(crate) fn convert_range_for_to_while(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
37+
let for_kw = ctx.find_token_syntax_at_offset(T![for])?;
38+
let for_ = ast::ForExpr::cast(for_kw.parent()?)?;
39+
let ast::Pat::IdentPat(pat) = for_.pat()? else { return None };
40+
let iterable = for_.iterable()?;
41+
let (start, end, step, inclusive) = extract_range(&iterable)?;
42+
let name = pat.name()?;
43+
let body = for_.loop_body()?;
44+
let last = previous_non_trivia_token(body.stmt_list()?.r_curly_token()?)?;
45+
46+
acc.add(
47+
AssistId::refactor("convert_range_for_to_while"),
48+
"Replace to while or loop",
49+
for_.syntax().text_range(),
50+
|builder| {
51+
let mut edit = builder.make_editor(for_.syntax());
52+
let make = SyntaxFactory::with_mappings();
53+
54+
let indent = for_.indent_level();
55+
let pat = make.ident_pat(pat.ref_token().is_some(), true, name.clone());
56+
let let_stmt = make.let_stmt(pat.into(), None, Some(start));
57+
edit.insert_all(
58+
Position::before(for_.syntax()),
59+
vec![
60+
let_stmt.syntax().syntax_element(),
61+
make.whitespace(&format!("\n{}", indent)).syntax_element(),
62+
],
63+
);
64+
65+
let mut elements = vec![];
66+
67+
let var_expr = make.expr_path(make.ident_path(&name.text()));
68+
let op = ast::BinaryOp::CmpOp(ast::CmpOp::Ord {
69+
ordering: ast::Ordering::Less,
70+
strict: !inclusive,
71+
});
72+
if let Some(end) = end {
73+
elements.extend([
74+
make.token(T![while]).syntax_element(),
75+
make.whitespace(" ").syntax_element(),
76+
make.expr_bin(var_expr.clone(), op, end).syntax().syntax_element(),
77+
]);
78+
} else {
79+
elements.push(make.token(T![loop]).syntax_element());
80+
}
81+
82+
edit.replace_all(
83+
for_kw.syntax_element()..=iterable.syntax().syntax_element(),
84+
elements,
85+
);
86+
87+
let op = ast::BinaryOp::Assignment { op: Some(ast::ArithOp::Add) };
88+
edit.insert_all(
89+
Position::after(last),
90+
vec![
91+
make.whitespace(&format!("\n{}", indent + 1)).syntax_element(),
92+
make.expr_bin(var_expr, op, step).syntax().syntax_element(),
93+
make.token(T![;]).syntax_element(),
94+
],
95+
);
96+
97+
edit.add_mappings(make.finish_with_mappings());
98+
builder.add_file_edits(ctx.vfs_file_id(), edit);
99+
},
100+
)
101+
}
102+
103+
fn extract_range(iterable: &ast::Expr) -> Option<(ast::Expr, Option<ast::Expr>, ast::Expr, bool)> {
104+
Some(match iterable {
105+
ast::Expr::ParenExpr(expr) => extract_range(&expr.expr()?)?,
106+
ast::Expr::RangeExpr(range) => {
107+
let inclusive = range.op_kind()? == ast::RangeOp::Inclusive;
108+
(range.start()?, range.end(), make::expr_literal("1").into(), inclusive)
109+
}
110+
ast::Expr::MethodCallExpr(call) if call.name_ref()?.text() == "step_by" => {
111+
let [step] = call.arg_list()?.args().collect_array()?;
112+
let (start, end, _, inclusive) = extract_range(&call.receiver()?)?;
113+
(start, end, step, inclusive)
114+
}
115+
_ => return None,
116+
})
117+
}
118+
119+
#[cfg(test)]
120+
mod tests {
121+
use crate::tests::{check_assist, check_assist_not_applicable};
122+
123+
use super::*;
124+
125+
#[test]
126+
fn test_convert_range_for_to_while() {
127+
check_assist(
128+
convert_range_for_to_while,
129+
"
130+
fn foo() {
131+
$0for i in 3..7 {
132+
foo(i);
133+
}
134+
}
135+
",
136+
"
137+
fn foo() {
138+
let mut i = 3;
139+
while i < 7 {
140+
foo(i);
141+
i += 1;
142+
}
143+
}
144+
",
145+
);
146+
}
147+
148+
#[test]
149+
fn test_convert_range_for_to_while_no_end_bound() {
150+
check_assist(
151+
convert_range_for_to_while,
152+
"
153+
fn foo() {
154+
$0for i in 3.. {
155+
foo(i);
156+
}
157+
}
158+
",
159+
"
160+
fn foo() {
161+
let mut i = 3;
162+
loop {
163+
foo(i);
164+
i += 1;
165+
}
166+
}
167+
",
168+
);
169+
}
170+
171+
#[test]
172+
fn test_convert_range_for_to_while_with_mut_binding() {
173+
check_assist(
174+
convert_range_for_to_while,
175+
"
176+
fn foo() {
177+
$0for mut i in 3..7 {
178+
foo(i);
179+
}
180+
}
181+
",
182+
"
183+
fn foo() {
184+
let mut i = 3;
185+
while i < 7 {
186+
foo(i);
187+
i += 1;
188+
}
189+
}
190+
",
191+
);
192+
}
193+
194+
#[test]
195+
fn test_convert_range_for_to_while_with_label() {
196+
check_assist(
197+
convert_range_for_to_while,
198+
"
199+
fn foo() {
200+
'a: $0for mut i in 3..7 {
201+
foo(i);
202+
}
203+
}
204+
",
205+
"
206+
fn foo() {
207+
let mut i = 3;
208+
'a: while i < 7 {
209+
foo(i);
210+
i += 1;
211+
}
212+
}
213+
",
214+
);
215+
}
216+
217+
#[test]
218+
fn test_convert_range_for_to_while_step_by() {
219+
check_assist(
220+
convert_range_for_to_while,
221+
"
222+
fn foo() {
223+
$0for mut i in (3..7).step_by(2) {
224+
foo(i);
225+
}
226+
}
227+
",
228+
"
229+
fn foo() {
230+
let mut i = 3;
231+
while i < 7 {
232+
foo(i);
233+
i += 2;
234+
}
235+
}
236+
",
237+
);
238+
}
239+
240+
#[test]
241+
fn test_convert_range_for_to_while_not_applicable_non_range() {
242+
check_assist_not_applicable(
243+
convert_range_for_to_while,
244+
"
245+
fn foo() {
246+
let ident = 3..7;
247+
$0for mut i in ident {
248+
foo(i);
249+
}
250+
}
251+
",
252+
);
253+
}
254+
}

src/tools/rust-analyzer/crates/ide-assists/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,7 @@ mod handlers {
131131
mod convert_match_to_let_else;
132132
mod convert_named_struct_to_tuple_struct;
133133
mod convert_nested_function_to_closure;
134+
mod convert_range_for_to_while;
134135
mod convert_to_guarded_return;
135136
mod convert_tuple_return_type_to_struct;
136137
mod convert_tuple_struct_to_named_struct;
@@ -268,6 +269,7 @@ mod handlers {
268269
convert_match_to_let_else::convert_match_to_let_else,
269270
convert_named_struct_to_tuple_struct::convert_named_struct_to_tuple_struct,
270271
convert_nested_function_to_closure::convert_nested_function_to_closure,
272+
convert_range_for_to_while::convert_range_for_to_while,
271273
convert_to_guarded_return::convert_to_guarded_return,
272274
convert_tuple_return_type_to_struct::convert_tuple_return_type_to_struct,
273275
convert_tuple_struct_to_named_struct::convert_tuple_struct_to_named_struct,

src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -731,6 +731,29 @@ fn main() {
731731
)
732732
}
733733

734+
#[test]
735+
fn doctest_convert_range_for_to_while() {
736+
check_doc_test(
737+
"convert_range_for_to_while",
738+
r#####"
739+
fn foo() {
740+
$0for i in 3..7 {
741+
foo(i);
742+
}
743+
}
744+
"#####,
745+
r#####"
746+
fn foo() {
747+
let mut i = 3;
748+
while i < 7 {
749+
foo(i);
750+
i += 1;
751+
}
752+
}
753+
"#####,
754+
)
755+
}
756+
734757
#[test]
735758
fn doctest_convert_to_guarded_return() {
736759
check_doc_test(

src/tools/rust-analyzer/crates/syntax/src/ast/make.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1355,7 +1355,7 @@ pub mod tokens {
13551355

13561356
pub(super) static SOURCE_FILE: LazyLock<Parse<SourceFile>> = LazyLock::new(|| {
13571357
SourceFile::parse(
1358-
"use crate::foo; const C: <()>::Item = ( true && true , true || true , 1 != 1, 2 == 2, 3 < 3, 4 <= 4, 5 > 5, 6 >= 6, !true, *p, &p , &mut p, async { let _ @ [] })\n;\n\nunsafe impl A for B where: {}",
1358+
"use crate::foo; const C: <()>::Item = ( true && true , true || true , 1 != 1, 2 == 2, 3 < 3, 4 <= 4, 5 > 5, 6 >= 6, !true, *p, &p , &mut p, async { let _ @ [] }, while loop {} {})\n;\n\nunsafe impl A for B where: {}",
13591359
Edition::CURRENT,
13601360
)
13611361
});

src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -644,6 +644,20 @@ impl SyntaxFactory {
644644
ast
645645
}
646646

647+
pub fn expr_loop(&self, body: ast::BlockExpr) -> ast::LoopExpr {
648+
let ast::Expr::LoopExpr(ast) = make::expr_loop(body.clone()).clone_for_update() else {
649+
unreachable!()
650+
};
651+
652+
if let Some(mut mapping) = self.mappings() {
653+
let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
654+
builder.map_node(body.syntax().clone(), ast.loop_body().unwrap().syntax().clone());
655+
builder.finish(&mut mapping);
656+
}
657+
658+
ast
659+
}
660+
647661
pub fn expr_while_loop(&self, condition: ast::Expr, body: ast::BlockExpr) -> ast::WhileExpr {
648662
let ast = make::expr_while_loop(condition.clone(), body.clone()).clone_for_update();
649663

0 commit comments

Comments
 (0)