Auto merge of rust-lang#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`

Author: bors

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 { .. } => {}
         }
     }
 }