shrink memory to minimum + refactor

lib/middleware-common/src/opcode_control.rs

@@ -1,14 +1,12 @@
 use wasmer_runtime_core::{
     codegen::{Event, EventSink, FunctionMiddleware, InternalEvent},
-    wasmparser::{Operator, Type as WpType, TypeOrFuncType as WpTypeOrFuncType},
-    vm::InternalField,
     module::ModuleInfo,
+    vm::InternalField,
+    wasmparser::{Operator, Type as WpType, TypeOrFuncType as WpTypeOrFuncType},
+    Instance,
 };
 
-use crate::runtime_breakpoints::{
-    push_runtime_breakpoint,
-    BREAKPOINT_VALUE_MEMORY_LIMIT,
-};
+use crate::runtime_breakpoints::{push_runtime_breakpoint, BREAKPOINT_VALUE_MEMORY_LIMIT};
 
 static FIELD_MEMORY_GROW_COUNT: InternalField = InternalField::allocate();
 
@@ -46,9 +44,7 @@ impl OpcodeControl {
         sink.push(Event::Internal(InternalEvent::GetInternal(
             FIELD_MEMORY_GROW_COUNT.index() as _,
         )));
-        sink.push(Event::WasmOwned(Operator::I64Const{
-            value: 1 as i64,
-        }));
+        sink.push(Event::WasmOwned(Operator::I64Const { value: 1 as i64 }));
         sink.push(Event::WasmOwned(Operator::I64Add));
         sink.push(Event::Internal(InternalEvent::SetInternal(
             FIELD_MEMORY_GROW_COUNT.index() as _,
@@ -118,3 +114,8 @@ impl FunctionMiddleware for OpcodeControl {
         Ok(())
     }
 }
+
+/// Sets the number of points used by an Instance.
+pub fn reset_memory_grow_count(instance: &mut Instance) {
+    instance.set_internal(&FIELD_MEMORY_GROW_COUNT, 0);
+}

lib/runtime-c-api/src/instance.rs

@@ -23,7 +23,7 @@ use wasmer_runtime_core::codegen::{MiddlewareChain, StreamingCompiler};
 #[cfg(not(feature = "cranelift-backend"))]
 use wasmer_middleware_common::metering;
 
-use wasmer_middleware_common::opcode_control;
+use wasmer_middleware_common::opcode_control::{self, reset_memory_grow_count};
 use wasmer_middleware_common::opcode_trace;
 use wasmer_middleware_common::runtime_breakpoints;
 
@@ -274,6 +274,8 @@ pub unsafe extern "C" fn wasmer_instantiate_reset(
         return wasmer_result_t::WASMER_ERROR;
     }
 
+    reset_memory_grow_count(instance);
+
     wasmer_result_t::WASMER_OK
 }
 

lib/runtime-core/src/backing.rs

@@ -110,12 +110,10 @@ impl LocalBacking {
     pub(crate) fn reset(&mut self, module_info: &ModuleInfo) -> RuntimeResult<()> {
         unsafe {
             if RESET_BACKING_MEMORIES {
-                println!("[WASMER] Resetting the memory");
                 Self::reset_memories(&module_info, &mut self.memories)?;
             }
 
             if RESET_BACKING_GLOBALS {
-                println!("[WASMER] Resetting the globals");
                 Self::reset_globals(&module_info, &mut self.globals)?;
             }
         }
@@ -126,10 +124,18 @@ impl LocalBacking {
         module_info: &ModuleInfo,
         memories: &mut SliceMap<LocalMemoryIndex, Memory>,
     ) -> RuntimeResult<()> {
+        Self::shrink_memories(memories)?;
         Self::zero_memories(memories);
         Self::reinitialize_memories(module_info, memories)
     }
 
+    fn shrink_memories(memories: &mut SliceMap<LocalMemoryIndex, Memory>) -> RuntimeResult<()> {
+        for (_index, memory) in memories.iter_mut() {
+            memory.shrink_to_minimum()?;
+        }
+        Ok(())
+    }
+
     fn zero_memories(memories: &mut SliceMap<LocalMemoryIndex, Memory>) {
         for (_index, memory) in memories.iter_mut() {
             let view = &memory.view::<u64>();

lib/runtime-core/src/memory/dynamic.rs

@@ -1,6 +1,5 @@
-use crate::error::GrowError;
 use crate::{
-    error::CreationError,
+    error::{CreationError, GrowError},
     sys,
     types::MemoryDescriptor,
     units::{Bytes, Pages},
@@ -26,6 +25,7 @@ pub struct DynamicMemory {
     memory: sys::Memory,
     current: Pages,
     max: Option<Pages>,
+    min: Pages,
 }
 
 impl DynamicMemory {
@@ -52,6 +52,7 @@ impl DynamicMemory {
             memory,
             current: desc.minimum,
             max: desc.maximum,
+            min: desc.minimum,
         });
         let storage_ptr: *mut DynamicMemory = &mut *storage;
 
@@ -106,6 +107,17 @@ impl DynamicMemory {
         Ok(old_pages)
     }
 
+    /// Shrink the memory to the minimum number of pages.
+    pub fn shrink_to_minimum(&mut self, local: &mut vm::LocalMemory) {
+        let min_size = self.min.bytes().0;
+        _= self.memory.split_at(min_size);
+
+        local.base = self.memory.as_ptr();
+        local.bound = min_size;
+
+        self.current = self.min;
+    }
+
     /// Get this memory represented as a slice of bytes.
     pub fn as_slice(&self) -> &[u8] {
         unsafe { &self.memory.as_slice()[0..self.current.bytes().0] }

lib/runtime-core/src/memory/mod.rs

@@ -1,7 +1,7 @@
 //! The memory module contains the implementation data structures and helper functions used to
 //! manipulate and access wasm memory.
 use crate::{
-    error::{CreationError, CreationResult, GrowError},
+    error::{CreationError, CreationResult, GrowError, RuntimeError, RuntimeResult},
     export::Export,
     import::IsExport,
     memory::dynamic::DYNAMIC_GUARD_SIZE,
@@ -130,6 +130,16 @@ impl Memory {
         }
     }
 
+    /// Shrink this memory to the minimum number of pages.
+    pub(crate) fn shrink_to_minimum(&self) -> RuntimeResult<()> {
+        match &self.variant {
+            MemoryVariant::Unshared(unshared_mem) => unshared_mem.shrink_to_minimum(),
+            MemoryVariant::Shared(_shared_mem) => Err(RuntimeError(Box::new(
+                "shrink_to_minimum is not supported for shared memories",
+            ))),
+        }
+    }
+
     /// Return a "view" of the currently accessible memory. By
     /// default, the view is unsynchronized, using regular memory
     /// accesses. You can force a memory view to use atomic accesses
@@ -313,6 +323,24 @@ impl UnsharedMemory {
         }
     }
 
+    /// Shrink the memory to the minimum number of pages.
+    pub fn shrink_to_minimum(&self) -> RuntimeResult<()> {
+        let mut storage = self.internal.storage.lock().unwrap();
+        let mut local = self.internal.local.get();
+
+        match &mut *storage {
+            UnsharedMemoryStorage::Dynamic(dynamic_memory) => {
+                dynamic_memory.shrink_to_minimum(&mut local)
+            }
+            UnsharedMemoryStorage::Static(_static_memory) => {
+                return Err(RuntimeError(Box::new("Cannot shrink static memory")));
+            }
+        }
+
+        self.internal.local.set(local);
+        Ok(())
+    }
+
     pub(crate) fn vm_local_memory(&self) -> *mut vm::LocalMemory {
         self.internal.local.as_ptr()
     }

lib/runtime-core/src/sys/unix/memory.rs

@@ -4,9 +4,9 @@ use crate::sys::{round_down_to_page_size, round_up_to_page_size};
 use errno;
 use nix::libc;
 use page_size;
+use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
 use std::ops::{Bound, RangeBounds};
 use std::{fs::File, os::unix::io::IntoRawFd, path::Path, ptr, slice, sync::Arc};
-use rkyv::{Archive, Serialize as RkyvSerialize, Deserialize as RkyvDeserialize};
 
 unsafe impl Send for Memory {}
 unsafe impl Sync for Memory {}
@@ -100,7 +100,10 @@ impl Memory {
 
     /// Create a new memory with the given contents size and protection.
     /// Used when the size of the contents must be tracked (e.g. for rkyv deserialization).
-    pub fn with_content_size_protect(content_size: u32, protection: Protect) -> Result<Self, String> {
+    pub fn with_content_size_protect(
+        content_size: u32,
+        protection: Protect,
+    ) -> Result<Self, String> {
         let mut memory = Self::with_size_protect(content_size as usize, protection)?;
         memory.set_content_size(content_size);
         Ok(memory)
@@ -194,7 +197,7 @@ impl Memory {
     }
 
     /// Split this memory into multiple memories by the given offset.
-    pub fn split_at(mut self, offset: usize) -> (Memory, Memory) {
+    pub fn split_at(&mut self, offset: usize) -> (&mut Memory, Memory) {
         let page_size = page_size::get();
         if offset % page_size == 0 {
             let second_ptr = unsafe { self.ptr.add(offset) };
@@ -283,7 +286,18 @@ impl Clone for Memory {
 }
 
 /// Kinds of memory protection.
-#[derive(Serialize, Deserialize, Debug, Copy, Clone, PartialEq, Eq, Archive, RkyvSerialize, RkyvDeserialize)]
+#[derive(
+    Serialize,
+    Deserialize,
+    Debug,
+    Copy,
+    Clone,
+    PartialEq,
+    Eq,
+    Archive,
+    RkyvSerialize,
+    RkyvDeserialize,
+)]
 #[allow(dead_code)]
 #[archive(compare(PartialEq))]
 #[archive_attr(derive(Debug))]