In highlight_related, when on an unsafe block, don't highlight unsafe operations of other unsafe blocks

Author: ChayimFriedman2

src/tools/rust-analyzer/crates/hir-ty/src/diagnostics/unsafe_check.rs

@@ -119,11 +119,11 @@ pub fn unsafe_operations(
     def: DefWithBodyId,
     body: &Body,
     current: ExprId,
-    callback: &mut dyn FnMut(InsideUnsafeBlock),
+    callback: &mut dyn FnMut(ExprOrPatId, InsideUnsafeBlock),
 ) {
     let mut visitor_callback = |diag| {
-        if let UnsafeDiagnostic::UnsafeOperation { inside_unsafe_block, .. } = diag {
-            callback(inside_unsafe_block);
+        if let UnsafeDiagnostic::UnsafeOperation { inside_unsafe_block, node, .. } = diag {
+            callback(node, inside_unsafe_block);
         }
     };
     let mut visitor = UnsafeVisitor::new(db, infer, body, def, &mut visitor_callback);

src/tools/rust-analyzer/crates/hir/src/semantics.rs

@@ -28,13 +28,13 @@ use hir_expand::{
     mod_path::{ModPath, PathKind},
     name::AsName,
 };
-use hir_ty::diagnostics::unsafe_operations_for_body;
+use hir_ty::diagnostics::{unsafe_operations, unsafe_operations_for_body};
 use intern::{Interned, Symbol, sym};
 use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
 use smallvec::{SmallVec, smallvec};
 use span::{Edition, FileId, SyntaxContext};
-use stdx::TupleExt;
+use stdx::{TupleExt, always};
 use syntax::{
     AstNode, AstToken, Direction, SyntaxKind, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange,
     TextSize,
@@ -1765,6 +1765,25 @@ impl<'db> SemanticsImpl<'db> {
         res
     }
 
+    pub fn get_unsafe_ops_for_unsafe_block(&self, block: ast::BlockExpr) -> Vec<ExprOrPatSource> {
+        always!(block.unsafe_token().is_some());
+        let block = self.wrap_node_infile(ast::Expr::from(block));
+        let Some(def) = self.body_for(block.syntax()) else { return Vec::new() };
+        let def = def.into();
+        let (body, source_map) = self.db.body_with_source_map(def);
+        let infer = self.db.infer(def);
+        let Some(ExprOrPatId::ExprId(block)) = source_map.node_expr(block.as_ref()) else {
+            return Vec::new();
+        };
+        let mut res = Vec::default();
+        unsafe_operations(self.db, &infer, def, &body, block, &mut |node, _| {
+            if let Ok(node) = source_map.expr_or_pat_syntax(node) {
+                res.push(node);
+            }
+        });
+        res
+    }
+
     pub fn is_unsafe_macro_call(&self, macro_call: &ast::MacroCall) -> bool {
         let Some(mac) = self.resolve_macro_call(macro_call) else { return false };
         if mac.is_asm_like(self.db) {

src/tools/rust-analyzer/crates/hir/src/source_analyzer.rs

@@ -1283,7 +1283,7 @@ impl<'db> SourceAnalyzer<'db> {
         {
             let mut is_unsafe = false;
             let mut walk_expr = |expr_id| {
-                unsafe_operations(db, infer, def, body, expr_id, &mut |inside_unsafe_block| {
+                unsafe_operations(db, infer, def, body, expr_id, &mut |_, inside_unsafe_block| {
                     is_unsafe |= inside_unsafe_block == InsideUnsafeBlock::No
                 })
             };

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

@@ -805,10 +805,8 @@ pub(crate) fn highlight_unsafe_points(
         push_to_highlights(unsafe_token_file_id, Some(unsafe_token.text_range()));
 
         // highlight unsafe operations
-        if let Some(block) = block_expr
-            && let Some(body) = sema.body_for(InFile::new(unsafe_token_file_id, block.syntax()))
-        {
-            let unsafe_ops = sema.get_unsafe_ops(body);
+        if let Some(block) = block_expr {
+            let unsafe_ops = sema.get_unsafe_ops_for_unsafe_block(block);
             for unsafe_op in unsafe_ops {
                 push_to_highlights(unsafe_op.file_id, Some(unsafe_op.value.text_range()));
             }
@@ -2535,4 +2533,21 @@ fn foo() {
 "#,
         );
     }
+
+    #[test]
+    fn different_unsafe_block() {
+        check(
+            r#"
+fn main() {
+    unsafe$0 {
+ // ^^^^^^
+        *(0 as *const u8)
+     // ^^^^^^^^^^^^^^^^^
+    };
+    unsafe { *(1 as *const u8) };
+    unsafe { *(2 as *const u8) };
+}
+        "#,
+        );
+    }
 }