miri: optimize zeroed alloc

Author: SpecificProtagonist

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -250,6 +250,20 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         self.insert_allocation(alloc, kind)
     }
 
+    pub fn allocate_zeroed_ptr(
+        &mut self,
+        size: Size,
+        align: Align,
+        kind: MemoryKind<M::MemoryKind>,
+    ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
+        let alloc = if M::PANIC_ON_ALLOC_FAIL {
+            Allocation::zeroed(size, align)
+        } else {
+            Allocation::try_zeroed(size, align)?
+        };
+        self.insert_allocation(alloc, kind)
+    }
+
     pub fn insert_allocation(
         &mut self,
         alloc: Allocation<M::Provenance, (), M::Bytes>,

compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -340,6 +340,49 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
         }
     }
 
+    fn zeroed_inner<R>(size: Size, align: Align, fail: impl FnOnce() -> R) -> Result<Self, R> {
+        // We raise an error if we cannot create the allocation on the host.
+        // This results in an error that can happen non-deterministically, since the memory
+        // available to the compiler can change between runs. Normally queries are always
+        // deterministic. However, we can be non-deterministic here because all uses of const
+        // evaluation (including ConstProp!) will make compilation fail (via hard error
+        // or ICE) upon encountering a `MemoryExhausted` error.
+        let bytes = Bytes::zeroed(size, align).ok_or_else(fail)?;
+
+        Ok(Allocation {
+            bytes,
+            provenance: ProvenanceMap::new(),
+            init_mask: InitMask::new(size, true),
+            align,
+            mutability: Mutability::Mut,
+            extra: (),
+        })
+    }
+
+    /// Try to create an Allocation of `size` zero-initialized bytes, failing if there is not enough memory
+    /// available to the compiler to do so.
+    pub fn try_zeroed<'tcx>(size: Size, align: Align) -> InterpResult<'tcx, Self> {
+        Self::zeroed_inner(size, align, || {
+            ty::tls::with(|tcx| tcx.dcx().delayed_bug("exhausted memory during interpretation"));
+            InterpErrorKind::ResourceExhaustion(ResourceExhaustionInfo::MemoryExhausted)
+        })
+        .into()
+    }
+
+    /// Try to create an Allocation of `size` zero-initialized bytes, panics if there is not enough memory
+    /// available to the compiler to do so.
+    pub fn zeroed(size: Size, align: Align) -> Self {
+        match Self::zeroed_inner(size, align, || {
+            panic!(
+                "interpreter ran out of memory: cannot create allocation of {} bytes",
+                size.bytes()
+            );
+        }) {
+            Ok(x) => x,
+            Err(x) => x,
+        }
+    }
+
     /// Add the extra.
     pub fn with_extra<Extra>(self, extra: Extra) -> Allocation<Prov, Extra, Bytes> {
         Allocation {

src/tools/miri/src/shims/alloc.rs

@@ -83,15 +83,11 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn malloc(&mut self, size: u64, zero_init: bool) -> InterpResult<'tcx, Pointer> {
         let this = self.eval_context_mut();
         let align = this.malloc_align(size);
-        let ptr = this.allocate_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into())?;
-        if zero_init {
-            // We just allocated this, the access is definitely in-bounds and fits into our address space.
-            this.write_bytes_ptr(
-                ptr.into(),
-                iter::repeat(0u8).take(usize::try_from(size).unwrap()),
-            )
-            .unwrap();
-        }
+        let ptr = if zero_init {
+            this.allocate_zeroed_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into())?
+        } else {
+            this.allocate_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into())?
+        };
         interp_ok(ptr.into())
     }
 

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

@@ -533,18 +533,11 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
                     this.check_rustc_alloc_request(size, align)?;
 
-                    let ptr = this.allocate_ptr(
+                    let ptr = this.allocate_zeroed_ptr(
                         Size::from_bytes(size),
                         Align::from_bytes(align).unwrap(),
                         MiriMemoryKind::Rust.into(),
                     )?;
-
-                    // We just allocated this, the access is definitely in-bounds.
-                    this.write_bytes_ptr(
-                        ptr.into(),
-                        iter::repeat(0u8).take(usize::try_from(size).unwrap()),
-                    )
-                    .unwrap();
                     this.write_pointer(ptr, dest)
                 });
             }

src/tools/miri/src/shims/unix/mem.rs

@@ -111,15 +111,12 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             return interp_ok(this.eval_libc("MAP_FAILED"));
         }
 
-        let ptr =
-            this.allocate_ptr(Size::from_bytes(map_length), align, MiriMemoryKind::Mmap.into())?;
-        // We just allocated this, the access is definitely in-bounds and fits into our address space.
         // mmap guarantees new mappings are zero-init.
-        this.write_bytes_ptr(
-            ptr.into(),
-            std::iter::repeat(0u8).take(usize::try_from(map_length).unwrap()),
-        )
-        .unwrap();
+        let ptr = this.allocate_zeroed_ptr(
+            Size::from_bytes(map_length),
+            align,
+            MiriMemoryKind::Mmap.into(),
+        )?;
 
         interp_ok(Scalar::from_pointer(ptr, this))
     }

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

@@ -258,17 +258,19 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 // Alignment is twice the pointer size.
                 // Source: <https://learn.microsoft.com/en-us/windows/win32/api/heapapi/nf-heapapi-heapalloc>
                 let align = this.tcx.pointer_size().bytes().strict_mul(2);
-                let ptr = this.allocate_ptr(
-                    Size::from_bytes(size),
-                    Align::from_bytes(align).unwrap(),
-                    MiriMemoryKind::WinHeap.into(),
-                )?;
-                if zero_init {
-                    this.write_bytes_ptr(
-                        ptr.into(),
-                        iter::repeat(0u8).take(usize::try_from(size).unwrap()),
-                    )?;
-                }
+                let ptr = if zero_init {
+                    this.allocate_zeroed_ptr(
+                        Size::from_bytes(size),
+                        Align::from_bytes(align).unwrap(),
+                        MiriMemoryKind::WinHeap.into(),
+                    )?
+                } else {
+                    this.allocate_ptr(
+                        Size::from_bytes(size),
+                        Align::from_bytes(align).unwrap(),
+                        MiriMemoryKind::WinHeap.into(),
+                    )?
+                };
                 this.write_pointer(ptr, dest)?;
             }
             "HeapFree" => {