don't call `align_offset` during const eval, ever

2 files changed, 90 insertions, 50 deletions

diff --git a/compiler/rustc_const_eval/src/const_eval/machine.rs b/compiler/rustc_const_eval/src/const_eval/machine.rs
index 97d168ceb01..6cf8a581d71 100644
--- a/compiler/rustc_const_eval/src/const_eval/machine.rs
+++ b/compiler/rustc_const_eval/src/const_eval/machine.rs
@@ -1,11 +1,10 @@
 use rustc_hir::def::DefKind;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::PointerArithmetic;
-use rustc_middle::ty::layout::FnAbiOf;
+use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use std::borrow::Borrow;
 use std::hash::Hash;
-use std::ops::ControlFlow;
 
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::fx::IndexEntry;
@@ -20,8 +19,8 @@ use rustc_target::abi::{Align, Size};
 use rustc_target::spec::abi::Abi as CallAbi;
 
 use crate::interpret::{
-    self, compile_time_machine, AllocId, ConstAllocation, FnVal, Frame, ImmTy, InterpCx,
-    InterpResult, OpTy, PlaceTy, Pointer, Scalar, StackPopUnwind,
+    self, compile_time_machine, AllocId, ConstAllocation, Frame, ImmTy, InterpCx, InterpResult,
+    OpTy, PlaceTy, Pointer, Scalar, StackPopUnwind,
 };
 
 use super::error::*;
@@ -156,7 +155,6 @@ impl<'mir, 'tcx: 'mir> CompileTimeEvalContext<'mir, 'tcx> {
     fn hook_special_const_fn(
         &mut self,
         instance: ty::Instance<'tcx>,
-        _abi: CallAbi,
         args: &[OpTy<'tcx>],
         dest: &PlaceTy<'tcx>,
         ret: Option<mir::BasicBlock>,
@@ -194,24 +192,21 @@ impl<'mir, 'tcx: 'mir> CompileTimeEvalContext<'mir, 'tcx> {
 
             return Ok(Some(new_instance));
         } else if Some(def_id) == self.tcx.lang_items().align_offset_fn() {
-            // For align_offset, we replace the function call if the pointer has no address.
-            match self.align_offset(instance, args, dest, ret)? {
-                ControlFlow::Continue(()) => return Ok(Some(instance)),
-                ControlFlow::Break(()) => return Ok(None),
-            }
+            // For align_offset, we replace the function call entirely.
+            self.align_offset(instance, args, dest, ret)?;
+            return Ok(None);
         }
         Ok(Some(instance))
     }
 
-    /// `align_offset(ptr, target_align)` needs special handling in const eval, because the pointer
-    /// may not have an address.
+    /// This function replaces `align_offset(ptr, target_align)` in const eval, because the
+    /// pointer may not have an address.
     ///
-    /// If `ptr` does have a known address, then we return `CONTINUE` and the function call should
-    /// proceed as normal.
+    /// If `ptr` does have a known address, we forward it to [`Self::align_offset_impl`].
     ///
     /// If `ptr` doesn't have an address, but its underlying allocation's alignment is at most
-    /// `target_align`, then we call the function again with an dummy address relative to the
-    /// allocation.
+    /// `target_align`, then we call [`Self::align_offset_impl`] with an dummy address relative
+    /// to the allocation.
     ///
     /// If `ptr` doesn't have an address and `target_align` is stricter than the underlying
     /// allocation's alignment, then we return `usize::MAX` immediately.
@@ -221,50 +216,100 @@ impl<'mir, 'tcx: 'mir> CompileTimeEvalContext<'mir, 'tcx> {
         args: &[OpTy<'tcx>],
         dest: &PlaceTy<'tcx>,
         ret: Option<mir::BasicBlock>,
-    ) -> InterpResult<'tcx, ControlFlow<()>> {
+    ) -> InterpResult<'tcx> {
         assert_eq!(args.len(), 2);
 
         let ptr = self.read_pointer(&args[0])?;
         let target_align = self.read_scalar(&args[1])?.to_machine_usize(self)?;
 
+        let pointee_ty = instance.substs.type_at(0);
+        let stride = self.layout_of(pointee_ty)?.size.bytes();
+
         if !target_align.is_power_of_two() {
             throw_ub_format!("`align_offset` called with non-power-of-two align: {}", target_align);
         }
 
-        match self.ptr_try_get_alloc_id(ptr) {
+        let mut align_offset = match self.ptr_try_get_alloc_id(ptr) {
             Ok((alloc_id, offset, _extra)) => {
+                // Extract the address relative to a base that is definitely sufficiently aligned.
                 let (_size, alloc_align, _kind) = self.get_alloc_info(alloc_id);
 
                 if target_align <= alloc_align.bytes() {
-                    // Extract the address relative to the allocation base that is definitely
-                    // sufficiently aligned and call `align_offset` again.
-                    let addr = ImmTy::from_uint(offset.bytes(), args[0].layout).into();
-                    let align = ImmTy::from_uint(target_align, args[1].layout).into();
-
-                    let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty())?;
-                    self.eval_fn_call(
-                        FnVal::Instance(instance),
-                        (CallAbi::Rust, fn_abi),
-                        &[addr, align],
-                        false,
-                        dest,
-                        ret,
-                        StackPopUnwind::NotAllowed,
-                    )?;
-                    Ok(ControlFlow::BREAK)
+                    // The pointer *is* alignable in const. We use an address relative to the
+                    // allocation base that is definitely sufficiently aligned.
+                    let addr = offset.bytes();
+                    Self::align_offset_impl(addr, stride, target_align)
                 } else {
-                    // Not alignable in const, return `usize::MAX`.
-                    let usize_max = Scalar::from_machine_usize(self.machine_usize_max(), self);
-                    self.write_scalar(usize_max, dest)?;
-                    self.return_to_block(ret)?;
-                    Ok(ControlFlow::BREAK)
+                    // The pointer *is not* alignable in const, return `usize::MAX`.
+                    // (We clamp this to machine `usize` below.)
+                    u64::MAX
                 }
             }
-            Err(_addr) => {
-                // The pointer has an address, continue with function call.
-                Ok(ControlFlow::CONTINUE)
+            Err(addr) => {
+                // The pointer has a known address.
+                Self::align_offset_impl(addr, stride, target_align)
             }
+        };
+
+        let usize_max = self.machine_usize_max();
+        if align_offset > usize_max {
+            align_offset = usize_max;
+        }
+
+        self.write_scalar(Scalar::from_machine_usize(align_offset, self), dest)?;
+        self.return_to_block(ret)?;
+
+        Ok(())
+    }
+
+    /// Const eval implementation of `#[lang = "align_offset"]`.
+    /// See the runtime version for a detailed explanation how this works.
+    fn align_offset_impl(addr: u64, stride: u64, align: u64) -> u64 {
+        assert!(align.is_power_of_two());
+
+        let addr_mod_align = addr % align;
+
+        if addr_mod_align == 0 {
+            // The address is already sufficiently aligned.
+            return 0;
         }
+
+        if stride == 0 {
+            // The address cannot be aligned.
+            return u64::MAX;
+        }
+
+        let byte_offset = align - addr_mod_align;
+
+        if align % stride == 0 {
+            if byte_offset % stride == 0 {
+                return byte_offset / stride;
+            } else {
+                return u64::MAX;
+            }
+        }
+
+        // This only works, because `align` is a power of two.
+        let gcd = 1u64 << (stride | align).trailing_zeros();
+
+        if addr % gcd != 0 {
+            // The address cannot be aligned.
+            return u64::MAX;
+        }
+
+        let align2 = align / gcd;
+        let stride2 = (stride / gcd) % align2;
+
+        let mut stride_inv = 1u64;
+        let mut mod_gate = 2u64;
+        let mut overflow = false;
+        while !overflow && mod_gate < align2 {
+            stride_inv =
+                stride_inv.wrapping_mul(2u64.wrapping_sub(stride2.wrapping_mul(stride_inv)));
+            (mod_gate, overflow) = mod_gate.overflowing_mul(mod_gate);
+        }
+
+        byte_offset.wrapping_mul(stride_inv) % align2
     }
 
     /// See documentation on the `ptr_guaranteed_cmp` intrinsic.
@@ -367,7 +412,7 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
                 }
             }
 
-            let Some(new_instance) = ecx.hook_special_const_fn(instance, _abi, args, dest, ret)? else {
+            let Some(new_instance) = ecx.hook_special_const_fn(instance, args, dest, ret)? else {
                 return Ok(None);
             };
 
diff --git a/library/core/src/ptr/mod.rs b/library/core/src/ptr/mod.rs
index 1ebd9dca098..e762837ff90 100644
--- a/library/core/src/ptr/mod.rs
+++ b/library/core/src/ptr/mod.rs
@@ -1590,9 +1590,8 @@ pub unsafe fn write_volatile<T>(dst: *mut T, src: T) {
 /// than trying to adapt this to accommodate that change.
 ///
 /// Any questions go to @nagisa.
-// #[cfg(not(bootstrap))] -- Calling this function in a const context from the bootstrap
-// compiler will always cause an error.
 #[lang = "align_offset"]
+#[rustc_do_not_const_check] // hooked by const-eval
 pub(crate) const unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usize {
     // FIXME(#75598): Direct use of these intrinsics improves codegen significantly at opt-level <=
     // 1, where the method versions of these operations are not inlined.
@@ -1652,13 +1651,9 @@ pub(crate) const unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usiz
         inverse & m_minus_one
     }
 
+    let addr = p.addr();
     let stride = mem::size_of::<T>();
 
-    // SAFETY: At runtime transmuting a pointer to `usize` is always safe, because they have the
-    // same layout. During const eval we hook this function to ensure that the pointer always has
-    // an address (only the standard library can do this).
-    let addr = unsafe { mem::transmute(p) };
-
     // SAFETY: `a` is a power-of-two, therefore non-zero.
     let a_minus_one = unsafe { unchecked_sub(a, 1) };