make const align_offset useful

Author: Lukas Markeffsky

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -1,5 +1,7 @@
 use rustc_hir::def::DefKind;
 use rustc_middle::mir;
+use rustc_middle::mir::interpret::PointerArithmetic;
+use rustc_middle::ty::layout::{FnAbiOf, LayoutOf};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use std::borrow::Borrow;
 use std::collections::hash_map::Entry;
@@ -17,58 +19,12 @@ use rustc_target::abi::{Align, Size};
 use rustc_target::spec::abi::Abi as CallAbi;
 
 use crate::interpret::{
-    self, compile_time_machine, AllocId, ConstAllocation, Frame, ImmTy, InterpCx, InterpResult,
-    OpTy, PlaceTy, Pointer, Scalar, StackPopUnwind,
+    self, compile_time_machine, AllocId, ConstAllocation, FnVal, Frame, ImmTy, InterpCx,
+    InterpResult, OpTy, PlaceTy, Pointer, Scalar, StackPopUnwind,
 };
 
 use super::error::*;
 
-impl<'mir, 'tcx> InterpCx<'mir, 'tcx, CompileTimeInterpreter<'mir, 'tcx>> {
-    /// "Intercept" a function call to a panic-related function
-    /// because we have something special to do for it.
-    /// If this returns successfully (`Ok`), the function should just be evaluated normally.
-    fn hook_special_const_fn(
-        &mut self,
-        instance: ty::Instance<'tcx>,
-        args: &[OpTy<'tcx>],
-    ) -> InterpResult<'tcx, Option<ty::Instance<'tcx>>> {
-        // All `#[rustc_do_not_const_check]` functions should be hooked here.
-        let def_id = instance.def_id();
-
-        if Some(def_id) == self.tcx.lang_items().panic_display()
-            || Some(def_id) == self.tcx.lang_items().begin_panic_fn()
-        {
-            // &str or &&str
-            assert!(args.len() == 1);
-
-            let mut msg_place = self.deref_operand(&args[0])?;
-            while msg_place.layout.ty.is_ref() {
-                msg_place = self.deref_operand(&msg_place.into())?;
-            }
-
-            let msg = Symbol::intern(self.read_str(&msg_place)?);
-            let span = self.find_closest_untracked_caller_location();
-            let (file, line, col) = self.location_triple_for_span(span);
-            return Err(ConstEvalErrKind::Panic { msg, file, line, col }.into());
-        } else if Some(def_id) == self.tcx.lang_items().panic_fmt() {
-            // For panic_fmt, call const_panic_fmt instead.
-            if let Some(const_panic_fmt) = self.tcx.lang_items().const_panic_fmt() {
-                return Ok(Some(
-                    ty::Instance::resolve(
-                        *self.tcx,
-                        ty::ParamEnv::reveal_all(),
-                        const_panic_fmt,
-                        self.tcx.intern_substs(&[]),
-                    )
-                    .unwrap()
-                    .unwrap(),
-                ));
-            }
-        }
-        Ok(None)
-    }
-}
-
 /// Extra machine state for CTFE, and the Machine instance
 pub struct CompileTimeInterpreter<'mir, 'tcx> {
     /// For now, the number of terminators that can be evaluated before we throw a resource
@@ -191,6 +147,120 @@ impl interpret::MayLeak for ! {
 }
 
 impl<'mir, 'tcx: 'mir> CompileTimeEvalContext<'mir, 'tcx> {
+    /// "Intercept" a function call to a panic-related function
+    /// because we have something special to do for it.
+    /// If this returns successfully (`Ok`), the function should just be evaluated normally.
+    fn hook_special_const_fn(
+        &mut self,
+        instance: ty::Instance<'tcx>,
+        _abi: CallAbi,
+        args: &[OpTy<'tcx>],
+        dest: &PlaceTy<'tcx>,
+        ret: Option<mir::BasicBlock>,
+    ) -> InterpResult<'tcx, Option<ty::Instance<'tcx>>> {
+        // All `#[rustc_do_not_const_check]` functions should be hooked here.
+        let def_id = instance.def_id();
+
+        if Some(def_id) == self.tcx.lang_items().panic_display()
+            || Some(def_id) == self.tcx.lang_items().begin_panic_fn()
+        {
+            // &str or &&str
+            assert!(args.len() == 1);
+
+            let mut msg_place = self.deref_operand(&args[0])?;
+            while msg_place.layout.ty.is_ref() {
+                msg_place = self.deref_operand(&msg_place.into())?;
+            }
+
+            let msg = Symbol::intern(self.read_str(&msg_place)?);
+            let span = self.find_closest_untracked_caller_location();
+            let (file, line, col) = self.location_triple_for_span(span);
+            return Err(ConstEvalErrKind::Panic { msg, file, line, col }.into());
+        } else if Some(def_id) == self.tcx.lang_items().panic_fmt() {
+            // For panic_fmt, call const_panic_fmt instead.
+            let Some(const_def_id) = self.tcx.lang_items().const_panic_fmt() else {
+                bug!("`const_panic_fmt` must be defined to call `panic_fmt` in const eval")
+            };
+            let new_instance = ty::Instance::resolve(
+                *self.tcx,
+                ty::ParamEnv::reveal_all(),
+                const_def_id,
+                instance.substs,
+            )
+            .unwrap()
+            .unwrap();
+
+            return Ok(Some(new_instance));
+        } else if Some(def_id) == self.tcx.lang_items().align_offset_fn() {
+            // For align_offset, we either call const_align_offset or return usize::MAX directly.
+
+            let Some(const_def_id) = self.tcx.lang_items().const_align_offset_fn() else {
+                bug!("`const_align_offset` must be defined to call `align_offset` in const eval")
+            };
+            let const_instance = ty::Instance::resolve(
+                *self.tcx,
+                ty::ParamEnv::reveal_all(),
+                const_def_id,
+                instance.substs,
+            )
+            .unwrap()
+            .unwrap();
+
+            self.align_offset(const_instance, args, dest, ret)?;
+
+            return Ok(None);
+        }
+        Ok(Some(instance))
+    }
+
+    fn align_offset(
+        &mut self,
+        const_instance: ty::Instance<'tcx>,
+        args: &[OpTy<'tcx>],
+        dest: &PlaceTy<'tcx>,
+        ret: Option<mir::BasicBlock>,
+    ) -> 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)?;
+
+        if !target_align.is_power_of_two() {
+            throw_ub_format!("`align_offset` called with non-power-of-two align: {}", target_align);
+        }
+
+        let addr = match self.ptr_try_get_alloc_id(ptr) {
+            Ok((alloc_id, offset, _extra)) => {
+                let (_size, alloc_align, _kind) = self.get_alloc_info(alloc_id);
+
+                if target_align > alloc_align.bytes() {
+                    let usize_max = Scalar::from_machine_usize(self.machine_usize_max(), self);
+                    self.write_scalar(usize_max, dest)?;
+                    self.return_to_block(ret)?;
+                    return Ok(());
+                } else {
+                    offset.bytes()
+                }
+            }
+            Err(addr) => addr,
+        };
+
+        let usize_layout = self.layout_of(self.tcx.types.usize)?;
+        let addr = ImmTy::from_uint(addr, usize_layout).into();
+        let align = ImmTy::from_uint(target_align, usize_layout).into();
+
+        let fn_abi = self.fn_abi_of_instance(const_instance, ty::List::empty())?;
+        self.eval_fn_call(
+            FnVal::Instance(const_instance),
+            (CallAbi::Rust, fn_abi),
+            &[addr, align],
+            false,
+            dest,
+            ret,
+            StackPopUnwind::NotAllowed,
+        )
+    }
+
     /// See documentation on the `ptr_guaranteed_cmp` intrinsic.
     fn guaranteed_cmp(&mut self, a: Scalar, b: Scalar) -> InterpResult<'tcx, u8> {
         Ok(match (a, b) {
@@ -271,8 +341,8 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
         instance: ty::Instance<'tcx>,
         _abi: CallAbi,
         args: &[OpTy<'tcx>],
-        _dest: &PlaceTy<'tcx>,
-        _ret: Option<mir::BasicBlock>,
+        dest: &PlaceTy<'tcx>,
+        ret: Option<mir::BasicBlock>,
         _unwind: StackPopUnwind, // unwinding is not supported in consts
     ) -> InterpResult<'tcx, Option<(&'mir mir::Body<'tcx>, ty::Instance<'tcx>)>> {
         debug!("find_mir_or_eval_fn: {:?}", instance);
@@ -291,7 +361,11 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
                 }
             }
 
-            if let Some(new_instance) = ecx.hook_special_const_fn(instance, args)? {
+            let Some(new_instance) = ecx.hook_special_const_fn(instance, _abi, args, dest, ret)? else {
+                return Ok(None);
+            };
+
+            if new_instance != instance {
                 // We call another const fn instead.
                 // However, we return the *original* instance to make backtraces work out
                 // (and we hope this does not confuse the FnAbi checks too much).
@@ -300,13 +374,14 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
                     new_instance,
                     _abi,
                     args,
-                    _dest,
-                    _ret,
+                    dest,
+                    ret,
                     _unwind,
                 )?
                 .map(|(body, _instance)| (body, instance)));
             }
         }
+
         // This is a const fn. Call it.
         Ok(Some((ecx.load_mir(instance.def, None)?, instance)))
     }

compiler/rustc_hir/src/lang_items.rs

@@ -283,7 +283,8 @@ language_item_table! {
     MaybeUninit,             sym::maybe_uninit,        maybe_uninit,               Target::Union,          GenericRequirement::None;
 
     /// Align offset for stride != 1; must not panic.
-    AlignOffset,             sym::align_offset,        align_offset_fn,            Target::Fn,             GenericRequirement::None;
+    AlignOffset,             sym::align_offset,        align_offset_fn,            Target::Fn,             GenericRequirement::Exact(1);
+    ConstAlignOffset,        sym::const_align_offset,  const_align_offset_fn,      Target::Fn,             GenericRequirement::Exact(1);
 
     Termination,             sym::termination,         termination,                Target::Trait,          GenericRequirement::None;
 

compiler/rustc_span/src/symbol.rs

@@ -510,6 +510,7 @@ symbols! {
         concat_macro,
         conservative_impl_trait,
         console,
+        const_align_offset,
         const_allocate,
         const_async_blocks,
         const_compare_raw_pointers,

library/core/src/ptr/const_ptr.rs

@@ -1303,6 +1303,21 @@ impl<T: ?Sized> *const T {
     /// ```
     #[stable(feature = "align_offset", since = "1.36.0")]
     #[rustc_const_unstable(feature = "const_align_offset", issue = "90962")]
+    #[cfg(not(bootstrap))]
+    pub const fn align_offset(self, align: usize) -> usize
+    where
+        T: Sized,
+    {
+        assert!(align.is_power_of_two(), "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) }
+    }
+
+    #[stable(feature = "align_offset", since = "1.36.0")]
+    #[rustc_const_unstable(feature = "const_align_offset", issue = "90962")]
+    #[allow(missing_docs)]
+    #[cfg(bootstrap)]
     pub const fn align_offset(self, align: usize) -> usize
     where
         T: Sized,

library/core/src/ptr/mod.rs

@@ -1554,20 +1554,46 @@ pub unsafe fn write_volatile<T>(dst: *mut T, src: T) {
 
 /// Align pointer `p`.
 ///
-/// Calculate offset (in terms of elements of `stride` stride) that has to be applied
+/// Calculate offset (in terms of elements of `size_of::<T>()` stride) that has to be applied
 /// to pointer `p` so that pointer `p` would get aligned to `a`.
 ///
+/// # Safety
+/// `a` must be a power of two.
+#[lang = "align_offset"]
+#[rustc_do_not_const_check]
+#[cfg(not(bootstrap))]
+pub(crate) const unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usize {
+    // SAFETY: Caller ensures that `a` is a power of two.
+    unsafe { const_align_offset::<T>(p.addr(), a) }
+}
+
+#[lang = "align_offset"]
+#[cfg(bootstrap)]
+pub(crate) unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usize {
+    // SAFETY: Caller ensures that `a` is a power of two.
+    unsafe { const_align_offset::<T>(p.addr(), a) }
+}
+
+/// Align address `addr`.
+///
+/// Calculate offset (in terms of elements of `size_of::<T>()` stride) that has to be applied
+/// to address `addr` so that `addr` would get aligned to `a`.
+///
 /// Note: This implementation has been carefully tailored to not panic. It is UB for this to panic.
 /// The only real change that can be made here is change of `INV_TABLE_MOD_16` and associated
 /// constants.
 ///
+/// # Safety
+/// `a` must be a power of two.
+///
 /// If we ever decide to make it possible to call the intrinsic with `a` that is not a
 /// power-of-two, it will probably be more prudent to just change to a naive implementation rather
 /// than trying to adapt this to accommodate that change.
 ///
 /// Any questions go to @nagisa.
-#[lang = "align_offset"]
-pub(crate) unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usize {
+#[cfg_attr(not(bootstrap), lang = "const_align_offset")]
+#[rustc_allow_const_fn_unstable(const_exact_div)]
+pub(crate) const unsafe fn const_align_offset<T: Sized>(addr: usize, 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.
     use intrinsics::{
@@ -1584,7 +1610,7 @@ pub(crate) unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usize {
     ///
     /// Implementation of this function shall not panic. Ever.
     #[inline]
-    unsafe fn mod_inv(x: usize, m: usize) -> usize {
+    const unsafe fn mod_inv(x: usize, m: usize) -> usize {
         /// Multiplicative modular inverse table modulo 2⁴ = 16.
         ///
         /// Note, that this table does not contain values where inverse does not exist (i.e., for
@@ -1624,7 +1650,6 @@ pub(crate) unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usize {
         }
     }
 
-    let addr = p.addr();
     let stride = mem::size_of::<T>();
     // SAFETY: `a` is a power-of-two, therefore non-zero.
     let a_minus_one = unsafe { unchecked_sub(a, 1) };

library/core/src/ptr/mut_ptr.rs

@@ -1574,6 +1574,21 @@ impl<T: ?Sized> *mut T {
     /// ```
     #[stable(feature = "align_offset", since = "1.36.0")]
     #[rustc_const_unstable(feature = "const_align_offset", issue = "90962")]
+    #[cfg(not(bootstrap))]
+    pub const fn align_offset(self, align: usize) -> usize
+    where
+        T: Sized,
+    {
+        assert!(align.is_power_of_two(), "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) }
+    }
+
+    #[stable(feature = "align_offset", since = "1.36.0")]
+    #[rustc_const_unstable(feature = "const_align_offset", issue = "90962")]
+    #[allow(missing_docs)]
+    #[cfg(bootstrap)]
     pub const fn align_offset(self, align: usize) -> usize
     where
         T: Sized,