fix(decoder): GC ref type decoding and custom page sizes support

kiln-build-core/src/wast_execution.rs

@@ -795,6 +795,7 @@ impl WastEngine {
                     limits: Limits { min: 1, max: Some(2) },
                     shared: false,
                     memory64: false,
+                    page_size: None,
                 };
                 validate_memory_import_compatibility(mem_type, &spectest_mem)?;
                 Ok(())
@@ -1208,6 +1209,7 @@ impl WastEngine {
                     limits: core_ty.limits,
                     shared: core_ty.shared,
                     memory64: false,
+                    page_size: None,
                 })
             } else {
                 None

kiln-component/src/adapter.rs

@@ -358,6 +358,7 @@ impl CoreModuleAdapter {
                     },
                     shared: mem_adapter.shared,
                     memory64: false,
+                    page_size: None,
                 })?),
                 kind: ExportKind::Value {
                     value_index: mem_adapter.core_index,

kiln-decoder/src/sections_no_std.rs

@@ -387,6 +387,7 @@ pub mod parsers {
                         },
                         shared: mem_limits.shared,
                         memory64: mem_limits.memory64,
+                        page_size: None,
                     };
                     KilnImportDesc::Memory(memory_type)
                 },
@@ -477,6 +478,7 @@ pub mod parsers {
                 limits: kiln_limits,
                 shared: limits.shared,
                 memory64: limits.memory64,
+                page_size: None,
             };
 
             memories

kiln-decoder/src/streaming_decoder.rs

@@ -1052,8 +1052,12 @@ impl<'a> StreamingDecoder<'a> {
                     offset += 1;
 
                     // Validate limits flags per WebAssembly spec:
-                    // Maximum valid flag for memory is 0x07 (has_max | shared | memory64)
-                    if flags > 0x07 {
+                    // - Bit 0: has max (0x01)
+                    // - Bit 1: shared (0x02) - threads proposal
+                    // - Bit 2: memory64 (0x04)
+                    // - Bit 3: custom page size (0x08) - custom-page-sizes proposal
+                    // All other bits must be zero. Maximum valid flag is 0x0F.
+                    if flags > 0x0F {
                         return Err(Error::parse_error("malformed limits flags"));
                     }
 
@@ -1119,6 +1123,25 @@ impl<'a> StreamingDecoder<'a> {
                         (min, max)
                     };
 
+                    // Custom page size (bit 3): read page size as LEB128 u32
+                    let custom_page_size = if (flags & 0x08) != 0 {
+                        let (ps, bytes_read) = read_leb128_u32(data, offset)?;
+                        offset += bytes_read;
+                        if ps > 65536 {
+                            return Err(Error::validation_error(
+                                "custom page size must be at most 65536",
+                            ));
+                        }
+                        if ps != 0 && (ps & (ps - 1)) != 0 {
+                            return Err(Error::validation_error(
+                                "custom page size must be a power of 2",
+                            ));
+                        }
+                        Some(ps)
+                    } else {
+                        None
+                    };
+
                     #[cfg(feature = "tracing")]
                     trace!(import_index = i, min_pages = min, max_pages = ?max, "import: memory");
 
@@ -1134,6 +1157,7 @@ impl<'a> StreamingDecoder<'a> {
                             limits,
                             shared: flags & 0x02 != 0, // bit 1 = shared
                             memory64: flags & 0x04 != 0, // bit 2 = memory64
+                            page_size: custom_page_size,
                         };
 
                         let import = Import {
@@ -1149,12 +1173,19 @@ impl<'a> StreamingDecoder<'a> {
                 },
                 0x03 => {
                     // Global import - need to parse global type
-                    // value_type (1 byte) + mutability (1 byte)
-                    if offset + 1 >= data.len() {
+                    // value_type (potentially multi-byte for GC ref types) + mutability (1 byte)
+                    if offset >= data.len() {
+                        return Err(Error::parse_error("Unexpected end of global import"));
+                    }
+
+                    // Parse value type using GC-aware parser to handle multi-byte
+                    // ref type encodings (0x63/0x64 + heap type index)
+                    let (value_type, new_offset) = self.parse_value_type(data, offset)?;
+                    offset = new_offset;
+
+                    if offset >= data.len() {
                         return Err(Error::parse_error("Unexpected end of global import"));
                     }
-                    let value_type_byte = data[offset];
-                    offset += 1;
                     let mutability_byte = data[offset];
                     offset += 1;
 
@@ -1163,19 +1194,6 @@ impl<'a> StreamingDecoder<'a> {
                         return Err(Error::parse_error("malformed mutability"));
                     }
 
-                    // Parse value type
-                    let value_type = match value_type_byte {
-                        0x7F => kiln_foundation::ValueType::I32,
-                        0x7E => kiln_foundation::ValueType::I64,
-                        0x7D => kiln_foundation::ValueType::F32,
-                        0x7C => kiln_foundation::ValueType::F64,
-                        0x7B => kiln_foundation::ValueType::V128,
-                        0x70 => kiln_foundation::ValueType::FuncRef,
-                        0x6F => kiln_foundation::ValueType::ExternRef,
-                        0x69 => kiln_foundation::ValueType::ExnRef,
-                        _ => return Err(Error::parse_error("Invalid global import value type")),
-                    };
-
                     #[cfg(feature = "tracing")]
                     trace!(import_index = i, value_type = ?value_type, mutable = (mutability_byte != 0), "import: global");
 
@@ -1519,11 +1537,12 @@ impl<'a> StreamingDecoder<'a> {
             offset += 1;
 
             // Validate limits flags per WebAssembly spec:
-            // - Bits 0: has max (0x01)
+            // - Bit 0: has max (0x01)
             // - Bit 1: shared (0x02) - threads proposal
             // - Bit 2: memory64 (0x04)
-            // All other bits must be zero. Maximum valid flag is 0x07.
-            if flags > 0x07 {
+            // - Bit 3: custom page size (0x08) - custom-page-sizes proposal
+            // All other bits must be zero. Maximum valid flag is 0x0F.
+            if flags > 0x0F {
                 return Err(Error::parse_error("malformed limits flags"));
             }
 
@@ -1597,11 +1616,31 @@ impl<'a> StreamingDecoder<'a> {
                 }
             }
 
+            // Custom page size (bit 3): read page size as LEB128 u32
+            let custom_page_size = if (flags & 0x08) != 0 {
+                let (ps, bytes_read) = read_leb128_u32(data, offset)?;
+                offset += bytes_read;
+                if ps > 65536 {
+                    return Err(Error::validation_error(
+                        "custom page size must be at most 65536",
+                    ));
+                }
+                if ps != 0 && (ps & (ps - 1)) != 0 {
+                    return Err(Error::validation_error(
+                        "custom page size must be a power of 2",
+                    ));
+                }
+                Some(ps)
+            } else {
+                None
+            };
+
             // Create memory type
             let memory_type = kiln_foundation::types::MemoryType {
                 limits: kiln_foundation::types::Limits { min, max },
                 shared,
                 memory64: is_memory64,
+                page_size: custom_page_size,
             };
 
             // Add to module
@@ -2303,44 +2342,42 @@ impl<'a> StreamingDecoder<'a> {
                 "code section: function locals"
             );
 
-            // Code section index i corresponds to module-defined function at index (num_imports + i)
-            let func_index = num_imports + i as usize;
-            if let Some(func) = self.module.functions.get_mut(func_index) {
-                // Parse local variable declarations and track total count
-                let mut total_locals: u64 = 0;
-
-                for _ in 0..local_count {
-                    let (count, bytes) = read_leb128_u32(&data[body_start..body_end], body_offset)?;
-                    body_offset += bytes;
+            // Parse local type groups before taking a mutable borrow on self.module.functions,
+            // because parse_value_type borrows &self and get_mut borrows &mut self.module.
+            let mut local_groups = Vec::new();
+            let mut total_locals: u64 = 0;
+            for _ in 0..local_count {
+                let (count, bytes) = read_leb128_u32(&data[body_start..body_end], body_offset)?;
+                body_offset += bytes;
 
-                    if body_offset >= body_size as usize {
-                        return Err(Error::parse_error("Unexpected end of function body"));
-                    }
+                if body_offset >= body_size as usize {
+                    return Err(Error::parse_error("Unexpected end of function body"));
+                }
 
-                    let value_type = data[body_start + body_offset];
-                    body_offset += 1;
-
-                    // Convert to ValueType and add to locals
-                    let vt = match value_type {
-                        0x7F => kiln_foundation::types::ValueType::I32,
-                        0x7E => kiln_foundation::types::ValueType::I64,
-                        0x7D => kiln_foundation::types::ValueType::F32,
-                        0x7C => kiln_foundation::types::ValueType::F64,
-                        0x7B => kiln_foundation::types::ValueType::V128,
-                        0x70 => kiln_foundation::types::ValueType::FuncRef,
-                        0x6F => kiln_foundation::types::ValueType::ExternRef,
-                        0x69 => kiln_foundation::types::ValueType::ExnRef,
-                        _ => return Err(Error::parse_error("Invalid local type")),
-                    };
+                // Parse value type using the full GC-aware parser to handle
+                // multi-byte ref type encodings (0x63/0x64 + heap type index)
+                let (vt, new_body_offset) = self.parse_value_type(
+                    &data[body_start..body_end],
+                    body_offset,
+                )?;
+                body_offset = new_body_offset;
+
+                // Validate total locals: sum of all declared locals must fit in u32
+                total_locals += count as u64;
+                if total_locals > u32::MAX as u64 {
+                    return Err(Error::parse_error("too many locals"));
+                }
 
-                    // Validate total locals: sum of all declared locals must fit in u32
-                    total_locals += count as u64;
-                    if total_locals > u32::MAX as u64 {
-                        return Err(Error::parse_error("too many locals"));
-                    }
+                local_groups.push((count, vt));
+            }
 
+            // Code section index i corresponds to module-defined function at index (num_imports + i)
+            let func_index = num_imports + i as usize;
+            if let Some(func) = self.module.functions.get_mut(func_index) {
+                // Apply parsed local declarations to the function
+                for (count, vt) in &local_groups {
                     // Validate total locals against platform limits before allocation
-                    let new_total = func.locals.len() + count as usize;
+                    let new_total = func.locals.len() + *count as usize;
                     if new_total > limits::MAX_FUNCTION_LOCALS {
                         return Err(Error::parse_error(
                             "Function exceeds maximum local count for platform",
@@ -2352,12 +2389,12 @@ impl<'a> StreamingDecoder<'a> {
                         AllocationPhase::Decode,
                         "streaming_decoder:func_locals",
                         "locals",
-                        count as usize
+                        *count as usize
                     );
 
                     // Add 'count' locals of this type
-                    for _ in 0..count {
-                        func.locals.push(vt);
+                    for _ in 0..*count {
+                        func.locals.push(*vt);
                     }
                 }
 

kiln-foundation/src/types.rs

@@ -3905,16 +3905,20 @@ pub struct MemoryType {
     pub shared: bool,
     /// Memory64 extension - uses i64 addresses instead of i32
     pub memory64: bool,
+    /// Custom page size (custom-page-sizes proposal).
+    /// When `Some`, the page size in bytes; must be a power of two and <= 65536.
+    /// When `None`, the standard 65536-byte page size is used.
+    pub page_size: Option<u32>,
 }
 
 impl MemoryType {
     pub const fn new(limits: Limits, shared: bool) -> Self {
-        Self { limits, shared, memory64: false }
+        Self { limits, shared, memory64: false, page_size: None }
     }
 
     /// Create a memory type with all options
     pub const fn new_with_memory64(limits: Limits, shared: bool, memory64: bool) -> Self {
-        Self { limits, shared, memory64 }
+        Self { limits, shared, memory64, page_size: None }
     }
 }
 
@@ -3923,6 +3927,10 @@ impl Checksummable for MemoryType {
         self.limits.update_checksum(checksum);
         checksum.update(self.shared as u8);
         checksum.update(self.memory64 as u8);
+        let ps = self.page_size.unwrap_or(0);
+        for byte in ps.to_le_bytes() {
+            checksum.update(byte);
+        }
     }
 }
 
@@ -3935,6 +3943,15 @@ impl ToBytes for MemoryType {
         self.limits.to_bytes_with_provider(writer, provider)?;
         writer.write_u8(self.shared as u8)?;
         writer.write_u8(self.memory64 as u8)?;
+        match self.page_size {
+            Some(ps) => {
+                writer.write_u8(1)?;
+                writer.write_u32_le(ps)?;
+            }
+            None => {
+                writer.write_u8(0)?;
+            }
+        }
         Ok(())
     }
     // Default to_bytes method will be used if #cfg(feature = "default-provider") is
@@ -3967,7 +3984,13 @@ impl FromBytes for MemoryType {
                 ));
             },
         };
-        Ok(MemoryType { limits, shared, memory64 })
+        let page_size_flag = reader.read_u8()?;
+        let page_size = if page_size_flag != 0 {
+            Some(reader.read_u32_le()?)
+        } else {
+            None
+        };
+        Ok(MemoryType { limits, shared, memory64, page_size })
     }
     // Default from_bytes method will be used if #cfg(feature = ")
     // is active

kiln-runtime/src/memory.rs

@@ -517,6 +517,7 @@ impl kiln_foundation::traits::FromBytes for Memory {
             },
             shared: false,
             memory64: false,
+            page_size: None,
         };
         Self::new(to_core_memory_type(&memory_type)).map(|boxed| *boxed)
     }

kiln-runtime/src/memory_partial.rs

@@ -406,6 +406,8 @@ impl kiln_foundation::traits::FromBytes for Memory {
         let memory_type = MemoryType {
             limits: Limits { min, max: if max == 0 { None } else { Some(max) } },
             shared: false,
+            memory64: false,
+            page_size: None,
         };
         Self::new(memory_type)
     }

kiln-runtime/src/module.rs

@@ -3709,6 +3709,7 @@ impl Module {
                         limits: KilnLimits { min: 0, max: None },  // Will be resolved via linking
                         shared: true,  // Component Model uses shared memory
                         memory64: false,
+                        page_size: None,
                     };
                     ExternType::Memory(memory_type)
                 },
@@ -3836,6 +3837,7 @@ impl Module {
                 },
                 shared: false,
                 memory64: false,
+                page_size: None,
             };
             runtime_module
                 .push_memory(MemoryWrapper::new(Memory::new(to_core_memory_type(