Skip to content

Commit cfbd307

Browse files
committed
Support WhileExpr and ForExpr for add_label_to_loop
Example --- ```rust fn main() { for$0 _ in 0..5 { break; continue; } } ``` **Before this PR** Assist not applicable **After this PR** ```rust fn main() { 'l: for _ in 0..5 { break 'l; continue 'l; } } ```
1 parent c937fcc commit cfbd307

File tree

3 files changed

+120
-5
lines changed

3 files changed

+120
-5
lines changed

crates/ide-assists/src/handlers/add_label_to_loop.rs

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,9 @@ use crate::{AssistContext, AssistId, Assists};
3535
// }
3636
// ```
3737
pub(crate) fn add_label_to_loop(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
38-
let loop_kw = ctx.find_token_syntax_at_offset(T![loop])?;
39-
let loop_expr = loop_kw.parent().and_then(ast::LoopExpr::cast)?;
40-
if loop_expr.label().is_some() {
38+
let loop_expr = ctx.find_node_at_offset::<ast::LoopLike>()?;
39+
let loop_kw = loop_expr.loop_token()?;
40+
if loop_expr.label().is_some() || !loop_kw.text_range().contains_inclusive(ctx.offset()) {
4141
return None;
4242
}
4343

@@ -123,6 +123,48 @@ fn main() {
123123
);
124124
}
125125

126+
#[test]
127+
fn add_label_to_while_expr() {
128+
check_assist(
129+
add_label_to_loop,
130+
r#"
131+
fn main() {
132+
while$0 true {
133+
break;
134+
continue;
135+
}
136+
}"#,
137+
r#"
138+
fn main() {
139+
'l: while true {
140+
break 'l;
141+
continue 'l;
142+
}
143+
}"#,
144+
);
145+
}
146+
147+
#[test]
148+
fn add_label_to_for_expr() {
149+
check_assist(
150+
add_label_to_loop,
151+
r#"
152+
fn main() {
153+
for$0 _ in 0..5 {
154+
break;
155+
continue;
156+
}
157+
}"#,
158+
r#"
159+
fn main() {
160+
'l: for _ in 0..5 {
161+
break 'l;
162+
continue 'l;
163+
}
164+
}"#,
165+
);
166+
}
167+
126168
#[test]
127169
fn add_label_to_outer_loop() {
128170
check_assist(
@@ -191,6 +233,31 @@ fn main() {
191233
break 'l;
192234
continue 'l;
193235
}
236+
}"#,
237+
);
238+
}
239+
240+
#[test]
241+
fn do_not_add_label_if_outside_keyword() {
242+
check_assist_not_applicable(
243+
add_label_to_loop,
244+
r#"
245+
fn main() {
246+
'l: loop {$0
247+
break 'l;
248+
continue 'l;
249+
}
250+
}"#,
251+
);
252+
253+
check_assist_not_applicable(
254+
add_label_to_loop,
255+
r#"
256+
fn main() {
257+
'l: while true {$0
258+
break 'l;
259+
continue 'l;
260+
}
194261
}"#,
195262
);
196263
}

crates/syntax/src/ast.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,9 @@ pub use self::{
2525
expr_ext::{ArrayExprKind, BlockModifier, CallableExpr, ElseBranch, LiteralKind},
2626
generated::{nodes::*, tokens::*},
2727
node_ext::{
28-
AttrKind, FieldKind, Macro, NameLike, NameOrNameRef, PathSegmentKind, SelfParamKind,
29-
SlicePatComponents, StructKind, TypeBoundKind, TypeOrConstParam, VisibilityKind,
28+
AttrKind, FieldKind, LoopLike, Macro, NameLike, NameOrNameRef, PathSegmentKind,
29+
SelfParamKind, SlicePatComponents, StructKind, TypeBoundKind, TypeOrConstParam,
30+
VisibilityKind,
3031
},
3132
operators::{ArithOp, BinaryOp, CmpOp, LogicOp, Ordering, RangeOp, UnaryOp},
3233
token_ext::{

crates/syntax/src/ast/node_ext.rs

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -880,6 +880,53 @@ impl AstNode for TypeOrConstParam {
880880

881881
impl HasAttrs for TypeOrConstParam {}
882882

883+
pub enum LoopLike {
884+
ForExpr(ast::ForExpr),
885+
LoopExpr(ast::LoopExpr),
886+
WhileExpr(ast::WhileExpr),
887+
}
888+
889+
impl LoopLike {
890+
pub fn loop_token(&self) -> Option<SyntaxToken> {
891+
match self {
892+
LoopLike::ForExpr(it) => it.for_token(),
893+
LoopLike::LoopExpr(it) => it.loop_token(),
894+
LoopLike::WhileExpr(it) => it.while_token(),
895+
}
896+
}
897+
}
898+
899+
impl AstNode for LoopLike {
900+
fn can_cast(kind: SyntaxKind) -> bool
901+
where
902+
Self: Sized,
903+
{
904+
matches!(kind, SyntaxKind::FOR_EXPR | SyntaxKind::LOOP_EXPR | SyntaxKind::WHILE_EXPR)
905+
}
906+
907+
fn cast(syntax: SyntaxNode) -> Option<Self>
908+
where
909+
Self: Sized,
910+
{
911+
match syntax.kind() {
912+
SyntaxKind::FOR_EXPR => ast::ForExpr::cast(syntax).map(Self::ForExpr),
913+
SyntaxKind::LOOP_EXPR => ast::LoopExpr::cast(syntax).map(Self::LoopExpr),
914+
SyntaxKind::WHILE_EXPR => ast::WhileExpr::cast(syntax).map(Self::WhileExpr),
915+
_ => None,
916+
}
917+
}
918+
919+
fn syntax(&self) -> &SyntaxNode {
920+
match self {
921+
LoopLike::ForExpr(it) => it.syntax(),
922+
LoopLike::LoopExpr(it) => it.syntax(),
923+
LoopLike::WhileExpr(it) => it.syntax(),
924+
}
925+
}
926+
}
927+
928+
impl ast::HasLoopBody for LoopLike {}
929+
883930
pub enum VisibilityKind {
884931
In(ast::Path),
885932
PubCrate,

0 commit comments

Comments
 (0)