minor: code review tweak

Author: gohome001

src/tools/rust-analyzer/crates/ide/src/highlight_related.rs

@@ -89,7 +89,7 @@ pub(crate) fn highlight_related(
         T![break] | T![loop] | T![while] | T![continue] if config.break_points => {
             highlight_break_points(sema, token).remove(&file_id)
         }
-        T![unsafe] if token.parent_ancestors().find_map(ast::BlockExpr::cast).is_some() => {
+        T![unsafe] if token.parent().and_then(ast::BlockExpr::cast).is_some() => {
             highlight_unsafe_points(sema, token).remove(&file_id)
         }
         T![|] if config.closure_captures => {
@@ -715,7 +715,7 @@ pub(crate) fn highlight_unsafe_points(
 ) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
     fn hl(
         sema: &Semantics<'_, RootDatabase>,
-        unsafe_token: Option<SyntaxToken>,
+        unsafe_token: &SyntaxToken,
         block_expr: Option<ast::BlockExpr>,
     ) -> Option<FxHashMap<EditionedFileId, Vec<HighlightedRange>>> {
         let mut highlights: FxHashMap<EditionedFileId, Vec<_>> = FxHashMap::default();
@@ -728,39 +728,23 @@ pub(crate) fn highlight_unsafe_points(
         };
 
         // highlight unsafe keyword itself
-        let unsafe_token = unsafe_token?;
         let unsafe_token_file_id = sema.hir_file_for(&unsafe_token.parent()?);
         push_to_highlights(unsafe_token_file_id, Some(unsafe_token.text_range()));
 
+        // highlight unsafe operations
         if let Some(block) = block_expr {
-            if let Some(node) = block.syntax().ancestors().find(|n| ast::Fn::can_cast(n.kind())) {
-                if let Some(function) = ast::Fn::cast(node) {
-                    // highlight unsafe keyword of the function
-                    if let Some(unsafe_token) = function.unsafe_token() {
-                        push_to_highlights(unsafe_token_file_id, Some(unsafe_token.text_range()));
-                    }
-                    // highlight unsafe operations
-                    if let Some(f) = sema.to_def(&function) {
-                        let unsafe_ops = sema.get_unsafe_ops(f.into());
-                        for unsafe_op in unsafe_ops {
-                            push_to_highlights(
-                                unsafe_op.file_id,
-                                Some(unsafe_op.value.text_range()),
-                            );
-                        }
-                    }
+            if let Some(body) = sema.body_for(InFile::new(unsafe_token_file_id, block.syntax())) {
+                let unsafe_ops = sema.get_unsafe_ops(body);
+                for unsafe_op in unsafe_ops {
+                    push_to_highlights(unsafe_op.file_id, Some(unsafe_op.value.text_range()));
                 }
             }
         }
 
         Some(highlights)
     }
 
-    let Some(block_expr) = token.parent().and_then(ast::BlockExpr::cast) else {
-        return FxHashMap::default();
-    };
-
-    hl(sema, Some(token), Some(block_expr)).unwrap_or_default()
+    hl(sema, &token, token.parent().and_then(ast::BlockExpr::cast)).unwrap_or_default()
 }
 
 #[cfg(test)]