avoid re-checking the offset while iterating an array/slice

4 files changed, 54 insertions, 15 deletions

diff --git a/compiler/rustc_const_eval/src/interpret/mod.rs b/compiler/rustc_const_eval/src/interpret/mod.rs
index 69eb22028fa..13664456987 100644
--- a/compiler/rustc_const_eval/src/interpret/mod.rs
+++ b/compiler/rustc_const_eval/src/interpret/mod.rs
@@ -26,7 +26,7 @@ pub use self::machine::{compile_time_machine, AllocMap, Machine, MayLeak, StackP
 pub use self::memory::{AllocKind, AllocRef, AllocRefMut, FnVal, Memory, MemoryKind};
 pub use self::operand::{ImmTy, Immediate, OpTy, Readable};
 pub use self::place::{MPlaceTy, MemPlaceMeta, PlaceTy, Writeable};
-pub use self::projection::Projectable;
+pub use self::projection::{OffsetMode, Projectable};
 pub use self::terminator::FnArg;
 pub use self::validity::{CtfeValidationMode, RefTracking};
 pub use self::visitor::ValueVisitor;
diff --git a/compiler/rustc_const_eval/src/interpret/operand.rs b/compiler/rustc_const_eval/src/interpret/operand.rs
index b504567989c..0f259f6c1a2 100644
--- a/compiler/rustc_const_eval/src/interpret/operand.rs
+++ b/compiler/rustc_const_eval/src/interpret/operand.rs
@@ -14,7 +14,8 @@ use rustc_target::abi::{self, Abi, Align, HasDataLayout, Size};
 
 use super::{
     alloc_range, from_known_layout, mir_assign_valid_types, AllocId, Frame, InterpCx, InterpResult,
-    MPlaceTy, Machine, MemPlace, MemPlaceMeta, PlaceTy, Pointer, Projectable, Provenance, Scalar,
+    MPlaceTy, Machine, MemPlace, MemPlaceMeta, OffsetMode, PlaceTy, Pointer, Projectable,
+    Provenance, Scalar,
 };
 
 /// An `Immediate` represents a single immediate self-contained Rust value.
@@ -297,6 +298,7 @@ impl<'tcx, Prov: Provenance> Projectable<'tcx, Prov> for ImmTy<'tcx, Prov> {
     fn offset_with_meta<'mir, M: Machine<'mir, 'tcx, Provenance = Prov>>(
         &self,
         offset: Size,
+        _mode: OffsetMode,
         meta: MemPlaceMeta<Prov>,
         layout: TyAndLayout<'tcx>,
         ecx: &InterpCx<'mir, 'tcx, M>,
@@ -391,12 +393,13 @@ impl<'tcx, Prov: Provenance> Projectable<'tcx, Prov> for OpTy<'tcx, Prov> {
     fn offset_with_meta<'mir, M: Machine<'mir, 'tcx, Provenance = Prov>>(
         &self,
         offset: Size,
+        mode: OffsetMode,
         meta: MemPlaceMeta<Prov>,
         layout: TyAndLayout<'tcx>,
         ecx: &InterpCx<'mir, 'tcx, M>,
     ) -> InterpResult<'tcx, Self> {
         match self.as_mplace_or_imm() {
-            Left(mplace) => Ok(mplace.offset_with_meta(offset, meta, layout, ecx)?.into()),
+            Left(mplace) => Ok(mplace.offset_with_meta(offset, mode, meta, layout, ecx)?.into()),
             Right(imm) => {
                 assert_matches!(meta, MemPlaceMeta::None); // no place to store metadata here
                 // Every part of an uninit is uninit.
diff --git a/compiler/rustc_const_eval/src/interpret/place.rs b/compiler/rustc_const_eval/src/interpret/place.rs
index 35ed899d3c8..f48417f207a 100644
--- a/compiler/rustc_const_eval/src/interpret/place.rs
+++ b/compiler/rustc_const_eval/src/interpret/place.rs
@@ -16,8 +16,8 @@ use rustc_target::abi::{Abi, Align, FieldIdx, HasDataLayout, Size, FIRST_VARIANT
 
 use super::{
     alloc_range, mir_assign_valid_types, AllocId, AllocRef, AllocRefMut, ImmTy, Immediate,
-    InterpCx, InterpResult, Machine, MemoryKind, OpTy, Operand, Pointer, PointerArithmetic,
-    Projectable, Provenance, Readable, Scalar,
+    InterpCx, InterpResult, Machine, MemoryKind, OffsetMode, OpTy, Operand, Pointer,
+    PointerArithmetic, Projectable, Provenance, Readable, Scalar,
 };
 
 #[derive(Copy, Clone, Hash, PartialEq, Eq, Debug)]
@@ -91,6 +91,7 @@ impl<Prov: Provenance> MemPlace<Prov> {
     fn offset_with_meta_<'mir, 'tcx, M: Machine<'mir, 'tcx, Provenance = Prov>>(
         self,
         offset: Size,
+        mode: OffsetMode,
         meta: MemPlaceMeta<Prov>,
         ecx: &InterpCx<'mir, 'tcx, M>,
     ) -> InterpResult<'tcx, Self> {
@@ -101,8 +102,12 @@ impl<Prov: Provenance> MemPlace<Prov> {
         if offset > ecx.data_layout().max_size_of_val() {
             throw_ub!(PointerArithOverflow);
         }
-        let offset: i64 = offset.bytes().try_into().unwrap();
-        let ptr = ecx.ptr_offset_inbounds(self.ptr, offset)?;
+        let ptr = match mode {
+            OffsetMode::Inbounds => {
+                ecx.ptr_offset_inbounds(self.ptr, offset.bytes().try_into().unwrap())?
+            }
+            OffsetMode::Wrapping => self.ptr.wrapping_offset(offset, ecx),
+        };
         Ok(MemPlace { ptr, meta })
     }
 }
@@ -194,12 +199,13 @@ impl<'tcx, Prov: Provenance> Projectable<'tcx, Prov> for MPlaceTy<'tcx, Prov> {
     fn offset_with_meta<'mir, M: Machine<'mir, 'tcx, Provenance = Prov>>(
         &self,
         offset: Size,
+        mode: OffsetMode,
         meta: MemPlaceMeta<Prov>,
         layout: TyAndLayout<'tcx>,
         ecx: &InterpCx<'mir, 'tcx, M>,
     ) -> InterpResult<'tcx, Self> {
         Ok(MPlaceTy {
-            mplace: self.mplace.offset_with_meta_(offset, meta, ecx)?,
+            mplace: self.mplace.offset_with_meta_(offset, mode, meta, ecx)?,
             align: self.align.restrict_for_offset(offset),
             layout,
         })
@@ -306,12 +312,13 @@ impl<'tcx, Prov: Provenance> Projectable<'tcx, Prov> for PlaceTy<'tcx, Prov> {
     fn offset_with_meta<'mir, M: Machine<'mir, 'tcx, Provenance = Prov>>(
         &self,
         offset: Size,
+        mode: OffsetMode,
         meta: MemPlaceMeta<Prov>,
         layout: TyAndLayout<'tcx>,
         ecx: &InterpCx<'mir, 'tcx, M>,
     ) -> InterpResult<'tcx, Self> {
         Ok(match self.as_mplace_or_local() {
-            Left(mplace) => mplace.offset_with_meta(offset, meta, layout, ecx)?.into(),
+            Left(mplace) => mplace.offset_with_meta(offset, mode, meta, layout, ecx)?.into(),
             Right((frame, local, old_offset)) => {
                 debug_assert!(layout.is_sized(), "unsized locals should live in memory");
                 assert_matches!(meta, MemPlaceMeta::None); // we couldn't store it anyway...
@@ -952,7 +959,13 @@ where
                     &mut Operand::Indirect(mplace) => mplace, // this already was an indirect local
                 };
                 if let Some(offset) = offset {
-                    whole_local.offset_with_meta_(offset, MemPlaceMeta::None, self)?
+                    // This offset is always inbounds, no need to check it again.
+                    whole_local.offset_with_meta_(
+                        offset,
+                        OffsetMode::Wrapping,
+                        MemPlaceMeta::None,
+                        self,
+                    )?
                 } else {
                     // Preserve wide place metadata, do not call `offset`.
                     whole_local
diff --git a/compiler/rustc_const_eval/src/interpret/projection.rs b/compiler/rustc_const_eval/src/interpret/projection.rs
index 5ff93d89485..f5666787fcd 100644
--- a/compiler/rustc_const_eval/src/interpret/projection.rs
+++ b/compiler/rustc_const_eval/src/interpret/projection.rs
@@ -19,6 +19,15 @@ use rustc_target::abi::{self, VariantIdx};
 
 use super::{InterpCx, InterpResult, MPlaceTy, Machine, MemPlaceMeta, OpTy, Provenance, Scalar};
 
+/// Describes the constraints placed on offset-projections.
+#[derive(Copy, Clone, Debug)]
+pub enum OffsetMode {
+    /// The offset has to be inbounds, like `ptr::offset`.
+    Inbounds,
+    /// No constraints, just wrap around the edge of the address space.
+    Wrapping,
+}
+
 /// A thing that we can project into, and that has a layout.
 pub trait Projectable<'tcx, Prov: Provenance>: Sized + std::fmt::Debug {
     /// Get the layout.
@@ -53,6 +62,7 @@ pub trait Projectable<'tcx, Prov: Provenance>: Sized + std::fmt::Debug {
     fn offset_with_meta<'mir, M: Machine<'mir, 'tcx, Provenance = Prov>>(
         &self,
         offset: Size,
+        mode: OffsetMode,
         meta: MemPlaceMeta<Prov>,
         layout: TyAndLayout<'tcx>,
         ecx: &InterpCx<'mir, 'tcx, M>,
@@ -65,7 +75,7 @@ pub trait Projectable<'tcx, Prov: Provenance>: Sized + std::fmt::Debug {
         ecx: &InterpCx<'mir, 'tcx, M>,
     ) -> InterpResult<'tcx, Self> {
         assert!(layout.is_sized());
-        self.offset_with_meta(offset, MemPlaceMeta::None, layout, ecx)
+        self.offset_with_meta(offset, OffsetMode::Inbounds, MemPlaceMeta::None, layout, ecx)
     }
 
     fn transmute<'mir, M: Machine<'mir, 'tcx, Provenance = Prov>>(
@@ -75,7 +85,7 @@ pub trait Projectable<'tcx, Prov: Provenance>: Sized + std::fmt::Debug {
     ) -> InterpResult<'tcx, Self> {
         assert!(self.layout().is_sized() && layout.is_sized());
         assert_eq!(self.layout().size, layout.size);
-        self.offset_with_meta(Size::ZERO, MemPlaceMeta::None, layout, ecx)
+        self.offset_with_meta(Size::ZERO, OffsetMode::Wrapping, MemPlaceMeta::None, layout, ecx)
     }
 
     /// Convert this to an `OpTy`. This might be an irreversible transformation, but is useful for
@@ -102,7 +112,17 @@ impl<'tcx, 'a, Prov: Provenance, P: Projectable<'tcx, Prov>> ArrayIterator<'tcx,
         ecx: &InterpCx<'mir, 'tcx, M>,
     ) -> InterpResult<'tcx, Option<(u64, P)>> {
         let Some(idx) = self.range.next() else { return Ok(None) };
-        Ok(Some((idx, self.base.offset(self.stride * idx, self.field_layout, ecx)?)))
+        // We use `Wrapping` here since the offset has already been checked when the iterator was created.
+        Ok(Some((
+            idx,
+            self.base.offset_with_meta(
+                self.stride * idx,
+                OffsetMode::Wrapping,
+                MemPlaceMeta::None,
+                self.field_layout,
+                ecx,
+            )?,
+        )))
     }
 }
 
@@ -157,7 +177,7 @@ where
             (MemPlaceMeta::None, offset)
         };
 
-        base.offset_with_meta(offset, meta, field_layout, self)
+        base.offset_with_meta(offset, OffsetMode::Inbounds, meta, field_layout, self)
     }
 
     /// Downcasting to an enum variant.
@@ -246,6 +266,9 @@ where
         };
         let len = base.len(self)?;
         let field_layout = base.layout().field(self, 0);
+        // Ensure that all the offsets are in-bounds once, up-front.
+        base.offset(len * stride, self.layout_of(self.tcx.types.unit).unwrap(), self)?;
+        // Create the iterator.
         Ok(ArrayIterator { base, range: 0..len, stride, field_layout, _phantom: PhantomData })
     }
 
@@ -303,7 +326,7 @@ where
         };
         let layout = self.layout_of(ty)?;
 
-        base.offset_with_meta(from_offset, meta, layout, self)
+        base.offset_with_meta(from_offset, OffsetMode::Inbounds, meta, layout, self)
     }
 
     /// Applying a general projection