Auto merge of #117473 - saethlin:codegen-alignment-checks, r=<try>

Move alignment checks to codegen

Implementing UB checks entirely in a MIR transform is quite limiting, we don't know for sure what all our types are so we need to make a lot of sacrifices. For example in here we used to emit MIR to compute the alignment mask at runtime, because the pointee type could be generic. Implementing the checks in codegen frees us from that requirement, because we get to deal with monomorphized types.

But I don't think we can move these checks entirely into codegen, because inserting the check needs to insert a new terminator into a basic block, which splits the previous basic block into two. We can't add control flow like this in codegen, but we can in MIR.

So now the MIR transform just inserts a `TerminatorKind::UbCheck` which is effectively a `Goto` that also reads an `Operand` (because it either goes to the target block or terminates), and codegen expands that new terminator into the actual check.

---

Also I'm writing this with the expectation that I implement the niche checks in the same manner, because they have the same problem with polymorphic MIR, possibly worse.

r? `@ghost`

compiler/rustc_borrowck/src/invalidation.rs

@@ -205,7 +205,8 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
             | TerminatorKind::UnwindTerminate(_)
             | TerminatorKind::Unreachable
             | TerminatorKind::FalseEdge { real_target: _, imaginary_target: _ }
-            | TerminatorKind::FalseUnwind { real_target: _, unwind: _ } => {
+            | TerminatorKind::FalseUnwind { real_target: _, unwind: _ }
+            | TerminatorKind::UbCheck { .. } => {
                 // no data used, thus irrelevant to borrowck
             }
         }

compiler/rustc_borrowck/src/lib.rs

@@ -785,7 +785,8 @@ impl<'cx, 'tcx, R> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx, R> for MirBorro
             | TerminatorKind::Return
             | TerminatorKind::CoroutineDrop
             | TerminatorKind::FalseEdge { real_target: _, imaginary_target: _ }
-            | TerminatorKind::FalseUnwind { real_target: _, unwind: _ } => {
+            | TerminatorKind::FalseUnwind { real_target: _, unwind: _ }
+            | TerminatorKind::UbCheck { .. } => {
                 // no data used, thus irrelevant to borrowck
             }
         }
@@ -835,7 +836,8 @@ impl<'cx, 'tcx, R> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx, R> for MirBorro
             | TerminatorKind::Goto { .. }
             | TerminatorKind::SwitchInt { .. }
             | TerminatorKind::Unreachable
-            | TerminatorKind::InlineAsm { .. } => {}
+            | TerminatorKind::InlineAsm { .. }
+            | TerminatorKind::UbCheck { .. } => {}
         }
     }
 }

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1358,7 +1358,8 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             | TerminatorKind::Drop { .. }
             | TerminatorKind::FalseEdge { .. }
             | TerminatorKind::FalseUnwind { .. }
-            | TerminatorKind::InlineAsm { .. } => {
+            | TerminatorKind::InlineAsm { .. }
+            | TerminatorKind::UbCheck { .. } => {
                 // no checks needed for these
             }
 
@@ -1690,6 +1691,9 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 }
                 self.assert_iscleanup_unwind(body, block_data, unwind, is_cleanup);
             }
+            TerminatorKind::UbCheck { target, .. } => {
+                self.assert_iscleanup(body, block_data, target, is_cleanup)
+            }
         }
     }
 

compiler/rustc_codegen_cranelift/src/base.rs

@@ -356,18 +356,6 @@ fn codegen_fn_body(fx: &mut FunctionCx<'_, '_, '_>, start_block: Block) {
                             source_info.span,
                         );
                     }
-                    AssertKind::MisalignedPointerDereference { ref required, ref found } => {
-                        let required = codegen_operand(fx, required).load_scalar(fx);
-                        let found = codegen_operand(fx, found).load_scalar(fx);
-                        let location = fx.get_caller_location(source_info).load_scalar(fx);
-
-                        codegen_panic_inner(
-                            fx,
-                            rustc_hir::LangItem::PanicMisalignedPointerDereference,
-                            &[required, found, location],
-                            source_info.span,
-                        );
-                    }
                     _ => {
                         let msg_str = msg.description();
                         codegen_panic(fx, msg_str, source_info);
@@ -488,6 +476,40 @@ fn codegen_fn_body(fx: &mut FunctionCx<'_, '_, '_>, start_block: Block) {
                 let target_block = fx.get_block(*target);
                 fx.bcx.ins().jump(target_block, &[]);
             }
+            TerminatorKind::UbCheck { target, kind: UbCheckKind::PointerAlignment { pointer } } => {
+                let location = fx.get_caller_location(source_info).load_scalar(fx);
+                let pointer = codegen_operand(fx, pointer);
+                let pointee_ty = pointer.layout().ty.builtin_deref(true).unwrap();
+                let pointee_layout = fx.layout_of(pointee_ty.ty);
+                let pointer = pointer.load_scalar(fx);
+
+                // Compute the alignment mask
+                let align = pointee_layout.align.abi.bytes() as i64;
+                let mask = fx.bcx.ins().iconst(fx.pointer_type, align - 1);
+                let required = fx.bcx.ins().iconst(fx.pointer_type, align);
+
+                // And the pointer with the mask
+                let masked = fx.bcx.ins().band(pointer, mask);
+
+                // Branch on whether the masked value is zero
+                let is_zero = fx.bcx.ins().icmp_imm(IntCC::Equal, masked, 0);
+
+                // Create a new block that we will jump to if !is_zero
+                let failure = fx.bcx.create_block();
+                fx.bcx.set_cold_block(failure);
+
+                let target = fx.get_block(*target);
+                fx.bcx.ins().brif(is_zero, target, &[], failure, &[]);
+
+                // Codegen the actual panic
+                fx.bcx.switch_to_block(failure);
+                codegen_panic_inner(
+                    fx,
+                    rustc_hir::LangItem::PanicMisalignedPointerDereference,
+                    &[required, pointer, location],
+                    source_info.span,
+                );
+            }
         };
     }
 }

compiler/rustc_codegen_cranelift/src/constant.rs

@@ -508,7 +508,8 @@ pub(crate) fn mir_operand_get_const_val<'tcx>(
                     | TerminatorKind::Return
                     | TerminatorKind::Unreachable
                     | TerminatorKind::Drop { .. }
-                    | TerminatorKind::Assert { .. } => {}
+                    | TerminatorKind::Assert { .. }
+                    | TerminatorKind::UbCheck { .. } => {}
                     TerminatorKind::Yield { .. }
                     | TerminatorKind::CoroutineDrop
                     | TerminatorKind::FalseEdge { .. }

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -277,7 +277,8 @@ pub fn cleanup_kinds(mir: &mir::Body<'_>) -> IndexVec<mir::BasicBlock, CleanupKi
                 | TerminatorKind::SwitchInt { .. }
                 | TerminatorKind::Yield { .. }
                 | TerminatorKind::FalseEdge { .. }
-                | TerminatorKind::FalseUnwind { .. } => { /* nothing to do */ }
+                | TerminatorKind::FalseUnwind { .. }
+                | TerminatorKind::UbCheck { .. } => { /* nothing to do */ }
                 TerminatorKind::Call { unwind, .. }
                 | TerminatorKind::InlineAsm { unwind, .. }
                 | TerminatorKind::Assert { unwind, .. }

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -12,7 +12,7 @@ use crate::MemFlags;
 use rustc_ast as ast;
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_hir::lang_items::LangItem;
-use rustc_middle::mir::{self, AssertKind, SwitchTargets, UnwindTerminateReason};
+use rustc_middle::mir::{self, AssertKind, SwitchTargets, UbCheckKind, UnwindTerminateReason};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
 use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths};
 use rustc_middle::ty::{self, Instance, Ty};
@@ -616,13 +616,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 // and `#[track_caller]` adds an implicit third argument.
                 (LangItem::PanicBoundsCheck, vec![index, len, location])
             }
-            AssertKind::MisalignedPointerDereference { ref required, ref found } => {
-                let required = self.codegen_operand(bx, required).immediate();
-                let found = self.codegen_operand(bx, found).immediate();
-                // It's `fn panic_misaligned_pointer_dereference(required: usize, found: usize)`,
-                // and `#[track_caller]` adds an implicit third argument.
-                (LangItem::PanicMisalignedPointerDereference, vec![required, found, location])
-            }
             _ => {
                 let msg = bx.const_str(msg.description());
                 // It's `pub fn panic(expr: &str)`, with the wide reference being passed
@@ -736,6 +729,79 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
     }
 
+    fn codegen_alignment_check(
+        &mut self,
+        helper: &TerminatorCodegenHelper<'tcx>,
+        bx: &mut Bx,
+        pointer: &mir::Operand<'tcx>,
+        source_info: mir::SourceInfo,
+        target: mir::BasicBlock,
+    ) -> MergingSucc {
+        let span = source_info.span;
+        let pointer = self.codegen_operand(bx, pointer);
+        let pointee_ty = pointer.layout.ty.builtin_deref(true).unwrap();
+        let pointee_layout = bx.layout_of(pointee_ty.ty);
+
+        let mk_usize = |v: u64| {
+            let layout = bx.layout_of(bx.tcx().types.usize);
+            let rustc_target::abi::Abi::Scalar(abi) = layout.abi else { unreachable!() };
+            let v = rustc_middle::mir::interpret::Scalar::from_target_usize(v, &bx.tcx());
+            bx.scalar_to_backend(v, abi, bx.cx().type_isize())
+        };
+
+        let align = pointee_layout.align.abi.bytes();
+        let mask = mk_usize(align - 1);
+        let zero = mk_usize(0);
+        let required = mk_usize(align);
+
+        let ptr_imm = match pointer.val {
+            crate::mir::OperandValue::Immediate(imm) => imm,
+            crate::mir::OperandValue::Pair(ptr, _) => ptr,
+            _ => {
+                unreachable!("{pointer:?}");
+            }
+        };
+        let int_imm = bx.ptrtoint(ptr_imm, bx.cx().type_isize());
+
+        let masked = bx.and(int_imm, mask);
+
+        let is_zero = bx.icmp(
+            crate::base::bin_op_to_icmp_predicate(mir::BinOp::Eq.to_hir_binop(), false),
+            masked,
+            zero,
+        );
+
+        let lltarget = helper.llbb_with_cleanup(self, target);
+        let panic_block = bx.append_sibling_block("panic");
+
+        bx.cond_br(is_zero, lltarget, panic_block);
+
+        bx.switch_to_block(panic_block);
+        self.set_debug_loc(bx, source_info);
+
+        let location = self.get_caller_location(bx, source_info).immediate();
+
+        let found = int_imm;
+
+        let (lang_item, args) =
+            (LangItem::PanicMisalignedPointerDereference, vec![required, found, location]);
+
+        let (fn_abi, llfn) = common::build_langcall(bx, Some(span), lang_item);
+        let merging_succ = helper.do_call(
+            self,
+            bx,
+            fn_abi,
+            llfn,
+            &args,
+            None,
+            mir::UnwindAction::Unreachable,
+            &[],
+            false,
+        );
+        assert_eq!(merging_succ, MergingSucc::False);
+        MergingSucc::False
+    }
+
     fn codegen_call_terminator(
         &mut self,
         helper: TerminatorCodegenHelper<'tcx>,
@@ -1291,6 +1357,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 self.instance,
                 mergeable_succ(),
             ),
+
+            mir::TerminatorKind::UbCheck {
+                target,
+                kind: UbCheckKind::PointerAlignment { ref pointer },
+            } => self.codegen_alignment_check(&helper, bx, pointer, terminator.source_info, target),
         }
     }
 

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -555,12 +555,6 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
             RemainderByZero(op) => RemainderByZero(eval_to_int(op)?),
             ResumedAfterReturn(coroutine_kind) => ResumedAfterReturn(*coroutine_kind),
             ResumedAfterPanic(coroutine_kind) => ResumedAfterPanic(*coroutine_kind),
-            MisalignedPointerDereference { ref required, ref found } => {
-                MisalignedPointerDereference {
-                    required: eval_to_int(required)?,
-                    found: eval_to_int(found)?,
-                }
-            }
         };
         Err(ConstEvalErrKind::AssertFailure(err).into())
     }

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -85,7 +85,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.pop_stack_frame(/* unwinding */ false)?
             }
 
-            Goto { target } => self.go_to_block(target),
+            Goto { target } | UbCheck { target, .. } => self.go_to_block(target),
 
             SwitchInt { ref discr, ref targets } => {
                 let discr = self.read_immediate(&self.eval_operand(discr, None)?)?;

compiler/rustc_const_eval/src/transform/check_consts/check.rs

@@ -1058,7 +1058,8 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
             | TerminatorKind::UnwindResume
             | TerminatorKind::Return
             | TerminatorKind::SwitchInt { .. }
-            | TerminatorKind::Unreachable => {}
+            | TerminatorKind::Unreachable
+            | TerminatorKind::UbCheck { .. } => {}
         }
     }
 }

compiler/rustc_const_eval/src/transform/check_consts/post_drop_elaboration.rs

@@ -118,7 +118,8 @@ impl<'tcx> Visitor<'tcx> for CheckLiveDrops<'_, 'tcx> {
             | mir::TerminatorKind::Return
             | mir::TerminatorKind::SwitchInt { .. }
             | mir::TerminatorKind::Unreachable
-            | mir::TerminatorKind::Yield { .. } => {}
+            | mir::TerminatorKind::Yield { .. }
+            | mir::TerminatorKind::UbCheck { .. } => {}
         }
     }
 }

compiler/rustc_const_eval/src/transform/validate.rs

@@ -393,7 +393,7 @@ impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
 
     fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
         match &terminator.kind {
-            TerminatorKind::Goto { target } => {
+            TerminatorKind::Goto { target } | TerminatorKind::UbCheck { target, .. } => {
                 self.check_edge(location, *target, EdgeKind::Normal);
             }
             TerminatorKind::SwitchInt { targets, discr: _ } => {
@@ -1298,7 +1298,8 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
             | TerminatorKind::UnwindResume
             | TerminatorKind::UnwindTerminate(_)
             | TerminatorKind::Return
-            | TerminatorKind::Unreachable => {}
+            | TerminatorKind::Unreachable
+            | TerminatorKind::UbCheck { .. } => {}
         }
 
         self.super_terminator(terminator, location);

compiler/rustc_middle/messages.ftl

@@ -14,9 +14,6 @@ middle_assert_divide_by_zero =
 
 middle_assert_gen_resume_after_panic = `gen` fn or block cannot be further iterated on after it panicked
 
-middle_assert_misaligned_ptr_deref =
-    misaligned pointer dereference: address must be a multiple of {$required} but is {$found}
-
 middle_assert_op_overflow =
     attempt to compute `{$left} {$op} {$right}`, which would overflow
 

compiler/rustc_middle/src/mir/pretty.rs

@@ -858,6 +858,9 @@ impl<'tcx> TerminatorKind<'tcx> {
                 }
                 write!(fmt, ", options({options:?}))")
             }
+            UbCheck { target: _, kind: UbCheckKind::PointerAlignment { pointer } } => {
+                write!(fmt, "pointer_alignment_check({pointer:?})")
+            }
         }
     }
 
@@ -866,7 +869,7 @@ impl<'tcx> TerminatorKind<'tcx> {
         use self::TerminatorKind::*;
         match *self {
             Return | UnwindResume | UnwindTerminate(_) | Unreachable | CoroutineDrop => vec![],
-            Goto { .. } => vec!["".into()],
+            Goto { .. } | UbCheck { .. } => vec!["".into()],
             SwitchInt { ref targets, .. } => targets
                 .values
                 .iter()

compiler/rustc_middle/src/mir/syntax.rs

@@ -588,7 +588,9 @@ impl CallSource {
 #[derive(Clone, TyEncodable, TyDecodable, Hash, HashStable, PartialEq, TypeFoldable, TypeVisitable)]
 pub enum TerminatorKind<'tcx> {
     /// Block has one successor; we continue execution there.
-    Goto { target: BasicBlock },
+    Goto {
+        target: BasicBlock,
+    },
 
     /// Switches based on the computed value.
     ///
@@ -655,7 +657,12 @@ pub enum TerminatorKind<'tcx> {
     /// The `replace` flag indicates whether this terminator was created as part of an assignment.
     /// This should only be used for diagnostic purposes, and does not have any operational
     /// meaning.
-    Drop { place: Place<'tcx>, target: BasicBlock, unwind: UnwindAction, replace: bool },
+    Drop {
+        place: Place<'tcx>,
+        target: BasicBlock,
+        unwind: UnwindAction,
+        replace: bool,
+    },
 
     /// Roughly speaking, evaluates the `func` operand and the arguments, and starts execution of
     /// the referred to function. The operand types must match the argument types of the function.
@@ -798,6 +805,17 @@ pub enum TerminatorKind<'tcx> {
         /// if and only if InlineAsmOptions::MAY_UNWIND is set.
         unwind: UnwindAction,
     },
+
+    UbCheck {
+        target: BasicBlock,
+        kind: UbCheckKind<'tcx>,
+    },
+}
+
+#[derive(Clone, TyEncodable, TyDecodable, Hash, HashStable, PartialEq, TypeFoldable, TypeVisitable)]
+#[derive(Debug)]
+pub enum UbCheckKind<'tcx> {
+    PointerAlignment { pointer: Operand<'tcx> },
 }
 
 impl TerminatorKind<'_> {
@@ -819,6 +837,7 @@ impl TerminatorKind<'_> {
             TerminatorKind::FalseEdge { .. } => "FalseEdge",
             TerminatorKind::FalseUnwind { .. } => "FalseUnwind",
             TerminatorKind::InlineAsm { .. } => "InlineAsm",
+            TerminatorKind::UbCheck { .. } => "UbCheck",
         }
     }
 }
@@ -885,7 +904,6 @@ pub enum AssertKind<O> {
     RemainderByZero(O),
     ResumedAfterReturn(CoroutineKind),
     ResumedAfterPanic(CoroutineKind),
-    MisalignedPointerDereference { required: O, found: O },
 }
 
 #[derive(Clone, Debug, PartialEq, TyEncodable, TyDecodable, Hash, HashStable)]