diff --git a/crates/ide-assists/src/handlers/add_label_to_loop.rs b/crates/ide-assists/src/handlers/add_label_to_loop.rs index d2b903447133..6a9c190e350d 100644 --- a/crates/ide-assists/src/handlers/add_label_to_loop.rs +++ b/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::()?; + 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; + } }"#, ); } diff --git a/crates/syntax/src/ast.rs b/crates/syntax/src/ast.rs index aea99a4389b9..e0b14531d770 100644 --- a/crates/syntax/src/ast.rs +++ b/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::{ diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs index af741d100f68..55c45a6ea852 100644 --- a/crates/syntax/src/ast/node_ext.rs +++ b/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 { + 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 + 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,