Merge pull request #5617 from wasmerio/feat/memory-size-fn

Implement Memory::size function

Author: Arshia001

lib/api/src/backend/js/entities/memory/js.rs

@@ -8,6 +8,12 @@ extern "C" {
     #[derive(Clone, Debug, PartialEq, Eq)]
     pub type JSMemory;
 
+    /// The `buffer` property of the `Memory` object, which is an `ArrayBuffer` or `SharedArrayBuffer`.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/WebAssembly/Reference/JavaScript_interface/Memory/buffer)
+    #[wasm_bindgen(method, getter, js_name = buffer, js_namespace = WebAssembly)]
+    pub fn buffer(this: &JSMemory) -> JsValue;
+
     /// The `grow()` protoype method of the `Memory` object increases the
     /// size of the memory instance by a specified number of WebAssembly
     /// pages.

lib/api/src/backend/js/entities/memory/mod.rs

@@ -8,7 +8,7 @@ pub(crate) mod buffer;
 pub(crate) use buffer::*;
 
 use wasm_bindgen::JsCast;
-use wasmer_types::{MemoryError, MemoryType, Pages};
+use wasmer_types::{MemoryError, MemoryType, Pages, WASM_PAGE_SIZE};
 
 use crate::{
     js::vm::memory::VMMemory,
@@ -88,6 +88,24 @@ impl Memory {
         MemoryView::new(self, store)
     }
 
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        let js_memory = &self.handle.memory;
+        let our_js_memory: &JSMemory = JsCast::unchecked_from_js_ref(js_memory);
+        let buffer = our_js_memory.buffer();
+
+        let byte_length = match buffer.dyn_into::<js_sys::ArrayBuffer>() {
+            Ok(array_buffer) => array_buffer.byte_length(),
+            Err(buffer) => match buffer.dyn_into::<js_sys::SharedArrayBuffer>() {
+                Ok(array_buffer) => array_buffer.byte_length(),
+                Err(_) => {
+                    unreachable!("Memory.buffer should be an array buffer or a shared array buffer")
+                }
+            },
+        };
+
+        Pages(byte_length.div_ceil(WASM_PAGE_SIZE as u32))
+    }
+
     pub fn grow<IntoPages>(
         &self,
         store: &mut impl AsStoreMut,

lib/api/src/backend/jsc/entities/memory/mod.rs

@@ -88,6 +88,23 @@ impl Memory {
         MemoryView::new(self, store)
     }
 
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        let store_ref = store.as_store_ref();
+        let engine = store_ref.engine();
+        let context = engine.as_jsc().context();
+
+        let js_memory = &self.handle.memory;
+        let buffer = js_memory
+            .get_property(context, "buffer")
+            .to_object(context)
+            .unwrap();
+        let byte_length = buffer
+            .get_property(context, "byteLength")
+            .to_number(context)
+            .unwrap() as u32;
+        Pages(byte_length.div_ceil(wasmer_types::WASM_PAGE_SIZE as u32))
+    }
+
     pub fn grow<IntoPages>(
         &self,
         store: &mut impl AsStoreMut,

lib/api/src/backend/sys/entities/memory/mod.rs

@@ -53,6 +53,12 @@ impl Memory {
             .ty()
     }
 
+    pub(crate) fn size(&self, store: &impl AsStoreRef) -> Pages {
+        self.handle
+            .get(store.as_store_ref().objects().as_sys())
+            .size()
+    }
+
     pub(crate) fn grow<IntoPages>(
         &self,
         store: &mut impl AsStoreMut,

lib/api/src/backend/v8/entities/memory/mod.rs

@@ -88,6 +88,12 @@ impl Memory {
         }
     }
 
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        check_isolate(store);
+        let size = unsafe { wasm_memory_size(self.handle) };
+        Pages(size)
+    }
+
     pub fn view<'a>(&self, store: &'a impl AsStoreRef) -> MemoryView<'a> {
         check_isolate(store);
         let store_ref = store.as_store_ref();

lib/api/src/backend/wamr/entities/memory/mod.rs

@@ -79,6 +79,11 @@ impl Memory {
         }
     }
 
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        let size = unsafe { wasm_memory_size(self.handle) };
+        Pages(size)
+    }
+
     pub fn view<'a>(&self, store: &'a impl AsStoreRef) -> MemoryView<'a> {
         MemoryView::new(self, store)
     }

lib/api/src/backend/wasmi/entities/memory/mod.rs

@@ -78,6 +78,11 @@ impl Memory {
         }
     }
 
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        let size = unsafe { wasm_memory_size(self.handle) };
+        Pages(size)
+    }
+
     pub fn view<'a>(&self, store: &'a impl AsStoreRef) -> MemoryView<'a> {
         MemoryView::new(self, store)
     }

lib/api/src/entities/memory/inner.rs

@@ -125,6 +125,13 @@ impl BackendMemory {
         })
     }
 
+    /// Retrieve the size of the memory in pages.
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        match_rt!(on self => s {
+            s.size(store)
+        })
+    }
+
     /// Grow memory by the specified amount of WebAssembly [`Pages`] and return
     /// the previous memory size.
     ///

lib/api/src/entities/memory/mod.rs

@@ -79,6 +79,11 @@ impl Memory {
         MemoryView::new(self, store)
     }
 
+    /// Retrieve the size of the memory in pages.
+    pub fn size(&self, store: &impl AsStoreRef) -> Pages {
+        self.0.size(store)
+    }
+
     /// Grow memory by the specified amount of WebAssembly [`Pages`] and return
     /// the previous memory size.
     ///

lib/api/tests/memory.rs

@@ -116,3 +116,29 @@ fn test_wasm_slice_issue_5444() {
         Some(MemoryAccessError::UnalignedPointerRead)
     ))
 }
+
+#[cfg_attr(
+    feature = "wamr",
+    ignore = "wamr reports memory size incorrectly in this scenario"
+)]
+#[test]
+fn test_wasm_memory_size() {
+    let mut store = Store::default();
+
+    // Test once with not-shared memory...
+    {
+        let memory = Memory::new(&mut store, MemoryType::new(10, Some(65536), false)).unwrap();
+        assert_eq!(memory.size(&store).0, 10);
+        memory.grow(&mut store, wasmer::Pages(1)).unwrap();
+        assert_eq!(memory.size(&store).0, 11);
+    }
+
+    // ... and once with shared memory, since JS and JSC (at least) have different
+    // representations for it
+    {
+        let memory = Memory::new(&mut store, MemoryType::new(10, Some(65536), true)).unwrap();
+        assert_eq!(memory.size(&store).0, 10);
+        memory.grow(&mut store, wasmer::Pages(1)).unwrap();
+        assert_eq!(memory.size(&store).0, 11);
+    }
+}