Fix wasm-mutate with table64 (bytecodealliance#1535)

Consult module information for the type of `table.*` expressions to know if it's 32-bit or 64-bit.
Robbepop · May 6, 2024 · cb87499 · cb87499
1 parent b1ee957
commit cb87499
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/crates/wasm-mutate/src/info.rs b/crates/wasm-mutate/src/info.rs
@@ -49,7 +49,7 @@ pub struct ModuleInfo<'a> {
     // function idx to type idx
     pub function_map: Vec<u32>,
     pub global_types: Vec<PrimitiveTypeInfo>,
-    pub table_elem_types: Vec<PrimitiveTypeInfo>,
+    pub table_types: Vec<wasmparser::TableType>,
     pub memory_types: Vec<wasmparser::MemoryType>,
 
     // raw_sections
@@ -119,7 +119,7 @@ impl<'a> ModuleInfo<'a> {
                             wasmparser::TypeRef::Table(ty) => {
                                 info.table_count += 1;
                                 info.imported_tables_count += 1;
-                                info.table_elem_types.push(ty.element_type.into());
+                                info.table_types.push(ty);
                             }
                             wasmparser::TypeRef::Tag(_ty) => {
                                 info.tag_count += 1;
@@ -143,8 +143,7 @@ impl<'a> ModuleInfo<'a> {
 
                     for table in reader {
                         let table = table?;
-                        let ty = PrimitiveTypeInfo::try_from(table.ty.element_type).unwrap();
-                        info.table_elem_types.push(ty);
+                        info.table_types.push(table.ty);
                     }
                 }
                 Payload::MemorySection(reader) => {

diff --git a/crates/wasm-mutate/src/mutators/peephole/eggsy/analysis.rs b/crates/wasm-mutate/src/mutators/peephole/eggsy/analysis.rs
@@ -24,7 +24,7 @@ pub struct PeepholeMutationAnalysis {
     /// function idx to type idx
     function_map: Vec<u32>,
     /// table index to the type of element it has
-    table_types: Vec<PrimitiveTypeInfo>,
+    table_types: Vec<wasmparser::TableType>,
     memory_types: Vec<wasmparser::MemoryType>,
 }
 
@@ -36,7 +36,7 @@ impl PeepholeMutationAnalysis {
             global_types: info.global_types.clone(),
             types_map: info.types_map.clone(),
             function_map: info.function_map.clone(),
-            table_types: info.table_elem_types.clone(),
+            table_types: info.table_types.clone(),
             memory_types: info.memory_types.clone(),
         }
     }
@@ -267,16 +267,30 @@ impl PeepholeMutationAnalysis {
             Lang::DataDrop(_) => Ok(PrimitiveTypeInfo::Empty),
             Lang::MemoryCopy { .. } => Ok(PrimitiveTypeInfo::Empty),
             Lang::MemoryFill { .. } => Ok(PrimitiveTypeInfo::Empty),
-            Lang::TableGrow { .. } => Ok(PrimitiveTypeInfo::I32),
-            Lang::TableSize { .. } => Ok(PrimitiveTypeInfo::I32),
+            Lang::TableGrow(table, _) => {
+                let ty = self.table_types[*table as usize];
+                if ty.table64 {
+                    Ok(PrimitiveTypeInfo::I64)
+                } else {
+                    Ok(PrimitiveTypeInfo::I32)
+                }
+            }
+            Lang::TableSize(table) => {
+                let ty = self.table_types[*table as usize];
+                if ty.table64 {
+                    Ok(PrimitiveTypeInfo::I64)
+                } else {
+                    Ok(PrimitiveTypeInfo::I32)
+                }
+            }
             Lang::TableInit { .. } => Ok(PrimitiveTypeInfo::Empty),
             Lang::ElemDrop(_) => Ok(PrimitiveTypeInfo::Empty),
             Lang::TableCopy { .. } => Ok(PrimitiveTypeInfo::Empty),
             Lang::TableFill { .. } => Ok(PrimitiveTypeInfo::Empty),
             Lang::TableSet(..) => Ok(PrimitiveTypeInfo::Empty),
             Lang::TableGet(idx, _) => {
                 let ty = self.table_types[*idx as usize].clone();
-                Ok(ty)
+                Ok(ty.element_type.into())
             }
             Lang::I32UseGlobal(_) => Ok(PrimitiveTypeInfo::I32),
             Lang::I64UseGlobal(_) => Ok(PrimitiveTypeInfo::I64),