miri: optimize zeroed alloc

Co-authored-by: Ralf Jung <[email protected]>

Author: SpecificProtagonist

src/shims/alloc.rs

@@ -1,5 +1,3 @@
-use std::iter;
-
 use rustc_abi::{Align, Size};
 use rustc_ast::expand::allocator::AllocatorKind;
 
@@ -80,18 +78,10 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         }
     }
 
-    fn malloc(&mut self, size: u64, zero_init: bool) -> InterpResult<'tcx, Pointer> {
+    fn malloc(&mut self, size: u64, init: AllocInit) -> 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 = this.allocate_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into(), init)?;
         interp_ok(ptr.into())
     }
 
@@ -115,6 +105,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 Size::from_bytes(size),
                 Align::from_bytes(align).unwrap(),
                 MiriMemoryKind::C.into(),
+                AllocInit::Uninit
             )?;
             this.write_pointer(ptr, &memptr)?;
             interp_ok(Scalar::from_i32(0))
@@ -134,7 +125,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let new_align = this.malloc_align(new_size);
         if this.ptr_is_null(old_ptr)? {
             // Here we must behave like `malloc`.
-            self.malloc(new_size, /*zero_init*/ false)
+            self.malloc(new_size, AllocInit::Uninit)
         } else {
             if new_size == 0 {
                 // C, in their infinite wisdom, made this UB.
@@ -147,6 +138,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     Size::from_bytes(new_size),
                     new_align,
                     MiriMemoryKind::C.into(),
+                    AllocInit::Uninit
                 )?;
                 interp_ok(new_ptr.into())
             }
@@ -187,6 +179,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     Size::from_bytes(size),
                     Align::from_bytes(align).unwrap(),
                     MiriMemoryKind::C.into(),
+                    AllocInit::Uninit
                 )?;
                 interp_ok(ptr.into())
             }

src/shims/foreign_items.rs

@@ -1,6 +1,5 @@
 use std::collections::hash_map::Entry;
 use std::io::Write;
-use std::iter;
 use std::path::Path;
 
 use rustc_abi::{Align, AlignFromBytesError, Size};
@@ -9,6 +8,7 @@ use rustc_ast::expand::allocator::alloc_error_handler_name;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::CrateNum;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::mir::interpret::AllocInit;
 use rustc_middle::ty::Ty;
 use rustc_middle::{mir, ty};
 use rustc_span::Symbol;
@@ -442,7 +442,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 let [size] = this.check_shim(abi, Conv::C, link_name, args)?;
                 let size = this.read_target_usize(size)?;
                 if size <= this.max_size_of_val().bytes() {
-                    let res = this.malloc(size, /*zero_init:*/ false)?;
+                    let res = this.malloc(size, AllocInit::Uninit)?;
                     this.write_pointer(res, dest)?;
                 } else {
                     // If this does not fit in an isize, return null and, on Unix, set errno.
@@ -457,7 +457,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 let items = this.read_target_usize(items)?;
                 let elem_size = this.read_target_usize(elem_size)?;
                 if let Some(size) = this.compute_size_in_bytes(Size::from_bytes(elem_size), items) {
-                    let res = this.malloc(size.bytes(), /*zero_init:*/ true)?;
+                    let res = this.malloc(size.bytes(), AllocInit::Zero)?;
                     this.write_pointer(res, dest)?;
                 } else {
                     // On size overflow, return null and, on Unix, set errno.
@@ -509,6 +509,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                         Size::from_bytes(size),
                         Align::from_bytes(align).unwrap(),
                         memory_kind.into(),
+                        AllocInit::Uninit
                     )?;
 
                     ecx.write_pointer(ptr, dest)
@@ -537,14 +538,8 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                         Size::from_bytes(size),
                         Align::from_bytes(align).unwrap(),
                         MiriMemoryKind::Rust.into(),
+                        AllocInit::Zero
                     )?;
-
-                    // 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)
                 });
             }
@@ -604,6 +599,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                         Size::from_bytes(new_size),
                         align,
                         MiriMemoryKind::Rust.into(),
+                        AllocInit::Uninit
                     )?;
                     this.write_pointer(new_ptr, dest)
                 });

src/shims/unix/fs.rs

@@ -1109,6 +1109,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     Size::from_bytes(size),
                     dirent_layout.align.abi,
                     MiriMemoryKind::Runtime.into(),
+                    AllocInit::Uninit
                 )?;
                 let entry: Pointer = entry.into();
 

src/shims/unix/linux/mem.rs

@@ -49,16 +49,8 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             Size::from_bytes(new_size),
             align,
             MiriMemoryKind::Mmap.into(),
+            AllocInit::Zero
         )?;
-        if let Some(increase) = new_size.checked_sub(old_size) {
-            // 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.wrapping_offset(Size::from_bytes(old_size), this).into(),
-                std::iter::repeat(0u8).take(usize::try_from(increase).unwrap()),
-            )
-            .unwrap();
-        }
 
         interp_ok(Scalar::from_pointer(ptr, this))
     }

src/shims/unix/mem.rs

@@ -111,15 +111,13 @@ 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_ptr(
+            Size::from_bytes(map_length),
+            align,
+            MiriMemoryKind::Mmap.into(),
+            // mmap guarantees new mappings are zero-init.
+            AllocInit::Zero
+        )?;
 
         interp_ok(Scalar::from_pointer(ptr, this))
     }

src/shims/windows/foreign_items.rs

@@ -253,22 +253,21 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 this.read_target_isize(handle)?;
                 let flags = this.read_scalar(flags)?.to_u32()?;
                 let size = this.read_target_usize(size)?;
-                let heap_zero_memory = 0x00000008; // HEAP_ZERO_MEMORY
-                let zero_init = (flags & heap_zero_memory) == heap_zero_memory;
+                const HEAP_ZERO_MEMORY: u32 = 0x00000008;
+                let init = if (flags & HEAP_ZERO_MEMORY) == HEAP_ZERO_MEMORY {
+                    AllocInit::Zero
+                } else {
+                    AllocInit::Uninit
+                };
                 // 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(),
+                    init
                 )?;
-                if zero_init {
-                    this.write_bytes_ptr(
-                        ptr.into(),
-                        iter::repeat(0u8).take(usize::try_from(size).unwrap()),
-                    )?;
-                }
                 this.write_pointer(ptr, dest)?;
             }
             "HeapFree" => {
@@ -300,6 +299,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     Size::from_bytes(size),
                     Align::from_bytes(align).unwrap(),
                     MiriMemoryKind::WinHeap.into(),
+                    AllocInit::Uninit
                 )?;
                 this.write_pointer(new_ptr, dest)?;
             }