Add support for the custom-page-sizes proposal (#1489)

* wasmparser: Add support for the custom-page-sizes proposal

* wasm-encoder: Add support for the custom-page-sizes proposal

* wast: Add support for the custom-page-sizes proposal

* wasm-smith: Add support for the custom-page-sizes proposal

* wasmprinter: Add support for the custom-page-sizes proposal

* Add support for the custom-page-sizes proposal to the roundtrip test

* Get the workspace compiling again

* Add roundtrip tests for the custom-page-sizes proposal

* Fix compiler errors

* More compiler errors and test expectation updates

* Address review feedback

* more review feedback

* use module quote

Author: fitzgen

crates/wasm-encoder/src/core/data.rs

@@ -18,6 +18,7 @@ use crate::{encode_section, encoding_size, ConstExpr, Encode, Section, SectionId
 ///     maximum: None,
 ///     memory64: false,
 ///     shared: false,
+///     page_size_log2: None,
 /// });
 ///
 /// let mut data = DataSection::new();

crates/wasm-encoder/src/core/imports.rs

@@ -103,6 +103,7 @@ impl TryFrom<wasmparser::TypeRef> for EntityType {
 ///         maximum: None,
 ///         memory64: false,
 ///         shared: false,
+///         page_size_log2: None,
 ///     }
 /// );
 ///

crates/wasm-encoder/src/core/memories.rs

@@ -15,6 +15,7 @@ use crate::{encode_section, Encode, Section, SectionId};
 ///     maximum: None,
 ///     memory64: false,
 ///     shared: false,
+///     page_size_log2: None,
 /// });
 ///
 /// let mut module = Module::new();
@@ -75,26 +76,41 @@ pub struct MemoryType {
     pub memory64: bool,
     /// Whether or not this memory is shared.
     pub shared: bool,
+    /// The log base 2 of a custom page size for this memory.
+    ///
+    /// The default page size for Wasm memories is 64KiB, i.e. 2<sup>16</sup> or
+    /// `65536`.
+    ///
+    /// After the introduction of [the custom-page-sizes
+    /// proposal](https://github.com/WebAssembly/custom-page-sizes), Wasm can
+    /// customize the page size. It must be a power of two that is less than or
+    /// equal to 64KiB. Attempting to encode an invalid page size may panic.
+    pub page_size_log2: Option<u32>,
 }
 
 impl Encode for MemoryType {
     fn encode(&self, sink: &mut Vec<u8>) {
         let mut flags = 0;
         if self.maximum.is_some() {
-            flags |= 0b001;
+            flags |= 0b0001;
         }
         if self.shared {
-            flags |= 0b010;
+            flags |= 0b0010;
         }
         if self.memory64 {
-            flags |= 0b100;
+            flags |= 0b0100;
+        }
+        if self.page_size_log2.is_some() {
+            flags |= 0b1000;
         }
-
         sink.push(flags);
         self.minimum.encode(sink);
         if let Some(max) = self.maximum {
             max.encode(sink);
         }
+        if let Some(p) = self.page_size_log2 {
+            p.encode(sink);
+        }
     }
 }
 
@@ -106,6 +122,7 @@ impl From<wasmparser::MemoryType> for MemoryType {
             maximum: memory_ty.maximum,
             memory64: memory_ty.memory64,
             shared: memory_ty.shared,
+            page_size_log2: memory_ty.page_size_log2,
         }
     }
 }

crates/wasm-mutate/src/mutators/translate.rs

@@ -160,6 +160,7 @@ pub fn memory_type(
         minimum: ty.initial,
         maximum: ty.maximum,
         shared: ty.shared,
+        page_size_log2: ty.page_size_log2,
     })
 }
 

crates/wasm-shrink/src/lib.rs

@@ -236,6 +236,7 @@ impl ShrinkRun {
             component_model: false,
             function_references: false,
             gc: false,
+            custom_page_sizes: false,
             component_model_values: false,
             component_model_nested_names: false,
 

crates/wasm-smith/src/config.rs

@@ -313,6 +313,12 @@ define_config! {
         /// Defaults to `false`.
         pub gc_enabled: bool = false,
 
+        /// Determines whether the custom-page-sizes proposal is enabled when
+        /// generating a Wasm module.
+        ///
+        /// Defaults to `false`.
+        pub custom_page_sizes_enabled: bool = false,
+
         /// Returns whether we should generate custom sections or not. Defaults
         /// to false.
         pub generate_custom_sections: bool = false,
@@ -715,6 +721,7 @@ impl<'a> Arbitrary<'a> for Config {
             export_everything: false,
             tail_call_enabled: false,
             gc_enabled: false,
+            custom_page_sizes_enabled: false,
             generate_custom_sections: false,
             allow_invalid_funcs: false,
         })

crates/wasm-smith/src/core.rs

@@ -2324,7 +2324,7 @@ impl Module {
 
         // Interesting values related to memory bounds.
         for m in self.memories.iter() {
-            let min = m.minimum.saturating_mul(crate::WASM_PAGE_SIZE);
+            let min = m.minimum.saturating_mul(crate::page_size(m).into());
             interesting(min);
             for i in 0..5 {
                 if let Some(x) = min.checked_add(1 << i) {
@@ -2336,7 +2336,7 @@ impl Module {
             }
 
             if let Some(max) = m.maximum {
-                let max = max.saturating_mul(crate::WASM_PAGE_SIZE);
+                let max = max.saturating_mul(crate::page_size(m).into());
                 interesting(max);
                 for i in 0..5 {
                     if let Some(x) = max.checked_add(1 << i) {
@@ -2518,6 +2518,11 @@ pub(crate) fn arbitrary_memtype(u: &mut Unstructured, config: &Config) -> Result
     // We want to favor memories <= 1gb in size, allocate at most 16k pages,
     // depending on the maximum number of memories.
     let memory64 = config.memory64_enabled && u.arbitrary()?;
+    let page_size = if config.custom_page_sizes_enabled && u.arbitrary()? {
+        Some(1 << u.int_in_range(0..=16)?)
+    } else {
+        None
+    };
     let max_inbounds = 16 * 1024 / u64::try_from(config.max_memories).unwrap();
     let min_pages = if config.disallow_traps { Some(1) } else { None };
     let max_pages = min_pages.unwrap_or(0).max(if memory64 {
@@ -2537,6 +2542,7 @@ pub(crate) fn arbitrary_memtype(u: &mut Unstructured, config: &Config) -> Result
         maximum,
         memory64,
         shared,
+        page_size_log2: page_size,
     })
 }
 

crates/wasm-smith/src/core/code_builder.rs

@@ -5203,10 +5203,12 @@ fn memory_offset(u: &mut Unstructured, module: &Module, memory_index: u32) -> Re
     assert!(a + b + c != 0);
 
     let memory_type = &module.memories[memory_index as usize];
-    let min = memory_type.minimum.saturating_mul(crate::WASM_PAGE_SIZE);
+    let min = memory_type
+        .minimum
+        .saturating_mul(crate::page_size(memory_type).into());
     let max = memory_type
         .maximum
-        .map(|max| max.saturating_mul(crate::WASM_PAGE_SIZE))
+        .map(|max| max.saturating_mul(crate::page_size(memory_type).into()))
         .unwrap_or(u64::MAX);
 
     let (min, max, true_max) = match (memory_type.memory64, module.config.disallow_traps) {

crates/wasm-smith/src/core/code_builder/no_traps.rs

@@ -37,7 +37,10 @@ pub(crate) fn load<'a>(
             // []
             insts.push(Instruction::MemorySize(memarg.memory_index));
             // [mem_size_in_pages:address_type]
-            insts.push(int_const_inst(address_type, 65_536));
+            insts.push(int_const_inst(
+                address_type,
+                crate::page_size(memory).into(),
+            ));
             // [mem_size_in_pages:address_type wasm_page_size:address_type]
             insts.push(int_mul_inst(address_type));
             // [mem_size_in_bytes:address_type]
@@ -116,7 +119,10 @@ pub(crate) fn store<'a>(
     // []
     insts.push(Instruction::MemorySize(memarg.memory_index));
     // [mem_size_in_pages:address_type]
-    insts.push(int_const_inst(address_type, 65_536));
+    insts.push(int_const_inst(
+        address_type,
+        crate::page_size(memory).into(),
+    ));
     // [mem_size_in_pages:address_type wasm_page_size:address_type]
     insts.push(int_mul_inst(address_type));
     // [mem_size_in_bytes:address_type]

crates/wasm-smith/src/lib.rs

@@ -63,11 +63,15 @@ use arbitrary::{Result, Unstructured};
 pub use component::Component;
 pub use config::{Config, MemoryOffsetChoices};
 use std::{collections::HashSet, fmt::Write, str};
+use wasm_encoder::MemoryType;
 
 #[cfg(feature = "_internal_cli")]
 pub use config::InternalOptionalConfig;
 
-const WASM_PAGE_SIZE: u64 = 65_536;
+pub(crate) fn page_size(mem: &MemoryType) -> u32 {
+    const DEFAULT_WASM_PAGE_SIZE: u32 = 65_536;
+    mem.page_size_log2.unwrap_or(DEFAULT_WASM_PAGE_SIZE)
+}
 
 /// Do something an arbitrary number of times.
 ///