Use HIR unsafety information for unsafe syntax highlightng

Author: Veykril

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

@@ -9,5 +9,8 @@ pub use crate::diagnostics::{
     expr::{
         record_literal_missing_fields, record_pattern_missing_fields, BodyValidationDiagnostic,
     },
-    unsafe_check::{missing_unsafe, unsafe_expressions, InsideUnsafeBlock, UnsafetyReason},
+    unsafe_check::{
+        missing_unsafe, unsafe_operations, unsafe_operations_for_body, InsideUnsafeBlock,
+        UnsafetyReason,
+    },
 };

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

@@ -95,7 +95,26 @@ enum UnsafeDiagnostic {
     DeprecatedSafe2024 { node: ExprId, inside_unsafe_block: InsideUnsafeBlock },
 }
 
-pub fn unsafe_expressions(
+pub fn unsafe_operations_for_body(
+    db: &dyn HirDatabase,
+    infer: &InferenceResult,
+    def: DefWithBodyId,
+    body: &Body,
+    callback: &mut dyn FnMut(ExprOrPatId),
+) {
+    let mut visitor_callback = |diag| {
+        if let UnsafeDiagnostic::UnsafeOperation { node, .. } = diag {
+            callback(node);
+        }
+    };
+    let mut visitor = UnsafeVisitor::new(db, infer, body, def, &mut visitor_callback);
+    visitor.walk_expr(body.body_expr);
+    for &param in &body.params {
+        visitor.walk_pat(param);
+    }
+}
+
+pub fn unsafe_operations(
     db: &dyn HirDatabase,
     infer: &InferenceResult,
     def: DefWithBodyId,
@@ -281,13 +300,6 @@ impl<'a> UnsafeVisitor<'a> {
                     self.on_unsafe_op(current.into(), UnsafetyReason::RawPtrDeref);
                 }
             }
-            Expr::Unsafe { .. } => {
-                let old_inside_unsafe_block =
-                    mem::replace(&mut self.inside_unsafe_block, InsideUnsafeBlock::Yes);
-                self.body.walk_child_exprs_without_pats(current, |child| self.walk_expr(child));
-                self.inside_unsafe_block = old_inside_unsafe_block;
-                return;
-            }
             &Expr::Assignment { target, value: _ } => {
                 let old_inside_assignment = mem::replace(&mut self.inside_assignment, true);
                 self.walk_pats_top(std::iter::once(target), current);
@@ -306,6 +318,20 @@ impl<'a> UnsafeVisitor<'a> {
                     }
                 }
             }
+            Expr::Unsafe { statements, .. } => {
+                let old_inside_unsafe_block =
+                    mem::replace(&mut self.inside_unsafe_block, InsideUnsafeBlock::Yes);
+                self.walk_pats_top(
+                    statements.iter().filter_map(|statement| match statement {
+                        &Statement::Let { pat, .. } => Some(pat),
+                        _ => None,
+                    }),
+                    current,
+                );
+                self.body.walk_child_exprs_without_pats(current, |child| self.walk_expr(child));
+                self.inside_unsafe_block = old_inside_unsafe_block;
+                return;
+            }
             Expr::Block { statements, .. } | Expr::Async { statements, .. } => {
                 self.walk_pats_top(
                     statements.iter().filter_map(|statement| match statement {

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

@@ -12,6 +12,7 @@ use std::{
 
 use either::Either;
 use hir_def::{
+    expr_store::ExprOrPatSource,
     hir::{Expr, ExprOrPatId},
     lower::LowerCtx,
     nameres::{MacroSubNs, ModuleOrigin},
@@ -30,6 +31,7 @@ use hir_expand::{
     name::AsName,
     ExpandResult, FileRange, InMacroFile, MacroCallId, MacroFileId, MacroFileIdExt,
 };
+use hir_ty::diagnostics::unsafe_operations_for_body;
 use intern::{sym, Symbol};
 use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -48,8 +50,8 @@ use crate::{
     db::HirDatabase,
     semantics::source_to_def::{ChildContainer, SourceToDefCache, SourceToDefCtx},
     source_analyzer::{name_hygiene, resolve_hir_path, SourceAnalyzer},
-    Access, Adjust, Adjustment, Adt, AutoBorrow, BindingMode, BuiltinAttr, Callable, Const,
-    ConstParam, Crate, DeriveHelper, Enum, Field, Function, GenericSubstitution, HasSource,
+    Adjust, Adjustment, Adt, AutoBorrow, BindingMode, BuiltinAttr, Callable, Const, ConstParam,
+    Crate, DefWithBody, DeriveHelper, Enum, Field, Function, GenericSubstitution, HasSource,
     HirFileId, Impl, InFile, InlineAsmOperand, ItemInNs, Label, LifetimeParam, Local, Macro,
     Module, ModuleDef, Name, OverloadedDeref, Path, ScopeDef, Static, Struct, ToolModule, Trait,
     TraitAlias, TupleField, Type, TypeAlias, TypeParam, Union, Variant, VariantDef,
@@ -1555,6 +1557,19 @@ impl<'db> SemanticsImpl<'db> {
             .matched_arm
     }
 
+    pub fn get_unsafe_ops(&self, def: DefWithBody) -> FxHashSet<ExprOrPatSource> {
+        let def = DefWithBodyId::from(def);
+        let (body, source_map) = self.db.body_with_source_map(def);
+        let infer = self.db.infer(def);
+        let mut res = FxHashSet::default();
+        unsafe_operations_for_body(self.db, &infer, def, &body, &mut |node| {
+            if let Ok(node) = source_map.expr_or_pat_syntax(node) {
+                res.insert(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_or_global_asm(self.db) {
@@ -1682,6 +1697,15 @@ impl<'db> SemanticsImpl<'db> {
         Some(res)
     }
 
+    pub fn body_for(&self, node: InFile<&SyntaxNode>) -> Option<DefWithBody> {
+        let container = self.with_ctx(|ctx| ctx.find_container(node))?;
+
+        match container {
+            ChildContainer::DefWithBodyId(def) => Some(def.into()),
+            _ => None,
+        }
+    }
+
     /// Returns none if the file of the node is not part of a crate.
     fn analyze(&self, node: &SyntaxNode) -> Option<SourceAnalyzer> {
         let node = self.find_file(node);
@@ -1783,91 +1807,6 @@ impl<'db> SemanticsImpl<'db> {
         InFile::new(file_id, node)
     }
 
-    pub fn is_unsafe_method_call(&self, method_call_expr: &ast::MethodCallExpr) -> bool {
-        method_call_expr
-            .receiver()
-            .and_then(|expr| {
-                let field_expr = match expr {
-                    ast::Expr::FieldExpr(field_expr) => field_expr,
-                    _ => return None,
-                };
-                let ty = self.type_of_expr(&field_expr.expr()?)?.original;
-                if !ty.is_packed(self.db) {
-                    return None;
-                }
-
-                let func = self.resolve_method_call(method_call_expr)?;
-                let res = match func.self_param(self.db)?.access(self.db) {
-                    Access::Shared | Access::Exclusive => true,
-                    Access::Owned => false,
-                };
-                Some(res)
-            })
-            .unwrap_or(false)
-    }
-
-    pub fn is_unsafe_ref_expr(&self, ref_expr: &ast::RefExpr) -> bool {
-        ref_expr
-            .expr()
-            .and_then(|expr| {
-                let field_expr = match expr {
-                    ast::Expr::FieldExpr(field_expr) => field_expr,
-                    _ => return None,
-                };
-                let expr = field_expr.expr()?;
-                self.type_of_expr(&expr)
-            })
-            // Binding a reference to a packed type is possibly unsafe.
-            .map(|ty| ty.original.is_packed(self.db))
-            .unwrap_or(false)
-
-        // FIXME This needs layout computation to be correct. It will highlight
-        // more than it should with the current implementation.
-    }
-
-    pub fn is_unsafe_ident_pat(&self, ident_pat: &ast::IdentPat) -> bool {
-        if ident_pat.ref_token().is_none() {
-            return false;
-        }
-
-        ident_pat
-            .syntax()
-            .parent()
-            .and_then(|parent| {
-                // `IdentPat` can live under `RecordPat` directly under `RecordPatField` or
-                // `RecordPatFieldList`. `RecordPatField` also lives under `RecordPatFieldList`,
-                // so this tries to lookup the `IdentPat` anywhere along that structure to the
-                // `RecordPat` so we can get the containing type.
-                let record_pat = ast::RecordPatField::cast(parent.clone())
-                    .and_then(|record_pat| record_pat.syntax().parent())
-                    .or_else(|| Some(parent.clone()))
-                    .and_then(|parent| {
-                        ast::RecordPatFieldList::cast(parent)?
-                            .syntax()
-                            .parent()
-                            .and_then(ast::RecordPat::cast)
-                    });
-
-                // If this doesn't match a `RecordPat`, fallback to a `LetStmt` to see if
-                // this is initialized from a `FieldExpr`.
-                if let Some(record_pat) = record_pat {
-                    self.type_of_pat(&ast::Pat::RecordPat(record_pat))
-                } else if let Some(let_stmt) = ast::LetStmt::cast(parent) {
-                    let field_expr = match let_stmt.initializer()? {
-                        ast::Expr::FieldExpr(field_expr) => field_expr,
-                        _ => return None,
-                    };
-
-                    self.type_of_expr(&field_expr.expr()?)
-                } else {
-                    None
-                }
-            })
-            // Binding a reference to a packed type is possibly unsafe.
-            .map(|ty| ty.original.is_packed(self.db))
-            .unwrap_or(false)
-    }
-
     /// Returns `true` if the `node` is inside an `unsafe` context.
     pub fn is_inside_unsafe(&self, expr: &ast::Expr) -> bool {
         let Some(enclosing_item) =

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

@@ -35,7 +35,7 @@ use hir_expand::{
 };
 use hir_ty::{
     diagnostics::{
-        record_literal_missing_fields, record_pattern_missing_fields, unsafe_expressions,
+        record_literal_missing_fields, record_pattern_missing_fields, unsafe_operations,
         InsideUnsafeBlock,
     },
     from_assoc_type_id,
@@ -1160,7 +1160,7 @@ impl SourceAnalyzer {
             if let Some(expanded_expr) = sm.macro_expansion_expr(macro_expr) {
                 let mut is_unsafe = false;
                 let mut walk_expr = |expr_id| {
-                    unsafe_expressions(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-assists/src/handlers/extract_function.rs

@@ -750,7 +750,10 @@ impl FunctionBody {
                         ast::Stmt::Item(_) => (),
                         ast::Stmt::LetStmt(stmt) => {
                             if let Some(pat) = stmt.pat() {
-                                walk_pat(&pat, cb);
+                                walk_pat(&pat, &mut |pat| {
+                                    cb(pat);
+                                    std::ops::ControlFlow::<(), ()>::Continue(())
+                                });
                             }
                             if let Some(expr) = stmt.initializer() {
                                 walk_patterns_in_expr(&expr, cb);

src/tools/rust-analyzer/crates/ide-db/src/syntax_helpers/node_ext.rs

@@ -1,4 +1,6 @@
 //! Various helper functions to work with SyntaxNodes.
+use std::ops::ControlFlow;
+
 use itertools::Itertools;
 use parser::T;
 use span::Edition;
@@ -119,7 +121,10 @@ pub fn walk_patterns_in_expr(start: &ast::Expr, cb: &mut dyn FnMut(ast::Pat)) {
         match ast::Stmt::cast(node.clone()) {
             Some(ast::Stmt::LetStmt(l)) => {
                 if let Some(pat) = l.pat() {
-                    walk_pat(&pat, cb);
+                    walk_pat(&pat, &mut |pat| {
+                        cb(pat);
+                        ControlFlow::<(), ()>::Continue(())
+                    });
                 }
                 if let Some(expr) = l.initializer() {
                     walk_patterns_in_expr(&expr, cb);
@@ -154,15 +159,21 @@ pub fn walk_patterns_in_expr(start: &ast::Expr, cb: &mut dyn FnMut(ast::Pat)) {
                     }
                 } else if let Some(pat) = ast::Pat::cast(node) {
                     preorder.skip_subtree();
-                    walk_pat(&pat, cb);
+                    walk_pat(&pat, &mut |pat| {
+                        cb(pat);
+                        ControlFlow::<(), ()>::Continue(())
+                    });
                 }
             }
         }
     }
 }
 
 /// Preorder walk all the pattern's sub patterns.
-pub fn walk_pat(pat: &ast::Pat, cb: &mut dyn FnMut(ast::Pat)) {
+pub fn walk_pat<T>(
+    pat: &ast::Pat,
+    cb: &mut dyn FnMut(ast::Pat) -> ControlFlow<T>,
+) -> ControlFlow<T> {
     let mut preorder = pat.syntax().preorder();
     while let Some(event) = preorder.next() {
         let node = match event {
@@ -173,10 +184,10 @@ pub fn walk_pat(pat: &ast::Pat, cb: &mut dyn FnMut(ast::Pat)) {
         match ast::Pat::cast(node) {
             Some(pat @ ast::Pat::ConstBlockPat(_)) => {
                 preorder.skip_subtree();
-                cb(pat);
+                cb(pat)?;
             }
             Some(pat) => {
-                cb(pat);
+                cb(pat)?;
             }
             // skip const args
             None if ast::GenericArg::can_cast(kind) => {
@@ -185,6 +196,7 @@ pub fn walk_pat(pat: &ast::Pat, cb: &mut dyn FnMut(ast::Pat)) {
             None => (),
         }
     }
+    ControlFlow::Continue(())
 }
 
 /// Preorder walk all the type's sub types.