Move MemoryImageSource::map_at to mmap module (bytecodealliance#9687)

* Simplify mmap interface slightly

Return a single `SendSyncPtr` -- the platform-independent context converts to
the various raw pointer types as desired. This simplifies upcoming work where I
wanted to return a `SendSyncPtr`.

* Move MemoryImageSource::map_at to mmap module

This is part of the work to centralize memory management into the `mmap`
module. This commit introduces a few structures which aid in that process, and
starts converting one of the functions (`MemoryImageSource::map_at`) into this
module.

The structures introduced are:

* `MemoryBase`: `RuntimeLinearMemory::base_ptr` is now
  `RuntimeLinearMemory::base`, which returns a `MemoryBase`. This is either a
  raw pointer or an `MmapOffset` as described below.

* `MmapOffset`: A combination of a reference to an mmap and an offset into it.
  Logically represents a pointer into a mapped section of memory.

In future work, we'll move more image-mapping code over to `Mmap` instances.

crates/wasmtime/src/runtime/trampoline/memory.rs

@@ -3,7 +3,7 @@ use crate::prelude::*;
 use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::{
     CompiledModuleId, GcHeapAllocationIndex, Imports, InstanceAllocationRequest, InstanceAllocator,
-    InstanceAllocatorImpl, Memory, MemoryAllocationIndex, ModuleRuntimeInfo,
+    InstanceAllocatorImpl, Memory, MemoryAllocationIndex, MemoryBase, ModuleRuntimeInfo,
     OnDemandInstanceAllocator, RuntimeLinearMemory, RuntimeMemoryCreator, SharedMemory, StorePtr,
     Table, TableAllocationIndex,
 };
@@ -89,8 +89,8 @@ impl RuntimeLinearMemory for LinearMemoryProxy {
         self.mem.grow_to(new_size)
     }
 
-    fn base_ptr(&self) -> *mut u8 {
-        self.mem.as_ptr()
+    fn base(&self) -> MemoryBase {
+        MemoryBase::new_raw(self.mem.as_ptr())
     }
 }
 

crates/wasmtime/src/runtime/vm.rs

@@ -73,7 +73,7 @@ pub use crate::runtime::vm::instance::{
 };
 pub use crate::runtime::vm::interpreter::*;
 pub use crate::runtime::vm::memory::{
-    Memory, RuntimeLinearMemory, RuntimeMemoryCreator, SharedMemory,
+    Memory, MemoryBase, RuntimeLinearMemory, RuntimeMemoryCreator, SharedMemory,
 };
 pub use crate::runtime::vm::mmap_vec::MmapVec;
 pub use crate::runtime::vm::mpk::MpkEnabled;
@@ -107,7 +107,7 @@ mod mmap;
 cfg_if::cfg_if! {
     if #[cfg(feature = "signals-based-traps")] {
         pub use crate::runtime::vm::byte_count::*;
-        pub use crate::runtime::vm::mmap::Mmap;
+        pub use crate::runtime::vm::mmap::{Mmap, MmapOffset};
         pub use self::cow::{MemoryImage, MemoryImageSlot, ModuleMemoryImages};
     } else {
         pub use self::cow_disabled::{MemoryImage, MemoryImageSlot, ModuleMemoryImages};

crates/wasmtime/src/runtime/vm/cow.rs

@@ -8,11 +8,10 @@
 use super::sys::DecommitBehavior;
 use crate::prelude::*;
 use crate::runtime::vm::sys::vm::{self, MemoryImageSource};
-use crate::runtime::vm::{host_page_size, HostAlignedByteCount, MmapVec, SendSyncPtr};
+use crate::runtime::vm::{host_page_size, HostAlignedByteCount, MmapOffset, MmapVec};
 use alloc::sync::Arc;
-use core::ffi::c_void;
 use core::ops::Range;
-use core::ptr::{self, NonNull};
+use core::ptr;
 use wasmtime_environ::{DefinedMemoryIndex, MemoryInitialization, Module, PrimaryMap, Tunables};
 
 /// Backing images for memories in a module.
@@ -131,13 +130,13 @@ impl MemoryImage {
         Ok(None)
     }
 
-    unsafe fn map_at(&self, base: *mut u8) -> Result<()> {
-        self.source.map_at(
-            base.add(self.linear_memory_offset.byte_count()),
-            self.len.byte_count(),
+    unsafe fn map_at(&self, mmap_base: &MmapOffset) -> Result<()> {
+        mmap_base.map_image_at(
+            &self.source,
             self.source_offset,
-        )?;
-        Ok(())
+            self.linear_memory_offset,
+            self.len,
+        )
     }
 
     unsafe fn remap_as_zeros_at(&self, base: *mut u8) -> Result<()> {
@@ -283,10 +282,9 @@ impl ModuleMemoryImages {
 /// with a fresh zero'd mmap, meaning that reuse is effectively not supported.
 #[derive(Debug)]
 pub struct MemoryImageSlot {
-    /// The base address in virtual memory of the actual heap memory.
-    ///
-    /// Bytes at this address are what is seen by the Wasm guest code.
-    base: SendSyncPtr<u8>,
+    /// The mmap and offset within it that contains the linear memory for this
+    /// slot.
+    base: MmapOffset,
 
     /// The maximum static memory size which `self.accessible` can grow to.
     static_size: usize,
@@ -337,12 +335,12 @@ impl MemoryImageSlot {
     /// and all memory from `accessible` from `static_size` should be mapped as
     /// `PROT_NONE` backed by zero-bytes.
     pub(crate) fn create(
-        base_addr: *mut c_void,
+        base: MmapOffset,
         accessible: HostAlignedByteCount,
         static_size: usize,
     ) -> Self {
         MemoryImageSlot {
-            base: NonNull::new(base_addr.cast()).unwrap().into(),
+            base,
             static_size,
             accessible,
             image: None,
@@ -463,7 +461,7 @@ impl MemoryImageSlot {
                 );
                 if !image.len.is_zero() {
                     unsafe {
-                        image.map_at(self.base.as_ptr())?;
+                        image.map_at(&self.base)?;
                     }
                 }
             }
@@ -480,7 +478,7 @@ impl MemoryImageSlot {
     pub(crate) fn remove_image(&mut self) -> Result<()> {
         if let Some(image) = &self.image {
             unsafe {
-                image.remap_as_zeros_at(self.base.as_ptr())?;
+                image.remap_as_zeros_at(self.base.as_mut_ptr())?;
             }
             self.image = None;
         }
@@ -589,7 +587,7 @@ impl MemoryImageSlot {
 
                     // This is memset (1)
                     ptr::write_bytes(
-                        self.base.as_ptr(),
+                        self.base.as_mut_ptr(),
                         0u8,
                         image.linear_memory_offset.byte_count(),
                     );
@@ -603,7 +601,7 @@ impl MemoryImageSlot {
 
                     // This is memset (3)
                     ptr::write_bytes(
-                        self.base.as_ptr().add(image_end.byte_count()),
+                        self.base.as_mut_ptr().add(image_end.byte_count()),
                         0u8,
                         remaining_memset.byte_count(),
                     );
@@ -639,7 +637,7 @@ impl MemoryImageSlot {
                     // Note that the memset may be zero bytes here.
 
                     // This is memset (1)
-                    ptr::write_bytes(self.base.as_ptr(), 0u8, keep_resident.byte_count());
+                    ptr::write_bytes(self.base.as_mut_ptr(), 0u8, keep_resident.byte_count());
 
                     // This is madvise (2)
                     self.restore_original_mapping(
@@ -657,7 +655,7 @@ impl MemoryImageSlot {
             // the rest.
             None => {
                 let size_to_memset = keep_resident.min(self.accessible);
-                ptr::write_bytes(self.base.as_ptr(), 0u8, size_to_memset.byte_count());
+                ptr::write_bytes(self.base.as_mut_ptr(), 0u8, size_to_memset.byte_count());
                 self.restore_original_mapping(
                     size_to_memset,
                     self.accessible
@@ -685,7 +683,10 @@ impl MemoryImageSlot {
             vm::decommit_behavior(),
             DecommitBehavior::RestoreOriginalMapping
         );
-        decommit(self.base.as_ptr().add(base.byte_count()), len.byte_count());
+        decommit(
+            self.base.as_mut_ptr().add(base.byte_count()),
+            len.byte_count(),
+        );
     }
 
     fn set_protection(&self, range: Range<HostAlignedByteCount>, readwrite: bool) -> Result<()> {
@@ -701,7 +702,7 @@ impl MemoryImageSlot {
         // TODO: use Mmap to change memory permissions instead of these free
         // functions.
         unsafe {
-            let start = self.base.as_ptr().add(range.start.byte_count());
+            let start = self.base.as_mut_ptr().add(range.start.byte_count());
             if readwrite {
                 vm::expose_existing_mapping(start, len.byte_count())?;
             } else {
@@ -731,7 +732,7 @@ impl MemoryImageSlot {
         }
 
         unsafe {
-            vm::erase_existing_mapping(self.base.as_ptr(), self.static_size)?;
+            vm::erase_existing_mapping(self.base.as_mut_ptr(), self.static_size)?;
         }
 
         self.image = None;
@@ -852,11 +853,8 @@ mod test {
         // 4 MiB mmap'd area, not accessible
         let mmap = mmap_4mib_inaccessible();
         // Create a MemoryImageSlot on top of it
-        let mut memfd = MemoryImageSlot::create(
-            mmap.as_mut_ptr() as *mut _,
-            HostAlignedByteCount::ZERO,
-            4 << 20,
-        );
+        let mut memfd =
+            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
         memfd.no_clear_on_drop();
         assert!(!memfd.is_dirty());
         // instantiate with 64 KiB initial size
@@ -903,11 +901,8 @@ mod test {
         // 4 MiB mmap'd area, not accessible
         let mmap = mmap_4mib_inaccessible();
         // Create a MemoryImageSlot on top of it
-        let mut memfd = MemoryImageSlot::create(
-            mmap.as_mut_ptr() as *mut _,
-            HostAlignedByteCount::ZERO,
-            4 << 20,
-        );
+        let mut memfd =
+            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
         memfd.no_clear_on_drop();
         // Create an image with some data.
         let image = Arc::new(create_memfd_with_data(page_size, &[1, 2, 3, 4]).unwrap());
@@ -996,11 +991,8 @@ mod test {
             ..Tunables::default_miri()
         };
         let mmap = mmap_4mib_inaccessible();
-        let mut memfd = MemoryImageSlot::create(
-            mmap.as_mut_ptr() as *mut _,
-            HostAlignedByteCount::ZERO,
-            4 << 20,
-        );
+        let mut memfd =
+            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
         memfd.no_clear_on_drop();
 
         // Test basics with the image
@@ -1066,11 +1058,8 @@ mod test {
         };
 
         let mmap = mmap_4mib_inaccessible();
-        let mut memfd = MemoryImageSlot::create(
-            mmap.as_mut_ptr() as *mut _,
-            HostAlignedByteCount::ZERO,
-            4 << 20,
-        );
+        let mut memfd =
+            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
         memfd.no_clear_on_drop();
         let image = Arc::new(create_memfd_with_data(page_size, &[1, 2, 3, 4]).unwrap());
         let initial = 64 << 10;

crates/wasmtime/src/runtime/vm/instance/allocator/pooling/memory_pool.rs

@@ -57,13 +57,12 @@ use super::{
 use crate::prelude::*;
 use crate::runtime::vm::{
     mmap::AlignedLength, CompiledModuleId, InstanceAllocationRequest, InstanceLimits, Memory,
-    MemoryImageSlot, Mmap, MpkEnabled, PoolingInstanceAllocatorConfig,
+    MemoryBase, MemoryImageSlot, Mmap, MmapOffset, MpkEnabled, PoolingInstanceAllocatorConfig,
 };
 use crate::{
     runtime::vm::mpk::{self, ProtectionKey, ProtectionMask},
     vm::HostAlignedByteCount,
 };
-use std::ffi::c_void;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 use wasmtime_environ::{DefinedMemoryIndex, Module, Tunables};
@@ -357,7 +356,7 @@ impl MemoryPool {
                     <= u64::try_from(self.layout.bytes_to_next_stripe_slot().byte_count()).unwrap()
             );
 
-            let base_ptr = self.get_base(allocation_index);
+            let base = self.get_base(allocation_index);
             let base_capacity = self.layout.max_memory_bytes;
 
             let mut slot = self.take_memory_image_slot(allocation_index);
@@ -385,7 +384,7 @@ impl MemoryPool {
             Memory::new_static(
                 ty,
                 tunables,
-                base_ptr,
+                MemoryBase::Mmap(base),
                 base_capacity.byte_count(),
                 slot,
                 unsafe { &mut *request.store.get().unwrap() },
@@ -471,20 +470,15 @@ impl MemoryPool {
         }
     }
 
-    fn get_base(&self, allocation_index: MemoryAllocationIndex) -> *mut u8 {
+    fn get_base(&self, allocation_index: MemoryAllocationIndex) -> MmapOffset {
         assert!(allocation_index.index() < self.layout.num_slots);
         let offset = self
             .layout
             .slot_bytes
             .checked_mul(allocation_index.index())
             .and_then(|c| c.checked_add(self.layout.pre_slab_guard_bytes))
             .expect("slot_bytes * index + pre_slab_guard_bytes overflows");
-        unsafe {
-            self.mapping
-                .as_ptr()
-                .offset(offset.byte_count() as isize)
-                .cast_mut()
-        }
+        self.mapping.offset(offset).expect("offset is in bounds")
     }
 
     /// Take ownership of the given image slot. Must be returned via
@@ -497,7 +491,7 @@ impl MemoryPool {
 
         maybe_slot.unwrap_or_else(|| {
             MemoryImageSlot::create(
-                self.get_base(allocation_index) as *mut c_void,
+                self.get_base(allocation_index),
                 HostAlignedByteCount::ZERO,
                 self.layout.max_memory_bytes.byte_count(),
             )
@@ -822,7 +816,7 @@ mod tests {
 
         for i in 0..5 {
             let index = MemoryAllocationIndex(i);
-            let ptr = pool.get_base(index);
+            let ptr = pool.get_base(index).as_mut_ptr();
             assert_eq!(
                 ptr as usize - base,
                 i as usize * pool.layout.slot_bytes.byte_count()