fix Display scope inlay hints after closing brace for more types of blocks #18833

crates/ide/src/inlay_hints/closing_brace.rs

@@ -16,6 +16,8 @@ use crate::{
     inlay_hints::LazyProperty,
 };
 
+const SNIPPET_DEFAULT_LIMIT: usize = 32;
+
 pub(super) fn hints(
     acc: &mut Vec<InlayHint>,
     sema: &Semantics<'_, RootDatabase>,
@@ -60,6 +62,11 @@ pub(super) fn hints(
 
         let module = ast::Module::cast(list.syntax().parent()?)?;
         (format!("mod {}", module.name()?), module.name().map(name))
+    } else if let Some(match_arm_list) = ast::MatchArmList::cast(node.clone()) {
+        closing_token = match_arm_list.r_curly_token()?;
+
+        let match_expr = ast::MatchExpr::cast(match_arm_list.syntax().parent()?)?;
+        (format_match_label(&match_expr, config), None)
     } else if let Some(label) = ast::Label::cast(node.clone()) {
         // in this case, `ast::Label` could be seen as a part of `ast::BlockExpr`
         // the actual number of lines in this case should be the line count of the parent BlockExpr,
@@ -91,7 +98,7 @@ pub(super) fn hints(
         match_ast! {
             match parent {
                 ast::Fn(it) => {
-                    (format!("fn {}", it.name()?), it.name().map(name))
+                    (format!("{}fn {}", fn_qualifiers(&it), it.name()?), it.name().map(name))
                 },
                 ast::Static(it) => (format!("static {}", it.name()?), it.name().map(name)),
                 ast::Const(it) => {
@@ -101,6 +108,35 @@ pub(super) fn hints(
                         (format!("const {}", it.name()?), it.name().map(name))
                     }
                 },
+                ast::LoopExpr(loop_expr) => {
+                    if loop_expr.label().is_some() {
+                        return None;
+                    }
+                    ("loop".into(), None)
+                },
+                ast::WhileExpr(while_expr) => {
+                    if while_expr.label().is_some() {
+                        return None;
+                    }
+                    (keyword_with_condition("while", while_expr.condition(), config), None)
+                },
+                ast::ForExpr(for_expr) => {
+                    if for_expr.label().is_some() {
+                        return None;
+                    }
+                    (format_for_label(&for_expr, config), None)
+                },
+                ast::IfExpr(if_expr) => {
+                    let label = label_for_if_block(&if_expr, &block, config)?;
+                    (label, None)
+                },
+                ast::LetElse(let_else) => (format_let_else_label(&let_else, config), None),
+                ast::MatchArm(match_arm) => {
+                    if !block_is_match_arm_body(&match_arm, &block) {
+                        return None;
+                    }
+                    (format_match_arm_label(&match_arm, config), None)
+                },
                 _ => return None,
             }
         }
@@ -154,11 +190,171 @@ pub(super) fn hints(
     None
 }
 
+fn fn_qualifiers(func: &ast::Fn) -> String {
+    let mut qualifiers = String::new();
+    if func.const_token().is_some() {
+        qualifiers.push_str("const ");
+    }
+    if func.async_token().is_some() {
+        qualifiers.push_str("async ");
+    }
+    if func.unsafe_token().is_some() {
+        qualifiers.push_str("unsafe ");
+    }
+    qualifiers
+}
+
+fn keyword_with_condition(
+    keyword: &str,
+    condition: Option<ast::Expr>,
+    config: &InlayHintsConfig<'_>,
+) -> String {
+    if let Some(expr) = condition {
+        let snippet = snippet_from_node(expr.syntax(), config);
+        if !snippet.is_empty() {
+            return format!("{keyword} {snippet}");
+        }
+    }
+    keyword.to_owned()
+}
+
+fn format_for_label(for_expr: &ast::ForExpr, config: &InlayHintsConfig<'_>) -> String {
+    let pat = for_expr.pat().and_then(|pat| {
+        let snippet = snippet_from_node(pat.syntax(), config);
+        (!snippet.is_empty()).then_some(snippet)
+    });
+    let iterable = for_expr.iterable().and_then(|expr| {
+        let snippet = snippet_from_node(expr.syntax(), config);
+        (!snippet.is_empty()).then_some(snippet)
+    });
+
+    match (pat, iterable) {
+        (Some(pat), Some(iterable)) => format!("for {pat} in {iterable}"),
+        (Some(pat), None) => format!("for {pat}"),
+        (None, Some(iterable)) => format!("for in {iterable}"),
+        _ => "for".into(),
+    }
+}
+
+fn format_match_label(match_expr: &ast::MatchExpr, config: &InlayHintsConfig<'_>) -> String {
+    match match_expr.expr() {
+        Some(expr) => {
+            let snippet = snippet_from_node(expr.syntax(), config);
+            if snippet.is_empty() { "match".into() } else { format!("match {snippet}") }
+        }
+        None => "match".into(),
+    }
+}
+
+fn format_match_arm_label(match_arm: &ast::MatchArm, config: &InlayHintsConfig<'_>) -> String {
+    let mut parts = Vec::new();
+    if let Some(pat) = match_arm.pat() {
+        let snippet = snippet_from_node(pat.syntax(), config);
+        if !snippet.is_empty() {
+            parts.push(snippet);
+        }
+    }
+    if let Some(guard) = match_arm.guard().and_then(|guard| guard.condition()) {
+        let snippet = snippet_from_node(guard.syntax(), config);
+        if !snippet.is_empty() {
+            parts.push(format!("if {snippet}"));
+        }
+    }
+
+    if parts.is_empty() { "match arm".into() } else { format!("{} =>", parts.join(" ")) }
+}
+
+fn block_is_match_arm_body(arm: &ast::MatchArm, block: &ast::BlockExpr) -> bool {
+    arm.expr()
+        .and_then(|expr| match expr {
+            ast::Expr::BlockExpr(expr_block) => Some(expr_block.syntax() == block.syntax()),
+            _ => None,
+        })
+        .unwrap_or(false)
+}
+
+fn label_for_if_block(
+    if_expr: &ast::IfExpr,
+    block: &ast::BlockExpr,
+    config: &InlayHintsConfig<'_>,
+) -> Option<String> {
+    if if_expr.then_branch().is_some_and(|then_branch| then_branch.syntax() == block.syntax()) {
+        Some(keyword_with_condition("if", if_expr.condition(), config))
+    } else if matches!(
+        if_expr.else_branch(),
+        Some(ast::ElseBranch::Block(else_block)) if else_block.syntax() == block.syntax()
+    ) {
+        Some("else".into())
+    } else {
+        None
+    }
+}
+
+fn format_let_else_label(let_else: &ast::LetElse, config: &InlayHintsConfig<'_>) -> String {
+    let let_stmt = let_else.syntax().parent().and_then(ast::LetStmt::cast);
+    let Some(stmt) = let_stmt else {
+        return "let ... else".into();
+    };
+
+    let pat = stmt.pat().and_then(|pat| {
+        let snippet = snippet_from_node(pat.syntax(), config);
+        (!snippet.is_empty()).then_some(snippet)
+    });
+    let initializer = stmt.initializer().and_then(|expr| {
+        let snippet = snippet_from_node(expr.syntax(), config);
+        (!snippet.is_empty()).then_some(snippet)
+    });
+
+    match (pat, initializer) {
+        (Some(pat), Some(init)) => format!("let {pat} = {init} else"),
+        (Some(pat), None) => format!("let {pat} else"),
+        _ => "let ... else".into(),
+    }
+}
+
+fn snippet_from_node(node: &SyntaxNode, config: &InlayHintsConfig<'_>) -> String {
+    let mut normalized = String::new();
+    let mut last_was_space = true;
+    for ch in node.text().to_string().chars() {
+        if ch.is_whitespace() {
+            if last_was_space {
+                continue;
+            }
+            normalized.push(' ');
+            last_was_space = true;
+        } else {
+            normalized.push(ch);
+            last_was_space = false;
+        }
+    }
+    let normalized = normalized.trim().to_owned();
+    if normalized.is_empty() {
+        return normalized;
+    }
+
+    let limit = config.max_length.unwrap_or(SNIPPET_DEFAULT_LIMIT);
+    if limit == 0 || normalized.chars().count() <= limit {
+        return normalized;
+    }
+
+    let mut truncated = String::new();
+    for (used, ch) in normalized.chars().enumerate() {
+        if used == limit {
+            break;
+        }
+        truncated.push(ch);
+    }
+    truncated.push_str("...");
+    truncated
+}
+
 #[cfg(test)]
 mod tests {
+    use expect_test::expect;
+
     use crate::{
         InlayHintsConfig,
-        inlay_hints::tests::{DISABLED_CONFIG, check_with_config},
+        inlay_hints::tests::{DISABLED_CONFIG, check_expect, check_with_config},
     };
 
     #[test]
@@ -179,6 +375,10 @@ fn h<T>(with: T, arguments: u8, ...) {
   }
 //^ fn h
 
+async fn async_fn() {
+  }
+//^ async fn async_fn
+
 trait Tr {
     fn f();
     fn g() {
@@ -260,4 +460,136 @@ fn test() {
 "#,
         );
     }
+
+    #[test]
+    fn hints_closing_brace_additional_blocks() {
+        check_expect(
+            InlayHintsConfig { closing_brace_hints_min_lines: Some(2), ..DISABLED_CONFIG },
+            r#"
+fn demo() {
+    loop {
+
+    }
+
+    while let Some(value) = next() {
+
+    }
+
+    for value in iter {
+
+    }
+
+    if cond {
+
+    }
+
+    if let Some(x) = maybe {
+
+    }
+
+    if other {
+    } else {
+
+    }
+
+    let Some(v) = maybe else {
+
+    };
+
+    match maybe {
+        Some(v) => {
+
+        }
+        value if check(value) => {
+
+        }
+        None => {}
+    }
+}
+"#,
+            expect![[r#"
+                [
+                    (
+                        364..365,
+                        [
+                            InlayHintLabelPart {
+                                text: "fn demo",
+                                linked_location: Some(
+                                    Computed(
+                                        FileRangeWrapper {
+                                            file_id: FileId(
+                                                0,
+                                            ),
+                                            range: 3..7,
+                                        },
+                                    ),
+                                ),
+                                tooltip: "",
+                            },
+                        ],
+                    ),
+                    (
+                        28..29,
+                        [
+                            "loop",
+                        ],
+                    ),
+                    (
+                        73..74,
+                        [
+                            "while let Some(value) = next()",
+                        ],
+                    ),
+                    (
+                        105..106,
+                        [
+                            "for value in iter",
+                        ],
+                    ),
+                    (
+                        127..128,
+                        [
+                            "if cond",
+                        ],
+                    ),
+                    (
+                        164..165,
+                        [
+                            "if let Some(x) = maybe",
+                        ],
+                    ),
+                    (
+                        200..201,
+                        [
+                            "else",
+                        ],
+                    ),
+                    (
+                        240..241,
+                        [
+                            "let Some(v) = maybe else",
+                        ],
+                    ),
+                    (
+                        362..363,
+                        [
+                            "match maybe",
+                        ],
+                    ),
+                    (
+                        291..292,
+                        [
+                            "Some(v) =>",
+                        ],
+                    ),
+                    (
+                        337..338,
+                        [
+                            "value if check(value) =>",
+                        ],
+                    ),
+                ]
+            "#]],
+        );
+    }
 }