feat: Custom memory page size (#22)

danielstuart14 · web-flow · commit 3c16adc4e128 · 2024-09-10T00:19:25.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/README.md b/README.md
@@ -27,17 +27,18 @@ TinyWasm passes all WebAssembly MVP tests from the [WebAssembly core testsuite](
 🚧 -- in development / partialy supported\
 🟢 -- fully supported
 
-| Proposal                                                                                                                   | Status | TinyWasm Version |
-| -------------------------------------------------------------------------------------------------------------------------- | ------ | ---------------- |
-| [**Mutable Globals**](https://github.com/WebAssembly/mutable-global/blob/master/proposals/mutable-global/Overview.md)      | 🟢     | 0.2.0            |
-| [**Non-trapping float-to-int Conversion**](https://github.com/WebAssembly/nontrapping-float-to-int-conversions)            | 🟢     | 0.2.0            |
-| [**Sign-extension operators**](https://github.com/WebAssembly/sign-extension-ops)                                          | 🟢     | 0.2.0            |
-| [**Multi-value**](https://github.com/WebAssembly/spec/blob/master/proposals/multi-value/Overview.md)                       | 🟢     | 0.2.0            |
-| [**Bulk Memory Operations**](https://github.com/WebAssembly/spec/blob/master/proposals/bulk-memory-operations/Overview.md) | 🟢     | 0.4.0            |
-| [**Reference Types**](https://github.com/WebAssembly/reference-types/blob/master/proposals/reference-types/Overview.md)    | 🟢     | 0.7.0            |
-| [**Multiple Memories**](https://github.com/WebAssembly/multi-memory/blob/master/proposals/multi-memory/Overview.md)        | 🟢     | 0.8.0            |
-| [**Memory64**](https://github.com/WebAssembly/memory64/blob/master/proposals/memory64/Overview.md)                         | 🚧     | N/A              |
-| [**Fixed-Width SIMD**](https://github.com/webassembly/simd)                                                                | 🌑     | N/A              |
+| Proposal                                                                                                                    | Status | TinyWasm Version |
+| --------------------------------------------------------------------------------------------------------------------------- | ------ | ---------------- |
+| [**Mutable Globals**](https://github.com/WebAssembly/mutable-global/blob/master/proposals/mutable-global/Overview.md)       | 🟢     | 0.2.0            |
+| [**Non-trapping float-to-int Conversion**](https://github.com/WebAssembly/nontrapping-float-to-int-conversions)             | 🟢     | 0.2.0            |
+| [**Sign-extension operators**](https://github.com/WebAssembly/sign-extension-ops)                                           | 🟢     | 0.2.0            |
+| [**Multi-value**](https://github.com/WebAssembly/spec/blob/master/proposals/multi-value/Overview.md)                        | 🟢     | 0.2.0            |
+| [**Bulk Memory Operations**](https://github.com/WebAssembly/spec/blob/master/proposals/bulk-memory-operations/Overview.md)  | 🟢     | 0.4.0            |
+| [**Reference Types**](https://github.com/WebAssembly/reference-types/blob/master/proposals/reference-types/Overview.md)     | 🟢     | 0.7.0            |
+| [**Multiple Memories**](https://github.com/WebAssembly/multi-memory/blob/master/proposals/multi-memory/Overview.md)         | 🟢     | 0.8.0            |
+| [**Custom Page Sizes**](https://github.com/WebAssembly/custom-page-sizes/blob/main/proposals/custom-page-sizes/Overview.md) | 🟢     | `next`           |
+| [**Memory64**](https://github.com/WebAssembly/memory64/blob/master/proposals/memory64/Overview.md)                          | 🚧     | N/A              |
+| [**Fixed-Width SIMD**](https://github.com/webassembly/simd)                                                                 | 🌑     | N/A              |
 
 ## Usage
 
diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
@@ -105,11 +105,12 @@ pub(crate) fn convert_module_memories<T: IntoIterator<Item = wasmparser::Result<
 }
 
 pub(crate) fn convert_module_memory(memory: wasmparser::MemoryType) -> MemoryType {
-    MemoryType {
-        arch: if memory.memory64 { MemoryArch::I64 } else { MemoryArch::I32 },
-        page_count_initial: memory.initial,
-        page_count_max: memory.maximum,
-    }
+    MemoryType::new(
+        if memory.memory64 { MemoryArch::I64 } else { MemoryArch::I32 },
+        memory.initial,
+        memory.maximum,
+        memory.page_size_log2.map(|x| 1 << x),
+    )
 }
 
 pub(crate) fn convert_module_tables<'a, T: IntoIterator<Item = wasmparser::Result<wasmparser::Table<'a>>>>(
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
@@ -63,6 +63,7 @@ impl Parser {
             multi_memory: true,
             memory64: false,
             simd: true,
+            custom_page_sizes: true,
 
             gc_types: true,
             component_model: false,
@@ -75,7 +76,6 @@ impl Parser {
             memory_control: false,
             relaxed_simd: false,
             threads: false,
-            custom_page_sizes: false,
             shared_everything_threads: false,
             component_model_multiple_returns: false,
             legacy_exceptions: false,
diff --git a/crates/tinywasm/Cargo.toml b/crates/tinywasm/Cargo.toml
@@ -73,6 +73,12 @@ name="test-wasm-simd"
 harness=false
 test=false
 
+
+[[test]]
+name="test-wasm-custom-page-sizes"
+harness=false
+test=false
+
 [[test]]
 name="test-wast"
 harness=false
diff --git a/crates/tinywasm/benches/argon2id.rs b/crates/tinywasm/benches/argon2id.rs
@@ -1,6 +1,6 @@
 use criterion::{criterion_group, criterion_main, Criterion};
 use eyre::Result;
-use tinywasm::{ModuleInstance, Store, types};
+use tinywasm::{types, ModuleInstance, Store};
 use types::{archive::AlignedVec, TinyWasmModule};
 
 const WASM: &[u8] = include_bytes!("../../../examples/rust/out/argon2id.opt.wasm");
diff --git a/crates/tinywasm/benches/fibonacci.rs b/crates/tinywasm/benches/fibonacci.rs
@@ -1,6 +1,6 @@
 use criterion::{criterion_group, criterion_main, Criterion};
 use eyre::Result;
-use tinywasm::{ModuleInstance, Store, types};
+use tinywasm::{types, ModuleInstance, Store};
 use types::{archive::AlignedVec, TinyWasmModule};
 
 const WASM: &[u8] = include_bytes!("../../../examples/rust/out/fibonacci.opt.wasm");
diff --git a/crates/tinywasm/benches/tinywasm.rs b/crates/tinywasm/benches/tinywasm.rs
@@ -1,6 +1,6 @@
 use criterion::{criterion_group, criterion_main, Criterion};
 use eyre::Result;
-use tinywasm::{Extern, FuncContext, Imports, ModuleInstance, Store, types};
+use tinywasm::{types, Extern, FuncContext, Imports, ModuleInstance, Store};
 use types::{archive::AlignedVec, TinyWasmModule};
 
 const WASM: &[u8] = include_bytes!("../../../examples/rust/out/tinywasm.opt.wasm");
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
@@ -208,7 +208,7 @@ impl From<&Import> for ExternName {
 /// imports
 ///     .define("my_module", "print_i32", print_i32)?
 ///     .define("my_module", "table", Extern::table(table_type, table_init))?
-///     .define("my_module", "memory", Extern::memory(MemoryType::new_32(1, Some(2))))?
+///     .define("my_module", "memory", Extern::memory(MemoryType::new_32(1, Some(2), None)))?
 ///     .define("my_module", "global_i32", Extern::global(WasmValue::I32(666), false))?
 ///     .link_module("my_other_module", 0)?;
 /// # Ok(())
@@ -317,22 +317,20 @@ impl Imports {
         actual: &MemoryType,
         real_size: Option<usize>,
     ) -> Result<()> {
-        Self::compare_types(import, &expected.arch, &actual.arch)?;
+        Self::compare_types(import, &expected.arch(), &actual.arch())?;
 
-        if actual.page_count_initial > expected.page_count_initial
-            && real_size.map_or(true, |size| actual.page_count_initial > size as u64)
+        if actual.page_count_initial() > expected.page_count_initial()
+            && real_size.map_or(true, |size| actual.page_count_initial() > size as u64)
         {
             return Err(LinkingError::incompatible_import_type(import).into());
         }
 
-        if expected.page_count_max.is_none() && actual.page_count_max.is_some() {
+        if expected.page_size() != actual.page_size() {
             return Err(LinkingError::incompatible_import_type(import).into());
         }
 
-        if let (Some(expected_max), Some(actual_max)) = (expected.page_count_max, actual.page_count_max) {
-            if actual_max < expected_max {
-                return Err(LinkingError::incompatible_import_type(import).into());
-            }
+        if expected.page_count_max() > actual.page_count_max() {
+            return Err(LinkingError::incompatible_import_type(import).into());
         }
 
         Ok(())
diff --git a/crates/tinywasm/src/store/memory.rs b/crates/tinywasm/src/store/memory.rs
@@ -4,10 +4,6 @@ use tinywasm_types::{MemoryType, ModuleInstanceAddr};
 
 use crate::{cold, log, Error, Result};
 
-const PAGE_SIZE: usize = 65536;
-const MAX_PAGES: usize = 65536;
-const MAX_SIZE: u64 = PAGE_SIZE as u64 * MAX_PAGES as u64;
-
 /// A WebAssembly Memory Instance
 ///
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#memory-instances>
@@ -21,13 +17,13 @@ pub(crate) struct MemoryInstance {
 
 impl MemoryInstance {
     pub(crate) fn new(kind: MemoryType, owner: ModuleInstanceAddr) -> Self {
-        assert!(kind.page_count_initial <= kind.page_count_max.unwrap_or(MAX_PAGES as u64));
-        log::debug!("initializing memory with {} pages", kind.page_count_initial);
+        assert!(kind.page_count_initial() <= kind.page_count_max());
+        log::debug!("initializing memory with {} pages of {} bytes", kind.page_count_initial(), kind.page_size());
 
         Self {
             kind,
-            data: vec![0; PAGE_SIZE * kind.page_count_initial as usize],
-            page_count: kind.page_count_initial as usize,
+            data: vec![0; kind.initial_size() as usize],
+            page_count: kind.page_count_initial() as usize,
             _owner: owner,
         }
     }
@@ -58,7 +54,7 @@ impl MemoryInstance {
     }
 
     pub(crate) fn max_pages(&self) -> usize {
-        self.kind.page_count_max.unwrap_or(MAX_PAGES as u64) as usize
+        self.kind.page_count_max() as usize
     }
 
     pub(crate) fn load(&self, addr: usize, len: usize) -> Result<&[u8]> {
@@ -133,12 +129,12 @@ impl MemoryInstance {
         let new_pages = current_pages as i64 + pages_delta as i64;
         debug_assert!(new_pages <= i32::MAX as i64, "page count should never be greater than i32::MAX");
 
-        if new_pages < 0 || new_pages > MAX_PAGES as i64 || new_pages as usize > self.max_pages() {
+        if new_pages < 0 || new_pages as usize > self.max_pages() {
             return None;
         }
 
-        let new_size = new_pages as usize * PAGE_SIZE;
-        if new_size as u64 > MAX_SIZE {
+        let new_size = (new_pages as u64 * self.kind.page_size()) as usize;
+        if new_size as u64 > self.kind.max_size() {
             return None;
         }
 
@@ -190,7 +186,7 @@ mod memory_instance_tests {
     use tinywasm_types::MemoryArch;
 
     fn create_test_memory() -> MemoryInstance {
-        let kind = MemoryType { arch: MemoryArch::I32, page_count_initial: 1, page_count_max: Some(2) };
+        let kind = MemoryType::new(MemoryArch::I32, 1, Some(2), None);
         let owner = ModuleInstanceAddr::default();
         MemoryInstance::new(kind, owner)
     }
@@ -249,7 +245,7 @@ mod memory_instance_tests {
     #[test]
     fn test_memory_grow_out_of_bounds() {
         let mut memory = create_test_memory();
-        assert!(memory.grow(MAX_PAGES as i32 + 1).is_none());
+        assert!(memory.grow(memory.kind.max_size() as i32 + 1).is_none());
     }
 
     #[test]
@@ -258,4 +254,29 @@ mod memory_instance_tests {
         assert_eq!(memory.grow(1), Some(1));
         assert_eq!(memory.grow(1), None);
     }
+
+    #[test]
+    fn test_memory_custom_page_size_out_of_bounds() {
+        let kind = MemoryType::new(MemoryArch::I32, 1, Some(2), Some(1));
+        let owner = ModuleInstanceAddr::default();
+        let mut memory = MemoryInstance::new(kind, owner);
+
+        let data_to_store = [1, 2];
+        assert!(memory.store(0, data_to_store.len(), &data_to_store).is_err());
+    }
+
+    #[test]
+    fn test_memory_custom_page_size_grow() {
+        let kind = MemoryType::new(MemoryArch::I32, 1, Some(2), Some(1));
+        let owner = ModuleInstanceAddr::default();
+        let mut memory = MemoryInstance::new(kind, owner);
+
+        assert_eq!(memory.grow(1), Some(1));
+
+        let data_to_store = [1, 2];
+        assert!(memory.store(0, data_to_store.len(), &data_to_store).is_ok());
+
+        let loaded_data = memory.load(0, data_to_store.len()).unwrap();
+        assert_eq!(loaded_data, &data_to_store);
+    }
 }
diff --git a/crates/tinywasm/src/store/mod.rs b/crates/tinywasm/src/store/mod.rs
@@ -243,7 +243,7 @@ impl Store {
         let mem_count = self.data.memories.len();
         let mut mem_addrs = Vec::with_capacity(mem_count);
         for (i, mem) in memories.into_iter().enumerate() {
-            if let MemoryArch::I64 = mem.arch {
+            if let MemoryArch::I64 = mem.arch() {
                 return Err(Error::UnsupportedFeature("64-bit memories".to_string()));
             }
             self.data.memories.push(MemoryInstance::new(mem, idx));
@@ -405,7 +405,7 @@ impl Store {
     }
 
     pub(crate) fn add_mem(&mut self, mem: MemoryType, idx: ModuleInstanceAddr) -> Result<MemAddr> {
-        if let MemoryArch::I64 = mem.arch {
+        if let MemoryArch::I64 = mem.arch() {
             return Err(Error::UnsupportedFeature("64-bit memories".to_string()));
         }
         self.data.memories.push(MemoryInstance::new(mem, idx));
diff --git a/crates/tinywasm/tests/generated/wasm-custom-page-sizes.csv b/crates/tinywasm/tests/generated/wasm-custom-page-sizes.csv
@@ -0,0 +1 @@
+0.8.0,57,0,[{"name":"custom-page-sizes-invalid.wast","passed":22,"failed":0},{"name":"custom-page-sizes.wast","passed":35,"failed":0}]
diff --git a/crates/tinywasm/tests/test-wasm-custom-page-sizes.rs b/crates/tinywasm/tests/test-wasm-custom-page-sizes.rs
@@ -0,0 +1,20 @@
+mod testsuite;
+use eyre::{eyre, Result};
+use owo_colors::OwoColorize;
+use testsuite::TestSuite;
+
+fn main() -> Result<()> {
+    let mut test_suite = TestSuite::new();
+
+    TestSuite::set_log_level(log::LevelFilter::Off);
+    test_suite.run_spec_group(wasm_testsuite::get_proposal_tests("custom-page-sizes"))?;
+    test_suite.save_csv("./tests/generated/wasm-custom-page-sizes.csv", env!("CARGO_PKG_VERSION"))?;
+
+    if test_suite.failed() {
+        println!();
+        Err(eyre!(format!("{}:\n{:#?}", "failed one or more tests".red().bold(), test_suite,)))
+    } else {
+        println!("\n\npassed all tests:\n{test_suite:#?}");
+        Ok(())
+    }
+}
diff --git a/crates/tinywasm/tests/testsuite/run.rs b/crates/tinywasm/tests/testsuite/run.rs
@@ -121,7 +121,11 @@ impl TestSuite {
         });
 
         imports
-            .define("spectest", "memory", Extern::memory(MemoryType::new_32(1, Some(2))))?
+            .define(
+                "spectest",
+                "memory",
+                Extern::memory(MemoryType::new(tinywasm_types::MemoryArch::I32, 1, Some(2), None)),
+            )?
             .define("spectest", "table", table)?
             .define("spectest", "global_i32", Extern::global(WasmValue::I32(666), false))?
             .define("spectest", "global_i64", Extern::global(WasmValue::I64(666), false))?
@@ -145,6 +149,7 @@ impl TestSuite {
 
     pub fn run_spec_group<T: AsRef<str>>(&mut self, tests: impl IntoIterator<Item = T>) -> Result<()> {
         tests.into_iter().for_each(|group| {
+            println!("running group: {}", group.as_ref());
             let group = group.as_ref();
             let group_wast = wasm_testsuite::get_test_wast(group).expect("failed to get test wast");
             if self.1.contains(&(*group).to_string()) {
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
@@ -12,6 +12,14 @@ extern crate alloc;
 use alloc::boxed::Box;
 use core::{fmt::Debug, ops::Range};
 
+// Memory defaults
+const MEM_PAGE_SIZE: u64 = 65536;
+const MAX_MEMORY_SIZE: u64 = 4294967296;
+
+const fn max_page_count(page_size: u64) -> u64 {
+    MAX_MEMORY_SIZE / page_size
+}
+
 // log for logging (optional).
 #[cfg(feature = "logging")]
 #[allow(clippy::single_component_path_imports, unused_imports)]
@@ -252,14 +260,39 @@ impl TableType {
 #[derive(Debug, Copy, Clone, PartialEq)]
 #[cfg_attr(feature = "archive", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize), archive(check_bytes))]
 pub struct MemoryType {
-    pub arch: MemoryArch,
-    pub page_count_initial: u64,
-    pub page_count_max: Option<u64>,
+    arch: MemoryArch,
+    page_count_initial: u64,
+    page_count_max: Option<u64>,
+    page_size: Option<u64>,
 }
 
 impl MemoryType {
-    pub fn new_32(page_count_initial: u64, page_count_max: Option<u64>) -> Self {
-        Self { arch: MemoryArch::I32, page_count_initial, page_count_max }
+    pub fn new(arch: MemoryArch, page_count_initial: u64, page_count_max: Option<u64>, page_size: Option<u64>) -> Self {
+        Self { arch, page_count_initial, page_count_max, page_size }
+    }
+
+    pub fn arch(&self) -> MemoryArch {
+        self.arch
+    }
+
+    pub fn page_count_initial(&self) -> u64 {
+        self.page_count_initial
+    }
+
+    pub fn page_count_max(&self) -> u64 {
+        self.page_count_max.unwrap_or_else(|| max_page_count(self.page_size()))
+    }
+
+    pub fn page_size(&self) -> u64 {
+        self.page_size.unwrap_or(MEM_PAGE_SIZE)
+    }
+
+    pub fn initial_size(&self) -> u64 {
+        self.page_count_initial * self.page_size()
+    }
+
+    pub fn max_size(&self) -> u64 {
+        self.page_count_max() * self.page_size()
     }
 }
 
diff --git a/crates/wasm-testsuite/lib.rs b/crates/wasm-testsuite/lib.rs
diff --git a/crates/wasm-testsuite/tests/proposals/custom-page-sizes/custom-page-sizes-invalid.wast b/crates/wasm-testsuite/tests/proposals/custom-page-sizes/custom-page-sizes-invalid.wast
diff --git a/crates/wasm-testsuite/tests/proposals/custom-page-sizes/custom-page-sizes.wast b/crates/wasm-testsuite/tests/proposals/custom-page-sizes/custom-page-sizes.wast
