Refactor the check_pointers interface and respect ZSTs in nullcheck

1c3t3a · 1c3t3a · commit 9b56e20c1d6a · 2024-12-20T16:37:59.000Z
diff --git a/compiler/rustc_mir_transform/src/check_alignment.rs b/compiler/rustc_mir_transform/src/check_alignment.rs
@@ -4,7 +4,7 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::{Ty, TyCtxt};
 use rustc_session::Session;
 
-use crate::check_pointers::check_pointers;
+use crate::check_pointers::{PointerCheck, check_pointers};
 
 pub(super) struct CheckAlignment;
 
@@ -24,35 +24,33 @@ impl<'tcx> crate::MirPass<'tcx> for CheckAlignment {
     }
 }
 
+/// Inserts the actual alignment check's logic. Returns a
+/// [AssertKind::MisalignedPointerDereference] on failure.
 fn insert_alignment_check<'tcx>(
     tcx: TyCtxt<'tcx>,
-    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
-    block_data: &mut BasicBlockData<'tcx>,
     pointer: Place<'tcx>,
     pointee_ty: Ty<'tcx>,
+    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
+    stmts: &mut Vec<Statement<'tcx>>,
     source_info: SourceInfo,
-    new_block: BasicBlock,
-) {
-    // Cast the pointer to a *const ()
+) -> PointerCheck<'tcx> {
+    // Cast the pointer to a *const ().
     let const_raw_ptr = Ty::new_imm_ptr(tcx, tcx.types.unit);
     let rvalue = Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(pointer), const_raw_ptr);
     let thin_ptr = local_decls.push(LocalDecl::with_source_info(const_raw_ptr, source_info)).into();
-    block_data
-        .statements
+    stmts
         .push(Statement { source_info, kind: StatementKind::Assign(Box::new((thin_ptr, rvalue))) });
 
-    // Transmute the pointer to a usize (equivalent to `ptr.addr()`)
+    // Transmute the pointer to a usize (equivalent to `ptr.addr()`).
     let rvalue = Rvalue::Cast(CastKind::Transmute, Operand::Copy(thin_ptr), tcx.types.usize);
     let addr = local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into();
-    block_data
-        .statements
-        .push(Statement { source_info, kind: StatementKind::Assign(Box::new((addr, rvalue))) });
+    stmts.push(Statement { source_info, kind: StatementKind::Assign(Box::new((addr, rvalue))) });
 
     // Get the alignment of the pointee
     let alignment =
         local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into();
     let rvalue = Rvalue::NullaryOp(NullOp::AlignOf, pointee_ty);
-    block_data.statements.push(Statement {
+    stmts.push(Statement {
         source_info,
         kind: StatementKind::Assign(Box::new((alignment, rvalue))),
     });
@@ -65,7 +63,7 @@ fn insert_alignment_check<'tcx>(
         user_ty: None,
         const_: Const::Val(ConstValue::Scalar(Scalar::from_target_usize(1, &tcx)), tcx.types.usize),
     }));
-    block_data.statements.push(Statement {
+    stmts.push(Statement {
         source_info,
         kind: StatementKind::Assign(Box::new((
             alignment_mask,
@@ -76,7 +74,7 @@ fn insert_alignment_check<'tcx>(
     // BitAnd the alignment mask with the pointer
     let alignment_bits =
         local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into();
-    block_data.statements.push(Statement {
+    stmts.push(Statement {
         source_info,
         kind: StatementKind::Assign(Box::new((
             alignment_bits,
@@ -94,29 +92,21 @@ fn insert_alignment_check<'tcx>(
         user_ty: None,
         const_: Const::Val(ConstValue::Scalar(Scalar::from_target_usize(0, &tcx)), tcx.types.usize),
     }));
-    block_data.statements.push(Statement {
+    stmts.push(Statement {
         source_info,
         kind: StatementKind::Assign(Box::new((
             is_ok,
             Rvalue::BinaryOp(BinOp::Eq, Box::new((Operand::Copy(alignment_bits), zero.clone()))),
         ))),
     });
 
-    // Set this block's terminator to our assert, continuing to new_block if we pass
-    block_data.terminator = Some(Terminator {
-        source_info,
-        kind: TerminatorKind::Assert {
-            cond: Operand::Copy(is_ok),
-            expected: true,
-            target: new_block,
-            msg: Box::new(AssertKind::MisalignedPointerDereference {
-                required: Operand::Copy(alignment),
-                found: Operand::Copy(addr),
-            }),
-            // This calls panic_misaligned_pointer_dereference, which is #[rustc_nounwind].
-            // We never want to insert an unwind into unsafe code, because unwinding could
-            // make a failing UB check turn into much worse UB when we start unwinding.
-            unwind: UnwindAction::Unreachable,
-        },
-    });
+    // Emit a check that asserts on the alignment and otherwise triggers a
+    // AssertKind::MisalignedPointerDereference.
+    PointerCheck {
+        cond: Operand::Copy(is_ok),
+        assert_kind: Box::new(AssertKind::MisalignedPointerDereference {
+            required: Operand::Copy(alignment),
+            found: Operand::Copy(addr),
+        }),
+    }
 }
diff --git a/compiler/rustc_mir_transform/src/check_null.rs b/compiler/rustc_mir_transform/src/check_null.rs
@@ -4,7 +4,7 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::{Ty, TyCtxt};
 use rustc_session::Session;
 
-use crate::check_pointers::check_pointers;
+use crate::check_pointers::{PointerCheck, check_pointers};
 
 pub(super) struct CheckNull;
 
@@ -20,54 +20,77 @@ impl<'tcx> crate::MirPass<'tcx> for CheckNull {
 
 fn insert_null_check<'tcx>(
     tcx: TyCtxt<'tcx>,
-    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
-    block_data: &mut BasicBlockData<'tcx>,
     pointer: Place<'tcx>,
-    _: Ty<'tcx>,
+    pointee_ty: Ty<'tcx>,
+    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
+    stmts: &mut Vec<Statement<'tcx>>,
     source_info: SourceInfo,
-    new_block: BasicBlock,
-) {
+) -> PointerCheck<'tcx> {
+    // Cast the pointer to a *const ().
     let const_raw_ptr = Ty::new_imm_ptr(tcx, tcx.types.unit);
     let rvalue = Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(pointer), const_raw_ptr);
     let thin_ptr = local_decls.push(LocalDecl::with_source_info(const_raw_ptr, source_info)).into();
-    block_data
-        .statements
+    stmts
         .push(Statement { source_info, kind: StatementKind::Assign(Box::new((thin_ptr, rvalue))) });
 
-    // Transmute the pointer to a usize (equivalent to `ptr.addr()`)
+    // Transmute the pointer to a usize (equivalent to `ptr.addr()`).
     let rvalue = Rvalue::Cast(CastKind::Transmute, Operand::Copy(thin_ptr), tcx.types.usize);
     let addr = local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into();
-    block_data
-        .statements
-        .push(Statement { source_info, kind: StatementKind::Assign(Box::new((addr, rvalue))) });
+    stmts.push(Statement { source_info, kind: StatementKind::Assign(Box::new((addr, rvalue))) });
 
-    // Check if the pointer is null.
-    let is_ok = local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into();
+    // Get size of the pointee (zero-sized reads and writes are allowed).
+    let rvalue = Rvalue::NullaryOp(NullOp::SizeOf, pointee_ty);
+    let sizeof_pointee =
+        local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into();
+    stmts.push(Statement {
+        source_info,
+        kind: StatementKind::Assign(Box::new((sizeof_pointee, rvalue))),
+    });
+
+    // Check that the pointee is a ZST.
     let zero = Operand::Constant(Box::new(ConstOperand {
         span: source_info.span,
         user_ty: None,
         const_: Const::Val(ConstValue::Scalar(Scalar::from_target_usize(0, &tcx)), tcx.types.usize),
     }));
-    block_data.statements.push(Statement {
+    let is_pointee_zst =
+        local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into();
+    stmts.push(Statement {
+        source_info,
+        kind: StatementKind::Assign(Box::new((
+            is_pointee_zst,
+            Rvalue::BinaryOp(BinOp::Eq, Box::new((Operand::Copy(sizeof_pointee), zero.clone()))),
+        ))),
+    });
+
+    // Check if the pointer is null.
+    let is_null = local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into();
+    stmts.push(Statement {
         source_info,
         kind: StatementKind::Assign(Box::new((
-            is_ok,
+            is_null,
             Rvalue::BinaryOp(BinOp::Ne, Box::new((Operand::Copy(addr), zero))),
         ))),
     });
 
-    // Set this block's terminator to our assert, continuing to new_block if we pass
-    block_data.terminator = Some(Terminator {
+    // Check if the pointer is null or points to a ZST.
+    let should_throw_exception =
+        local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into();
+    stmts.push(Statement {
         source_info,
-        kind: TerminatorKind::Assert {
-            cond: Operand::Copy(is_ok),
-            expected: true,
-            target: new_block,
-            msg: Box::new(AssertKind::NullPointerDereference),
-            // This calls panic_misaligned_pointer_dereference, which is #[rustc_nounwind].
-            // We never want to insert an unwind into unsafe code, because unwinding could
-            // make a failing UB check turn into much worse UB when we start unwinding.
-            unwind: UnwindAction::Unreachable,
-        },
+        kind: StatementKind::Assign(Box::new((
+            is_null,
+            Rvalue::BinaryOp(
+                BinOp::BitOr,
+                Box::new((Operand::Copy(is_null), Operand::Copy(is_pointee_zst))),
+            ),
+        ))),
     });
+
+    // Emit a PointerCheck that asserts on the condition and otherwise triggers
+    // a AssertKind::NullPointerDereference.
+    PointerCheck {
+        cond: Operand::Copy(should_throw_exception),
+        assert_kind: Box::new(AssertKind::NullPointerDereference),
+    }
 }
diff --git a/compiler/rustc_mir_transform/src/check_pointers.rs b/compiler/rustc_mir_transform/src/check_pointers.rs
@@ -5,21 +5,40 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use tracing::{debug, trace};
 
+/// Details of a pointer check, the condition on which we decide whether to
+/// fail the assert and an [AssertKind] that defines the behavior on failure.
+pub(crate) struct PointerCheck<'tcx> {
+    pub(crate) cond: Operand<'tcx>,
+    pub(crate) assert_kind: Box<AssertKind<Operand<'tcx>>>,
+}
+
+/// Utility for adding a check for read/write on every sized, unsafe pointer.
+///
+/// Visits every read/write access to a [Sized], unsafe pointer and inserts a
+/// new basic block directly before the pointer access. Then calls `on_finding`
+/// to insert the actual logic for a pointer check (e.g. check for alignment).
+/// This utility takes care of the right order of blocks, the only thing a
+/// caller must do in `on_finding` is:
+/// - Append [Statement]s to `stmts`.
+/// - Append [LocalDecl]s to `local_decls`.
+/// - Return a [PointerCheck] that contains the condition and an [AssertKind].
+///   The AssertKind must be a panic with `#[rustc_nounwind]`.
+/// This utility will insert a terminator block that asserts on the condition
+/// and panics on failure.
 pub(crate) fn check_pointers<'a, 'tcx, F>(
     tcx: TyCtxt<'tcx>,
     body: &mut Body<'tcx>,
     excluded_pointees: &'a [Ty<'tcx>],
     on_finding: F,
 ) where
     F: Fn(
-        TyCtxt<'tcx>,
-        &mut IndexVec<Local, LocalDecl<'tcx>>,
-        &mut BasicBlockData<'tcx>,
-        Place<'tcx>,
-        Ty<'tcx>,
-        SourceInfo,
-        BasicBlock,
-    ),
+        /* tcx: */ TyCtxt<'tcx>,
+        /* pointer: */ Place<'tcx>,
+        /* pointee_ty: */ Ty<'tcx>,
+        /* local_decls: */ &mut IndexVec<Local, LocalDecl<'tcx>>,
+        /* stmts: */ &mut Vec<Statement<'tcx>>,
+        /* source_info: */ SourceInfo,
+    ) -> PointerCheck<'tcx>,
 {
     // This pass emits new panics. If for whatever reason we do not have a panic
     // implementation, running this pass may cause otherwise-valid code to not compile.
@@ -48,15 +67,33 @@ pub(crate) fn check_pointers<'a, 'tcx, F>(
             for (local, ty) in finder.into_found_pointers() {
                 debug!("Inserting check for {:?}", ty);
                 let new_block = split_block(basic_blocks, location);
-                on_finding(
+
+                // Invoke `on_finding` which appends to `local_decls` and the
+                // blocks statements. It returns information about the assert
+                // we're performing in the Terminator.
+                let block_data = &mut basic_blocks[block];
+                let pointer_check = on_finding(
                     tcx,
-                    local_decls,
-                    &mut basic_blocks[block],
                     local,
                     ty,
+                    local_decls,
+                    &mut block_data.statements,
                     source_info,
-                    new_block,
                 );
+                block_data.terminator = Some(Terminator {
+                    source_info,
+                    kind: TerminatorKind::Assert {
+                        cond: pointer_check.cond,
+                        expected: true,
+                        target: new_block,
+                        msg: pointer_check.assert_kind,
+                        // This calls a panic function associated with the pointer check, which
+                        // is #[rustc_nounwind]. We never want to insert an unwind into unsafe
+                        // code, because unwinding could make a failing UB check turn into much
+                        // worse UB when we start unwinding.
+                        unwind: UnwindAction::Unreachable,
+                    },
+                });
             }
         }
     }
@@ -128,7 +165,7 @@ impl<'a, 'tcx> Visitor<'tcx> for PointerFinder<'a, 'tcx> {
             return;
         }
 
-        // We don't need to look for str and slices, we already rejected unsized types above
+        // We don't need to look for slices, we already rejected unsized types above.
         let element_ty = match pointee_ty.kind() {
             ty::Array(ty, _) => *ty,
             _ => pointee_ty,
@@ -150,7 +187,7 @@ fn split_block(
 ) -> BasicBlock {
     let block_data = &mut basic_blocks[location.block];
 
-    // Drain every statement after this one and move the current terminator to a new basic block
+    // Drain every statement after this one and move the current terminator to a new basic block.
     let new_block = BasicBlockData {
         statements: block_data.statements.split_off(location.statement_index),
         terminator: block_data.terminator.take(),
diff --git a/tests/ui/mir/null/zero_sized_access.rs b/tests/ui/mir/null/zero_sized_access.rs
@@ -0,0 +1,7 @@
+//@ run-pass
+//@ compile-flags: -C debug-assertions
+
+fn main() {
+    let ptr: *const () = std::ptr::null();
+    let _ptr = unsafe { *ptr };
+}
