miri: support 'promising' alignment for symbolic alignment check

Author: RalfJung

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -12,12 +12,12 @@ use rustc_middle::mir;
 use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_span::def_id::DefId;
-use rustc_target::abi::Size;
+use rustc_target::abi::{Align, Size};
 use rustc_target::spec::abi::Abi as CallAbi;
 
 use super::{
- AllocBytes, AllocId, AllocRange, Allocation, ConstAllocation, FnArg, Frame, ImmTy, InterpCx,
- InterpResult, MPlaceTy, MemoryKind, OpTy, PlaceTy, Pointer, Provenance,
+ AllocBytes, AllocId, AllocKind, AllocRange, Allocation, ConstAllocation, FnArg, Frame, ImmTy,
+ InterpCx, InterpResult, MPlaceTy, MemoryKind, Misalignment, OpTy, PlaceTy, Pointer, Provenance,
 };
 
 /// Data returned by Machine::stack_pop,
@@ -143,11 +143,18 @@ pub trait Machine<'mir, 'tcx: 'mir>: Sized {
  /// Whether memory accesses should be alignment-checked.
  fn enforce_alignment(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool;
 
- /// Whether, when checking alignment, we should look at the actual address and thus support
- /// custom alignment logic based on whatever the integer address happens to be.
- ///
- /// If this returns true, Provenance::OFFSET_IS_ADDR must be true.
- fn use_addr_for_alignment_check(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool;
+ /// Gives the machine a chance to detect more misalignment than the built-in checks would catch.
+ #[inline(always)]
+ fn alignment_check(
+ _ecx: &InterpCx<'mir, 'tcx, Self>,
+ _alloc_id: AllocId,
+ _alloc_align: Align,
+ _alloc_kind: AllocKind,
+ _offset: Size,
+ _align: Align,
+ ) -> Option<Misalignment> {
+ None
+ }
 
  /// Whether to enforce the validity invariant for a specific layout.
  fn enforce_validity(ecx: &InterpCx<'mir, 'tcx, Self>, layout: TyAndLayout<'tcx>) -> bool;
@@ -519,12 +526,6 @@ pub macro compile_time_machine(<$mir: lifetime, $tcx: lifetime>) {
  type FrameExtra = ();
  type Bytes = Box<[u8]>;
 
- #[inline(always)]
- fn use_addr_for_alignment_check(_ecx: &InterpCx<$mir, $tcx, Self>) -> bool {
- // We do not support `use_addr`.
- false
- }
-
  #[inline(always)]
  fn ignore_optional_overflow_checks(_ecx: &InterpCx<$mir, $tcx, Self>) -> bool {
  false

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -58,6 +58,7 @@ impl<T: fmt::Display> fmt::Display for MemoryKind<T> {
 }
 
 /// The return value of `get_alloc_info` indicates the "kind" of the allocation.
+#[derive(Copy, Clone, PartialEq, Debug)]
 pub enum AllocKind {
  /// A regular live data allocation.
  LiveData,
@@ -473,8 +474,12 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
  match self.ptr_try_get_alloc_id(ptr) {
  Err(addr) => offset_misalignment(addr, align),
  Ok((alloc_id, offset, _prov)) => {
- let (_size, alloc_align, _kind) = self.get_alloc_info(alloc_id);
- if M::use_addr_for_alignment_check(self) {
+ let (_size, alloc_align, kind) = self.get_alloc_info(alloc_id);
+ if let Some(misalign) =
+ M::alignment_check(self, alloc_id, alloc_align, kind, offset, align)
+ {
+ Some(misalign)
+ } else if M::Provenance::OFFSET_IS_ADDR {
  // `use_addr_for_alignment_check` can only be true if `OFFSET_IS_ADDR` is true.
  offset_misalignment(ptr.addr().bytes(), align)
  } else {

library/core/src/ptr/const_ptr.rs

@@ -1368,10 +1368,36 @@ impl<T: ?Sized> *const T {
  panic!("align_offset: align is not a power-of-two");
  }
 
- {
- // SAFETY: `align` has been checked to be a power of 2 above
- unsafe { align_offset(self, align) }
+ // SAFETY: `align` has been checked to be a power of 2 above
+ let ret = unsafe { align_offset(self, align) };
+
+ // Inform Miri that we want to consider the resulting pointer to be suitably aligned.
+ #[cfg(miri)]
+ if ret != usize::MAX {
+ fn runtime(ptr: *const (), align: usize) {
+ extern "Rust" {
+ pub fn miri_promise_symbolic_alignment(ptr: *const (), align: usize);
+ }
+
+ // SAFETY: this call is always safe.
+ unsafe {
+ miri_promise_symbolic_alignment(ptr, align);
+ }
+ }
+
+ const fn compiletime(_ptr: *const (), _align: usize) {}
+
+ // SAFETY: the extra behavior at runtime is for UB checks only.
+ unsafe {
+ intrinsics::const_eval_select(
+ (self.wrapping_add(ret).cast(), align),
+ compiletime,
+ runtime,
+ );
+ }
  }
+
+ ret
  }
 
  /// Returns whether the pointer is properly aligned for `T`.

library/core/src/ptr/mut_ptr.rs

@@ -1635,10 +1635,36 @@ impl<T: ?Sized> *mut T {
  panic!("align_offset: align is not a power-of-two");
  }
 
- {
- // SAFETY: `align` has been checked to be a power of 2 above
- unsafe { align_offset(self, align) }
+ // SAFETY: `align` has been checked to be a power of 2 above
+ let ret = unsafe { align_offset(self, align) };
+
+ // Inform Miri that we want to consider the resulting pointer to be suitably aligned.
+ #[cfg(miri)]
+ if ret != usize::MAX {
+ fn runtime(ptr: *const (), align: usize) {
+ extern "Rust" {
+ pub fn miri_promise_symbolic_alignment(ptr: *const (), align: usize);
+ }
+
+ // SAFETY: this call is always safe.
+ unsafe {
+ miri_promise_symbolic_alignment(ptr, align);
+ }
+ }
+
+ const fn compiletime(_ptr: *const (), _align: usize) {}
+
+ // SAFETY: the extra behavior at runtime is for UB checks only.
+ unsafe {
+ intrinsics::const_eval_select(
+ (self.wrapping_add(ret).cast_const().cast(), align),
+ compiletime,
+ runtime,
+ );
+ }
  }
+
+ ret
  }
 
  /// Returns whether the pointer is properly aligned for `T`.

library/core/src/slice/mod.rs

@@ -3868,6 +3868,18 @@ impl<T> [T] {
  } else {
  let (left, rest) = self.split_at(offset);
  let (us_len, ts_len) = rest.align_to_offsets::<U>();
+ // Inform Miri that we want to consider the "middle" pointer to be suitably aligned.
+ #[cfg(miri)]
+ {
+ extern "Rust" {
+ pub fn miri_promise_symbolic_alignment(ptr: *const (), align: usize);
+ }
+
+ // SAFETY: this call is always safe.
+ unsafe {
+ miri_promise_symbolic_alignment(rest.as_ptr().cast(), mem::align_of::<U>());
+ }
+ }
  // SAFETY: now `rest` is definitely aligned, so `from_raw_parts` below is okay,
  // since the caller guarantees that we can transmute `T` to `U` safely.
  unsafe {
@@ -3938,6 +3950,21 @@ impl<T> [T] {
  let (us_len, ts_len) = rest.align_to_offsets::<U>();
  let rest_len = rest.len();
  let mut_ptr = rest.as_mut_ptr();
+ // Inform Miri that we want to consider the "middle" pointer to be suitably aligned.
+ #[cfg(miri)]
+ {
+ extern "Rust" {
+ pub fn miri_promise_symbolic_alignment(ptr: *const (), align: usize);
+ }
+
+ // SAFETY: this call is always safe.
+ unsafe {
+ miri_promise_symbolic_alignment(
+ mut_ptr.cast() as *const (),
+ mem::align_of::<U>(),
+ );
+ }
+ }
  // We can't use `rest` again after this, that would invalidate its alias `mut_ptr`!
  // SAFETY: see comments for `align_to`.
  unsafe {

src/tools/miri/src/machine.rs

@@ -2,7 +2,7 @@
 //! `Machine` trait.
 
 use std::borrow::Cow;
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::fmt;
 use std::path::Path;
 use std::process;
@@ -309,11 +309,20 @@ pub struct AllocExtra<'tcx> {
  /// if this allocation is leakable. The backtrace is not
  /// pruned yet; that should be done before printing it.
  pub backtrace: Option<Vec<FrameInfo<'tcx>>>,
+ /// An offset inside this allocation that was deemed aligned even for symbolic alignment checks.
+ /// Invariant: the promised alignment will never be less than the native alignment of this allocation.
+ pub symbolic_alignment: Cell<Option<(Size, Align)>>,
 }
 
 impl VisitProvenance for AllocExtra<'_> {
  fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
- let AllocExtra { borrow_tracker, data_race, weak_memory, backtrace: _ } = self;
+ let AllocExtra {
+ borrow_tracker,
+ data_race,
+ weak_memory,
+ backtrace: _,
+ symbolic_alignment: _,
+ } = self;
 
  borrow_tracker.visit_provenance(visit);
  data_race.visit_provenance(visit);
@@ -907,8 +916,45 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
  }
 
  #[inline(always)]
- fn use_addr_for_alignment_check(ecx: &MiriInterpCx<'mir, 'tcx>) -> bool {
- ecx.machine.check_alignment == AlignmentCheck::Int
+ fn alignment_check(
+ ecx: &MiriInterpCx<'mir, 'tcx>,
+ alloc_id: AllocId,
+ alloc_align: Align,
+ alloc_kind: AllocKind,
+ offset: Size,
+ align: Align,
+ ) -> Option<Misalignment> {
+ if ecx.machine.check_alignment != AlignmentCheck::Symbolic {
+ // Just use the built-in check.
+ return None;
+ }
+ if alloc_kind != AllocKind::LiveData {
+ // Can't have any extra info here.
+ return None;
+ }
+ // Let's see which alignment we have been promised for this allocation.
+ let alloc_info = ecx.get_alloc_extra(alloc_id).unwrap(); // cannot fail since the allocation is live
+ let (promised_offset, promised_align) =
+ alloc_info.symbolic_alignment.get().unwrap_or((Size::ZERO, alloc_align));
+ if promised_align < align {
+ // Definitely not enough.
+ Some(Misalignment { has: promised_align, required: align })
+ } else {
+ // What's the offset between us and the promised alignment?
+ let distance = offset.bytes().wrapping_sub(promised_offset.bytes());
+ // That must also be aligned.
+ if distance % align.bytes() == 0 {
+ // All looking good!
+ None
+ } else {
+ // The biggest power of two through which `distance` is divisible.
+ let distance_pow2 = 1 << distance.trailing_zeros();
+ Some(Misalignment {
+ has: Align::from_bytes(distance_pow2).unwrap(),
+ required: align,
+ })
+ }
+ }
  }
 
  #[inline(always)]
@@ -1112,6 +1158,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
  data_race: race_alloc,
  weak_memory: buffer_alloc,
  backtrace,
+ symbolic_alignment: Cell::new(None),
  },
  |ptr| ecx.global_base_pointer(ptr),
  )?;

src/tools/miri/src/shims/foreign_items.rs

@@ -467,7 +467,7 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
  let ptr = this.read_pointer(ptr)?;
  let (alloc_id, _, _) = this.ptr_get_alloc_id(ptr).map_err(|_e| {
  err_machine_stop!(TerminationInfo::Abort(format!(
- "pointer passed to miri_get_alloc_id must not be dangling, got {ptr:?}"
+ "pointer passed to `miri_get_alloc_id` must not be dangling, got {ptr:?}"
  )))
  })?;
  this.write_scalar(Scalar::from_u64(alloc_id.0.get()), dest)?;
@@ -499,7 +499,7 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
  let (alloc_id, offset, _) = this.ptr_get_alloc_id(ptr)?;
  if offset != Size::ZERO {
  throw_unsup_format!(
- "pointer passed to miri_static_root must point to beginning of an allocated block"
+ "pointer passed to `miri_static_root` must point to beginning of an allocated block"
  );
  }
  this.machine.static_roots.push(alloc_id);
@@ -556,6 +556,39 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
  };
  }
 
+ // Promises that a pointer has a given symbolic alignment.
+ "miri_promise_symbolic_alignment" => {
+ let [ptr, align] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+ let ptr = this.read_pointer(ptr)?;
+ let align = this.read_target_usize(align)?;
+ let Ok(align) = Align::from_bytes(align) else {
+ throw_unsup_format!(
+ "`miri_promise_symbolic_alignment`: alignment must be a power of 2"
+ );
+ };
+ let (_, addr) = ptr.into_parts(); // we know the offset is absolute
+ if addr.bytes() % align.bytes() != 0 {
+ throw_unsup_format!(
+ "`miri_promise_symbolic_alignment`: pointer is not actually aligned"
+ );
+ }
+ if let Ok((alloc_id, offset, ..)) = this.ptr_try_get_alloc_id(ptr) {
+ let (_size, alloc_align, _kind) = this.get_alloc_info(alloc_id);
+ // Not `get_alloc_extra_mut`, need to handle read-only allocations!
+ let alloc_extra = this.get_alloc_extra(alloc_id)?;
+ // If the newly promised alignment is bigger than the native alignment of this
+ // allocation, and bigger than the previously promised alignment, then set it.
+ if align > alloc_align
+ && !alloc_extra
+ .symbolic_alignment
+ .get()
+ .is_some_and(|(_, old_align)| align <= old_align)
+ {
+ alloc_extra.symbolic_alignment.set(Some((offset, align)));
+ }
+ }
+ }
+
  // Standard C allocation
  "malloc" => {
  let [size] = this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;