feat(runtime): Memory swapping

crates/wasmtime/src/runtime/memory.rs

@@ -680,6 +680,72 @@ impl Memory {
         }
     }
 
+    /// Hot-swaps the backing storage of this memory with `other`'s, in O(1) and
+    /// without copying either memory's contents.
+    ///
+    /// After this returns, reads and writes through `self` observe what `other`
+    /// previously held and vice versa, and both memories' `VMContext`s are
+    /// updated so that running wasm sees the swap immediately. The two memories
+    /// exchange their entire backing allocations; their lengths and capacities
+    /// come along with them.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the two memories have different [types](Memory::ty) or
+    /// different byte capacities, or if they are defined in the same instance
+    /// (swapping operates across instances). Swapping a memory with itself is a
+    /// no-op and returns `Ok`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if either memory does not belong to `store`.
+    pub fn swap(&self, mut store: impl AsContextMut, other: &Memory) -> Result<()> {
+        let store = store.as_context_mut().0;
+        assert!(
+            self.comes_from_same_store(store),
+            "memory used with the wrong store"
+        );
+        assert!(
+            other.comes_from_same_store(store),
+            "memory used with the wrong store"
+        );
+
+        let a = self.instance.instance();
+        let b = other.instance.instance();
+        if a == b {
+            // Same instance: only the no-op self-swap is supported; swapping two
+            // memories within one instance isn't (it has no use case and would
+            // need a different disjoint-borrow path).
+            if self.index == other.index {
+                return Ok(());
+            }
+            bail!("cannot swap two memories defined in the same instance");
+        }
+
+        // Types (page size, limits, shared bit, index type) must match so the
+        // swapped-in memory is interchangeable.
+        if self.wasmtime_ty(store) != other.wasmtime_ty(store) {
+            bail!("cannot swap memories with different types");
+        }
+
+        // Capacities (reservations) must match so the swapped-in base satisfies
+        // the bounds checks baked into compiled code. Equal types under one
+        // engine already imply this, but check explicitly to keep the invariant
+        // local to the swap.
+        let a_cap = store[self.instance]
+            .get_defined_memory(self.index)
+            .byte_capacity();
+        let b_cap = store[other.instance]
+            .get_defined_memory(other.index)
+            .byte_capacity();
+        if a_cap != b_cap {
+            bail!("cannot swap memories with different byte capacities");
+        }
+
+        store.swap_defined_memories((a, self.index), (b, other.index));
+        Ok(())
+    }
+
     /// Creates a new memory from its raw component parts.
     ///
     /// # Safety
@@ -1156,4 +1222,78 @@ mod tests {
 
         Ok(())
     }
+
+    // Two host memories of the same type exchange their contents in place.
+    #[test]
+    fn swap_exchanges_contents() -> Result<()> {
+        let mut store = Store::<()>::default();
+        let ty = MemoryType::new(1, None);
+        let a = Memory::new(&mut store, ty.clone())?;
+        let b = Memory::new(&mut store, ty)?;
+
+        a.data_mut(&mut store)[0] = 0xAA;
+        b.data_mut(&mut store)[0] = 0xBB;
+
+        a.swap(&mut store, &b)?;
+        assert_eq!(a.data(&store)[0], 0xBB);
+        assert_eq!(b.data(&store)[0], 0xAA);
+
+        a.swap(&mut store, &b)?;
+        assert_eq!(a.data(&store)[0], 0xAA);
+        assert_eq!(b.data(&store)[0], 0xBB);
+
+        Ok(())
+    }
+
+    // After swap, compiled wasm observes the new backing.
+    // i.e. the swap rewrites the in-`VMContext` `VMMemoryDefinition`.
+    #[test]
+    fn swap_is_visible_to_running_wasm() -> Result<()> {
+        let mut store = Store::<()>::default();
+        let module = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (memory (export "m") 1 1)
+                    (func (export "load8") (param i32) (result i32)
+                        local.get 0
+                        i32.load8_u))
+            "#,
+        )?;
+        let instance = Instance::new(&mut store, &module, &[])?;
+        let m = instance.get_memory(&mut store, "m").unwrap();
+        let load8 = instance.get_typed_func::<i32, i32>(&mut store, "load8")?;
+
+        // A detached buffer of the matching type holding different contents.
+        let buf_ty = m.ty(&store);
+        let buf = Memory::new(&mut store, buf_ty)?;
+        m.data_mut(&mut store)[0] = 11;
+        buf.data_mut(&mut store)[0] = 99;
+
+        assert_eq!(load8.call(&mut store, 0)?, 11);
+
+        m.swap(&mut store, &buf)?;
+
+        // Compiled code now reads the swapped-in backing.
+        assert_eq!(load8.call(&mut store, 0)?, 99);
+        assert_eq!(m.data(&store)[0], 99);
+        assert_eq!(buf.data(&store)[0], 11);
+
+        Ok(())
+    }
+
+    // Swapping a memory with itself is a no-op; mismatched types are rejected.
+    #[test]
+    fn swap_self_noop_and_type_mismatch() -> Result<()> {
+        let mut store = Store::<()>::default();
+        let a = Memory::new(&mut store, MemoryType::new(1, None))?;
+        a.data_mut(&mut store)[0] = 7;
+        a.swap(&mut store, &a)?;
+        assert_eq!(a.data(&store)[0], 7);
+
+        let big = Memory::new(&mut store, MemoryType::new(2, None))?;
+        assert!(a.swap(&mut store, &big).is_err());
+
+        Ok(())
+    }
 }

crates/wasmtime/src/runtime/store.rs

@@ -110,7 +110,9 @@ use core::pin::Pin;
 use core::ptr::NonNull;
 #[cfg(any(feature = "async", feature = "gc"))]
 use core::task::Poll;
-use wasmtime_environ::{DefinedGlobalIndex, DefinedTableIndex, EntityRef, TripleExt};
+use wasmtime_environ::{
+    DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, EntityRef, TripleExt,
+};
 
 mod context;
 pub use self::context::*;
@@ -1660,6 +1662,43 @@ impl StoreOpaque {
         self.instances[id].handle.get_mut()
     }
 
+    /// Hot-swap the backing storage of two defined memories living in two
+    /// different instances within this store without copying contents.
+    /// The `(allocation, Memory)` entries are swapped wholesale (so the
+    /// instance allocator's bookkeeping travels with each memory, keeping both
+    /// the on-demand and pooling allocators correct), and each instance's
+    /// `VMContext` is refreshed so compiled code observes the swap.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `a.0 == b.0` (the two memories must be in different instances);
+    /// callers (`Memory::swap`) reject that case with an error beforehand. The
+    /// caller must also have validated that the two memories have matching types
+    /// and byte capacities.
+    pub(crate) fn swap_defined_memories(
+        &mut self,
+        a: (InstanceId, DefinedMemoryIndex),
+        b: (InstanceId, DefinedMemoryIndex),
+    ) {
+        assert_ne!(
+            a.0, b.0,
+            "swap_defined_memories requires different instances"
+        );
+        // SAFETY: `a.0 != b.0` so the two instance handles are disjoint, and we
+        // only touch each instance's own defined memory (never traversing
+        // laterally between instances), satisfying the contract of
+        // `optional_gc_store_and_instances_mut`.
+        unsafe {
+            let (_gc, [mut ia, mut ib]) = self.optional_gc_store_and_instances_mut([a.0, b.0]);
+            core::mem::swap(
+                ia.as_mut().defined_memory_entry_mut(a.1),
+                ib.as_mut().defined_memory_entry_mut(b.1),
+            );
+            ia.as_mut().refresh_defined_memory(a.1);
+            ib.as_mut().refresh_defined_memory(b.1);
+        }
+    }
+
     /// Accessor from `InstanceId` to both `Pin<&mut vm::Instance>`
     /// and `&ModuleRegistry`.
     #[inline]

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

@@ -1034,6 +1034,27 @@ impl Instance {
         &self.memories[index].1
     }
 
+    /// Mutable access to the `index`th defined memory's `(allocation, memory)`
+    /// entry. Used by the store to hot-swap a memory's backing storage between
+    /// two instances; after swapping the entry, call
+    /// [`Instance::refresh_defined_memory`] to update the `VMContext`.
+    pub(crate) fn defined_memory_entry_mut(
+        self: Pin<&mut Self>,
+        index: DefinedMemoryIndex,
+    ) -> &mut (MemoryAllocationIndex, Memory) {
+        &mut self.memories_mut()[index]
+    }
+
+    /// Recompute the `index`th defined memory's `VMMemoryDefinition` (base +
+    /// length) from its current allocation and write it into the `VMContext`, so
+    /// compiled code observes the memory's current backing. This is the same
+    /// fix-up [`Instance::memory_grow`] performs after a grow; it must be called
+    /// after the entry is swapped via [`Instance::defined_memory_entry_mut`].
+    pub(crate) fn refresh_defined_memory(mut self: Pin<&mut Self>, index: DefinedMemoryIndex) {
+        let vmmemory = self.as_mut().get_defined_memory_mut(index).vmmemory();
+        self.set_memory(index, vmmemory);
+    }
+
     pub fn get_defined_memory_vmimport(&self, index: DefinedMemoryIndex) -> VMMemoryImport {
         crate::runtime::vm::VMMemoryImport {
             from: self.memory_ptr(index).into(),

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

@@ -385,6 +385,18 @@ impl Memory {
         }
     }
 
+    /// Returns the number of bytes this memory's current allocation can address
+    /// without relocating its base pointer (i.e. its reservation). Two memories
+    /// can only be hot-swapped (see the embedder `Memory::swap`) if their
+    /// capacities match.
+    pub fn byte_capacity(&self) -> usize {
+        match self {
+            Memory::Local(mem) => mem.byte_capacity(),
+            // Shared memories are never hot-swapped; report the logical size.
+            Memory::Shared(mem) => mem.byte_size(),
+        }
+    }
+
     /// Returns whether or not this memory needs initialization. It
     /// may not if it already has initial content thanks to a CoW
     /// mechanism.
@@ -733,6 +745,10 @@ impl LocalMemory {
         self.alloc.byte_size()
     }
 
+    pub fn byte_capacity(&self) -> usize {
+        self.alloc.byte_capacity()
+    }
+
     pub fn needs_init(&self) -> bool {
         match &self.memory_image {
             Some(image) => !image.has_image(),