Add support for deferred casting for the invalid_reference_casting lint

6 files changed, 80 insertions, 32 deletions

diff --git a/compiler/rustc_lint/messages.ftl b/compiler/rustc_lint/messages.ftl
index 16e17fc9d6a..95f0b651962 100644
--- a/compiler/rustc_lint/messages.ftl
+++ b/compiler/rustc_lint/messages.ftl
@@ -319,6 +319,7 @@ lint_invalid_nan_comparisons_eq_ne = incorrect NaN comparison, NaN cannot be dir
 lint_invalid_nan_comparisons_lt_le_gt_ge = incorrect NaN comparison, NaN is not orderable
 
 lint_invalid_reference_casting = casting `&T` to `&mut T` is undefined behavior, even if the reference is unused, consider instead using an `UnsafeCell`
+    .label = casting happend here
 
 lint_lintpass_by_hand = implementing `LintPass` by hand
     .help = try using `declare_lint_pass!` or `impl_lint_pass!` instead
diff --git a/compiler/rustc_lint/src/lib.rs b/compiler/rustc_lint/src/lib.rs
index 53089294fe2..96fd3ccf774 100644
--- a/compiler/rustc_lint/src/lib.rs
+++ b/compiler/rustc_lint/src/lib.rs
@@ -218,7 +218,7 @@ late_lint_methods!(
             BoxPointers: BoxPointers,
             PathStatements: PathStatements,
             LetUnderscore: LetUnderscore,
-            InvalidReferenceCasting: InvalidReferenceCasting,
+            InvalidReferenceCasting: InvalidReferenceCasting::default(),
             // Depends on referenced function signatures in expressions
             UnusedResults: UnusedResults,
             NonUpperCaseGlobals: NonUpperCaseGlobals,
diff --git a/compiler/rustc_lint/src/lints.rs b/compiler/rustc_lint/src/lints.rs
index 968172693a9..6d774070ec4 100644
--- a/compiler/rustc_lint/src/lints.rs
+++ b/compiler/rustc_lint/src/lints.rs
@@ -746,7 +746,10 @@ pub enum InvalidFromUtf8Diag {
 // reference_casting.rs
 #[derive(LintDiagnostic)]
 #[diag(lint_invalid_reference_casting)]
-pub struct InvalidReferenceCastingDiag;
+pub struct InvalidReferenceCastingDiag {
+    #[label]
+    pub orig_cast: Option<Span>,
+}
 
 // hidden_unicode_codepoints.rs
 #[derive(LintDiagnostic)]
diff --git a/compiler/rustc_lint/src/reference_casting.rs b/compiler/rustc_lint/src/reference_casting.rs
index 1a172dd6942..e6cab4cebe7 100644
--- a/compiler/rustc_lint/src/reference_casting.rs
+++ b/compiler/rustc_lint/src/reference_casting.rs
@@ -1,7 +1,8 @@
 use rustc_ast::Mutability;
-use rustc_hir::{Expr, ExprKind, MutTy, TyKind, UnOp};
-use rustc_middle::ty;
-use rustc_span::sym;
+use rustc_data_structures::fx::FxHashMap;
+use rustc_hir::{def::Res, Expr, ExprKind, HirId, Local, QPath, StmtKind, UnOp};
+use rustc_middle::ty::{self, TypeAndMut};
+use rustc_span::{sym, Span};
 
 use crate::{lints::InvalidReferenceCastingDiag, LateContext, LateLintPass, LintContext};
 
@@ -33,42 +34,74 @@ declare_lint! {
     "casts of `&T` to `&mut T` without interior mutability"
 }
 
-declare_lint_pass!(InvalidReferenceCasting => [INVALID_REFERENCE_CASTING]);
+#[derive(Default)]
+pub struct InvalidReferenceCasting {
+    casted: FxHashMap<HirId, Span>,
+}
+
+impl_lint_pass!(InvalidReferenceCasting => [INVALID_REFERENCE_CASTING]);
 
 impl<'tcx> LateLintPass<'tcx> for InvalidReferenceCasting {
-    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'tcx>) {
-        let ExprKind::Unary(UnOp::Deref, e) = &expr.kind else {
+    fn check_stmt(&mut self, cx: &LateContext<'tcx>, stmt: &'tcx rustc_hir::Stmt<'tcx>) {
+        let StmtKind::Local(local) = stmt.kind else {
             return;
         };
-
-        let e = e.peel_blocks();
-        let e = if let ExprKind::Cast(e, t) = e.kind
-            && let TyKind::Ptr(MutTy { mutbl: Mutability::Mut, .. }) = t.kind {
-            e
-        } else if let ExprKind::MethodCall(_, expr, [], _) = e.kind
-            && let Some(def_id) = cx.typeck_results().type_dependent_def_id(e.hir_id)
-            && cx.tcx.is_diagnostic_item(sym::ptr_cast_mut, def_id) {
-            expr
-        } else {
+        let Local { init: Some(init), els: None, .. } = local else {
             return;
         };
 
-        let e = e.peel_blocks();
-        let e = if let ExprKind::Cast(e, t) = e.kind
-            && let TyKind::Ptr(MutTy { mutbl: Mutability::Not, .. }) = t.kind {
-            e
-        } else if let ExprKind::Call(path, [arg]) = e.kind
-            && let ExprKind::Path(ref qpath) = path.kind
-            && let Some(def_id) = cx.qpath_res(qpath, path.hir_id).opt_def_id()
-            && cx.tcx.is_diagnostic_item(sym::ptr_from_ref, def_id) {
-            arg
-        } else {
+        if is_cast_from_const_to_mut(cx, init) {
+            self.casted.insert(local.pat.hir_id, init.span);
+        }
+    }
+
+    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'tcx>) {
+        let ExprKind::Unary(UnOp::Deref, e) = &expr.kind else {
             return;
         };
 
-        let e = e.peel_blocks();
-        if let ty::Ref(..) = cx.typeck_results().node_type(e.hir_id).kind() {
-            cx.emit_spanned_lint(INVALID_REFERENCE_CASTING, expr.span, InvalidReferenceCastingDiag);
+        if is_cast_from_const_to_mut(cx, e) {
+            cx.emit_spanned_lint(INVALID_REFERENCE_CASTING, expr.span, InvalidReferenceCastingDiag { orig_cast: None });
+        } else if let ExprKind::Path(QPath::Resolved(_, path)) = e.kind
+            && let Res::Local(hir_id) = &path.res
+            && let Some(orig_cast) = self.casted.get(hir_id) {
+            cx.emit_spanned_lint(INVALID_REFERENCE_CASTING, expr.span, InvalidReferenceCastingDiag { orig_cast: Some(*orig_cast) });
         }
     }
 }
+
+fn is_cast_from_const_to_mut<'tcx>(cx: &LateContext<'tcx>, e: &'tcx Expr<'tcx>) -> bool {
+    let e = e.peel_blocks();
+
+    // <expr> as *mut ...
+    let e = if let ExprKind::Cast(e, t) = e.kind
+        && let ty::RawPtr(TypeAndMut { mutbl: Mutability::Mut, .. }) = cx.typeck_results().node_type(t.hir_id).kind() {
+        e
+    // <expr>.cast_mut()
+    } else if let ExprKind::MethodCall(_, expr, [], _) = e.kind
+        && let Some(def_id) = cx.typeck_results().type_dependent_def_id(e.hir_id)
+        && cx.tcx.is_diagnostic_item(sym::ptr_cast_mut, def_id) {
+        expr
+    } else {
+        return false;
+    };
+
+    let e = e.peel_blocks();
+
+    // <expr> as *const ...
+    let e = if let ExprKind::Cast(e, t) = e.kind
+        && let ty::RawPtr(TypeAndMut { mutbl: Mutability::Not, .. }) = cx.typeck_results().node_type(t.hir_id).kind() {
+        e
+    // ptr::from_ref(<expr>)
+    } else if let ExprKind::Call(path, [arg]) = e.kind
+        && let ExprKind::Path(ref qpath) = path.kind
+        && let Some(def_id) = cx.qpath_res(qpath, path.hir_id).opt_def_id()
+        && cx.tcx.is_diagnostic_item(sym::ptr_from_ref, def_id) {
+        arg
+    } else {
+        return false;
+    };
+
+    let e = e.peel_blocks();
+    matches!(cx.typeck_results().node_type(e.hir_id).kind(), ty::Ref(..))
+}
diff --git a/tests/ui/lint/reference_casting.rs b/tests/ui/lint/reference_casting.rs
index 745d7070143..c3168cd4c01 100644
--- a/tests/ui/lint/reference_casting.rs
+++ b/tests/ui/lint/reference_casting.rs
@@ -36,6 +36,9 @@ fn main() {
         //~^ ERROR casting `&T` to `&mut T` is undefined behavior
         *(std::ptr::from_ref({ num }) as *mut i32) += 1;
         //~^ ERROR casting `&T` to `&mut T` is undefined behavior
+        let value = num as *const i32 as *mut i32;
+        *value = 1;
+        //~^ ERROR casting `&T` to `&mut T` is undefined behavior
 
         // Shouldn't be warned against
         println!("{}", *(num as *const _ as *const i16));
diff --git a/tests/ui/lint/reference_casting.stderr b/tests/ui/lint/reference_casting.stderr
index d1dd1b32ff4..d9ce4b38387 100644
--- a/tests/ui/lint/reference_casting.stderr
+++ b/tests/ui/lint/reference_casting.stderr
@@ -60,5 +60,13 @@ error: casting `&T` to `&mut T` is undefined behavior, even if the reference is
 LL |         *(std::ptr::from_ref({ num }) as *mut i32) += 1;
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-error: aborting due to 10 previous errors
+error: casting `&T` to `&mut T` is undefined behavior, even if the reference is unused, consider instead using an `UnsafeCell`
+  --> $DIR/reference_casting.rs:40:9
+   |
+LL |         let value = num as *const i32 as *mut i32;
+   |                     ----------------------------- casting happend here
+LL |         *value = 1;
+   |         ^^^^^^
+
+error: aborting due to 11 previous errors