From 265ee58d0129e4151a7da5c3ef3ed432e8562f09 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Sat, 20 Jan 2024 15:06:45 +0300
Subject: [PATCH] feat(wasm-smith): add `reserved_memory_size` and tweak
 `memory_offset_choices`

---
 crates/wasm-smith/src/config.rs            | 13 ++++++++++++-
 crates/wasm-smith/src/core/code_builder.rs |  5 +++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/crates/wasm-smith/src/config.rs b/crates/wasm-smith/src/config.rs
index ae85c6cad7..64adbf758e 100644
--- a/crates/wasm-smith/src/config.rs
+++ b/crates/wasm-smith/src/config.rs
@@ -323,6 +323,11 @@ pub trait Config: 'static + std::fmt::Debug {
         false
     }
 
+    /// The size of reserved memory on the last memory page. Defaults to None.
+    fn reserved_memory_size(&self) -> Option<u64> {
+        None
+    }
+
     /// Determines whether the tail calls proposal is enabled for generating
     /// instructions.
     ///
@@ -560,6 +565,7 @@ pub struct SwarmConfig {
     pub min_uleb_size: u8,
     pub multi_value_enabled: bool,
     pub reference_types_enabled: bool,
+    pub reserved_memory_size: Option<u64>,
     pub tail_call_enabled: bool,
     pub relaxed_simd_enabled: bool,
     pub saturating_float_to_int_enabled: bool,
@@ -595,6 +601,7 @@ impl<'a> Arbitrary<'a> for SwarmConfig {
             max_tables,
             max_memory_pages: u.arbitrary()?,
             min_uleb_size: u.int_in_range(0..=5)?,
+            reserved_memory_size: None,
             bulk_memory_enabled: reference_types_enabled || u.arbitrary()?,
             reference_types_enabled,
             simd_enabled: u.arbitrary()?,
@@ -637,7 +644,7 @@ impl<'a> Arbitrary<'a> for SwarmConfig {
             max_modules: 0,
             max_components: 0,
             max_values: 0,
-            memory_offset_choices: (75, 24, 1),
+            memory_offset_choices: (75, 25, 0),
             allow_start_export: true,
             relaxed_simd_enabled: false,
             exceptions_enabled: false,
@@ -784,6 +791,10 @@ impl Config for SwarmConfig {
         self.reference_types_enabled
     }
 
+    fn reserved_memory_size(&self) -> Option<u64> {
+        self.reserved_memory_size
+    }
+
     fn tail_call_enabled(&self) -> bool {
         self.tail_call_enabled
     }
diff --git a/crates/wasm-smith/src/core/code_builder.rs b/crates/wasm-smith/src/core/code_builder.rs
index 81ac4123ac..ea16bc1f6b 100644
--- a/crates/wasm-smith/src/core/code_builder.rs
+++ b/crates/wasm-smith/src/core/code_builder.rs
@@ -4625,6 +4625,11 @@ fn memory_offset(u: &mut Unstructured, module: &Module, memory_index: u32) -> Re
 
     let choice = u.int_in_range(0..=a + b + c - 1)?;
     if choice < a {
+        let min = module
+            .config
+            .reserved_memory_size()
+            .map(|reserved| min.saturating_sub(reserved).saturating_sub(16))
+            .unwrap_or(min);
         u.int_in_range(0..=min)
     } else if choice < a + b {
         u.int_in_range(min..=max)