librustc_mir: Add support for const fn offset/arith_offset

Miri's pointer_offset_inbounds implementation has been moved into
librustc_mir as ptr_offset_inbounds (to avoid breaking miri on a
nightly update). The comments have been slightly reworked to better
match `offset`'s external documentation about what causes UB.

The intrinsic implementations are taken directly from miri.

Signed-off-by: Joe Richey <joerichey@google.com>

2 files changed, 53 insertions, 2 deletions

diff --git a/src/librustc_mir/interpret/intrinsics.rs b/src/librustc_mir/interpret/intrinsics.rs
index fc4be82ad90..0db9c702658 100644
--- a/src/librustc_mir/interpret/intrinsics.rs
+++ b/src/librustc_mir/interpret/intrinsics.rs
@@ -10,11 +10,11 @@ use rustc_middle::mir::{
 };
 use rustc_middle::ty;
 use rustc_middle::ty::subst::SubstsRef;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{Ty, TyCtxt};
 use rustc_span::symbol::{sym, Symbol};
 use rustc_target::abi::{Abi, LayoutOf as _, Primitive, Size};
 
-use super::{ImmTy, InterpCx, Machine, OpTy, PlaceTy};
+use super::{CheckInAllocMsg, ImmTy, InterpCx, Machine, OpTy, PlaceTy};
 
 mod caller_location;
 mod type_name;
@@ -279,7 +279,24 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 let result = Scalar::from_uint(truncated_bits, layout.size);
                 self.write_scalar(result, dest)?;
             }
+            sym::offset => {
+                let ptr = self.read_scalar(args[0])?.not_undef()?;
+                let offset_count = self.read_scalar(args[1])?.to_machine_isize(self)?;
+                let pointee_ty = substs.type_at(0);
 
+                let offset_ptr = self.ptr_offset_inbounds(ptr, pointee_ty, offset_count)?;
+                self.write_scalar(offset_ptr, dest)?;
+            }
+            sym::arith_offset => {
+                let ptr = self.read_scalar(args[0])?.not_undef()?;
+                let offset_count = self.read_scalar(args[1])?.to_machine_isize(self)?;
+                let pointee_ty = substs.type_at(0);
+
+                let pointee_size = i64::try_from(self.layout_of(pointee_ty)?.size.bytes()).unwrap();
+                let offset_bytes = offset_count.wrapping_mul(pointee_size);
+                let offset_ptr = ptr.ptr_wrapping_signed_offset(offset_bytes, self);
+                self.write_scalar(offset_ptr, dest)?;
+            }
             sym::ptr_offset_from => {
                 let a = self.read_immediate(args[0])?.to_scalar()?;
                 let b = self.read_immediate(args[1])?.to_scalar()?;
@@ -409,4 +426,36 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         // `Rem` says this is all right, so we can let `Div` do its job.
         self.binop_ignore_overflow(BinOp::Div, a, b, dest)
     }
+
+    /// Offsets a pointer by some multiple of its type, returning an error if the pointer leaves its
+    /// allocation. For integer pointers, we consider each of them their own tiny allocation of size
+    /// 0, so offset-by-0 (and only 0) is okay -- except that NULL cannot be offset by _any_ value.
+    pub fn ptr_offset_inbounds(
+        &self,
+        ptr: Scalar<M::PointerTag>,
+        pointee_ty: Ty<'tcx>,
+        offset_count: i64,
+    ) -> InterpResult<'tcx, Scalar<M::PointerTag>> {
+        let pointee_size = i64::try_from(self.layout_of(pointee_ty)?.size.bytes()).unwrap();
+        // The computed offset, in bytes, cannot overflow an isize.
+        let offset_bytes = offset_count
+            .checked_mul(pointee_size)
+            .ok_or(err_ub_format!("inbounds pointer arithmetic: overflow computing offset"))?;
+        // The offset being in bounds cannot rely on "wrapping around" the address space.
+        // So, first rule out overflows in the pointer arithmetic.
+        let offset_ptr = ptr.ptr_signed_offset(offset_bytes, self)?;
+        // ptr and offset_ptr must be in bounds of the same allocated object. This means all of the
+        // memory between these pointers must be accessible. Note that we do not require the
+        // pointers to be properly aligned (unlike a read/write operation).
+        let min_ptr = if offset_bytes >= 0 { ptr } else { offset_ptr };
+        let size = offset_bytes.checked_abs().unwrap();
+        // This call handles checking for integer/NULL pointers.
+        self.memory.check_ptr_access_align(
+            min_ptr,
+            Size::from_bytes(size),
+            None,
+            CheckInAllocMsg::InboundsTest,
+        )?;
+        Ok(offset_ptr)
+    }
 }
diff --git a/src/librustc_span/symbol.rs b/src/librustc_span/symbol.rs
index 6a6098710e8..9b055beb99d 100644
--- a/src/librustc_span/symbol.rs
+++ b/src/librustc_span/symbol.rs
@@ -147,6 +147,7 @@ symbols! {
         Arc,
         Arguments,
         ArgumentV1,
+        arith_offset,
         arm_target_feature,
         asm,
         assert,
@@ -516,6 +517,7 @@ symbols! {
         not,
         note,
         object_safe_for_dispatch,
+        offset,
         Ok,
         omit_gdb_pretty_printer_section,
         on,