Support WhileExpr and ForExpr for add_label_to_loop

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

@@ -1,8 +1,5 @@
 use ide_db::syntax_helpers::node_ext::for_each_break_and_continue_expr;
-use syntax::{
-    T,
-    ast::{self, AstNode, HasLoopBody},
-};
+use syntax::ast::{self, AstNode, HasLoopBody};
 
 use crate::{AssistContext, AssistId, Assists};
 
@@ -28,9 +25,9 @@ use crate::{AssistContext, AssistId, Assists};
 // }
 // ```
 pub(crate) fn add_label_to_loop(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let loop_kw = ctx.find_token_syntax_at_offset(T![loop])?;
-    let loop_expr = loop_kw.parent().and_then(ast::LoopExpr::cast)?;
-    if loop_expr.label().is_some() {
+    let loop_expr = ctx.find_node_at_offset::<ast::LoopLike>()?;
+    let loop_kw = loop_expr.loop_token()?;
+    if loop_expr.label().is_some() || !loop_kw.text_range().contains_inclusive(ctx.offset()) {
         return None;
     }
 
@@ -90,6 +87,48 @@ fn main() {
         );
     }
 
+    #[test]
+    fn add_label_to_while_expr() {
+        check_assist(
+            add_label_to_loop,
+            r#"
+fn main() {
+    while$0 true {
+        break;
+        continue;
+    }
+}"#,
+            r#"
+fn main() {
+    'l: while true {
+        break 'l;
+        continue 'l;
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn add_label_to_for_expr() {
+        check_assist(
+            add_label_to_loop,
+            r#"
+fn main() {
+    for$0 _ in 0..5 {
+        break;
+        continue;
+    }
+}"#,
+            r#"
+fn main() {
+    'l: for _ in 0..5 {
+        break 'l;
+        continue 'l;
+    }
+}"#,
+        );
+    }
+
     #[test]
     fn add_label_to_outer_loop() {
         check_assist(
@@ -158,6 +197,31 @@ fn main() {
         break 'l;
         continue 'l;
     }
+}"#,
+        );
+    }
+
+    #[test]
+    fn do_not_add_label_if_outside_keyword() {
+        check_assist_not_applicable(
+            add_label_to_loop,
+            r#"
+fn main() {
+    'l: loop {$0
+        break 'l;
+        continue 'l;
+    }
+}"#,
+        );
+
+        check_assist_not_applicable(
+            add_label_to_loop,
+            r#"
+fn main() {
+    'l: while true {$0
+        break 'l;
+        continue 'l;
+    }
 }"#,
         );
     }

crates/syntax/src/ast.rs

@@ -25,8 +25,9 @@ pub use self::{
     expr_ext::{ArrayExprKind, BlockModifier, CallableExpr, ElseBranch, LiteralKind},
     generated::{nodes::*, tokens::*},
     node_ext::{
-        AttrKind, FieldKind, Macro, NameLike, NameOrNameRef, PathSegmentKind, SelfParamKind,
-        SlicePatComponents, StructKind, TypeBoundKind, TypeOrConstParam, VisibilityKind,
+        AttrKind, FieldKind, LoopLike, Macro, NameLike, NameOrNameRef, PathSegmentKind,
+        SelfParamKind, SlicePatComponents, StructKind, TypeBoundKind, TypeOrConstParam,
+        VisibilityKind,
     },
     operators::{ArithOp, BinaryOp, CmpOp, LogicOp, Ordering, RangeOp, UnaryOp},
     token_ext::{

crates/syntax/src/ast/node_ext.rs

@@ -880,6 +880,53 @@ impl AstNode for TypeOrConstParam {
 
 impl HasAttrs for TypeOrConstParam {}
 
+pub enum LoopLike {
+    ForExpr(ast::ForExpr),
+    LoopExpr(ast::LoopExpr),
+    WhileExpr(ast::WhileExpr),
+}
+
+impl LoopLike {
+    pub fn loop_token(&self) -> Option<SyntaxToken> {
+        match self {
+            LoopLike::ForExpr(it) => it.for_token(),
+            LoopLike::LoopExpr(it) => it.loop_token(),
+            LoopLike::WhileExpr(it) => it.while_token(),
+        }
+    }
+}
+
+impl AstNode for LoopLike {
+    fn can_cast(kind: SyntaxKind) -> bool
+    where
+        Self: Sized,
+    {
+        matches!(kind, SyntaxKind::FOR_EXPR | SyntaxKind::LOOP_EXPR | SyntaxKind::WHILE_EXPR)
+    }
+
+    fn cast(syntax: SyntaxNode) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        match syntax.kind() {
+            SyntaxKind::FOR_EXPR => ast::ForExpr::cast(syntax).map(Self::ForExpr),
+            SyntaxKind::LOOP_EXPR => ast::LoopExpr::cast(syntax).map(Self::LoopExpr),
+            SyntaxKind::WHILE_EXPR => ast::WhileExpr::cast(syntax).map(Self::WhileExpr),
+            _ => None,
+        }
+    }
+
+    fn syntax(&self) -> &SyntaxNode {
+        match self {
+            LoopLike::ForExpr(it) => it.syntax(),
+            LoopLike::LoopExpr(it) => it.syntax(),
+            LoopLike::WhileExpr(it) => it.syntax(),
+        }
+    }
+}
+
+impl ast::HasLoopBody for LoopLike {}
+
 pub enum VisibilityKind {
     In(ast::Path),
     PubCrate,