reencoder: allow implementers to fail to reencode indices and such (#2179)

I have found myself wanting to return an error from
`Reencode::type_index`. This is admittedly a pretty invasive change
though and I punted on component stuff.

Author: nagisa

crates/wasm-encoder/src/reencode.rs

@@ -40,14 +40,14 @@ impl<'cfg, 'wasm> Reencode for InitTranslator<'cfg, 'wasm> {
 
     /// Handle `elem`s with values of the `ElementItem::Func` kind. This function will not be
     /// called for values of the `ElementItem::Expr` kind.
-    fn function_index(&mut self, idx: u32) -> u32 {
+    fn function_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         if self.kind != ConstExpressionMutator::ElementFunc || !self.should_process() {
-            return idx;
+            return Ok(idx);
         }
 
         log::trace!("... replacing referenced function index with 0");
         // FIXME: generate random function indices when `!config.reduce`.
-        0
+        Ok(0)
     }
 
     /// Handle `global` initializers and `elem`s with values of the `ElementItem::Expr` kind.

crates/wasm-encoder/src/reencode/component.rs

@@ -187,7 +187,7 @@ impl RemoveItem<'_> {
     ///   try something else.
     /// * Finally our index is larger than the one being removed which means we
     ///   now decrement it by one to account for the removed item.
-    fn remap(&mut self, item: Item, idx: u32) -> u32 {
+    fn remap(&mut self, item: Item, idx: u32) -> ReencodeResult<u32> {
         // If we're before the code section then all function references, no
         // matter where they are, are considered "referencing functions" so we
         // save the indices of that which is referenced.
@@ -197,7 +197,7 @@ impl RemoveItem<'_> {
             }
         }
 
-        if item != self.item || idx < self.idx {
+        Ok(if item != self.item || idx < self.idx {
             // Different kind of item or a later item was removed, index doesn't change
             idx
         } else if idx == self.idx {
@@ -210,42 +210,42 @@ impl RemoveItem<'_> {
             // Otherwise this item comes after the item being removed, so
             // this item's index has decreased by one.
             idx - 1
-        }
+        })
     }
 }
 
 impl Reencode for RemoveItem<'_> {
     type Error = Error;
 
-    fn function_index(&mut self, idx: u32) -> u32 {
+    fn function_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Function, idx)
     }
 
-    fn table_index(&mut self, idx: u32) -> u32 {
+    fn table_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Table, idx)
     }
 
-    fn memory_index(&mut self, idx: u32) -> u32 {
+    fn memory_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Memory, idx)
     }
 
-    fn tag_index(&mut self, idx: u32) -> u32 {
+    fn tag_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Tag, idx)
     }
 
-    fn global_index(&mut self, idx: u32) -> u32 {
+    fn global_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Global, idx)
     }
 
-    fn type_index(&mut self, idx: u32) -> u32 {
+    fn type_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Type, idx)
     }
 
-    fn data_index(&mut self, idx: u32) -> u32 {
+    fn data_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Data, idx)
     }
 
-    fn element_index(&mut self, idx: u32) -> u32 {
+    fn element_index(&mut self, idx: u32) -> ReencodeResult<u32> {
         self.remap(Item::Element, idx)
     }
 
@@ -312,7 +312,7 @@ impl Reencode for RemoveItem<'_> {
         section: wasmparser::FunctionSectionReader<'_>,
     ) -> ReencodeResult<()> {
         self.filter_out(section, Item::Function, |me, func| {
-            funcs.function(me.type_index(func));
+            funcs.function(me.type_index(func)?);
             Ok(())
         })
     }
@@ -334,7 +334,7 @@ impl Reencode for RemoveItem<'_> {
         section: wasmparser::MemorySectionReader<'_>,
     ) -> ReencodeResult<()> {
         self.filter_out(section, Item::Memory, |me, memory| {
-            memories.memory(me.memory_type(memory));
+            memories.memory(me.memory_type(memory)?);
             Ok(())
         })
     }
@@ -378,23 +378,23 @@ impl Reencode for RemoveItem<'_> {
         section: wasmparser::TagSectionReader<'_>,
     ) -> ReencodeResult<()> {
         self.filter_out(section, Item::Tag, |me, tag| {
-            tags.tag(me.tag_type(tag));
+            tags.tag(me.tag_type(tag)?);
             Ok(())
         })
     }
 
-    fn data_count(&mut self, count: u32) -> u32 {
+    fn data_count(&mut self, count: u32) -> ReencodeResult<u32> {
         // Note that the data count section is decremented here if
         // we're removing a data item, otherwise it's preserved
         // as-is.
-        if self.item == Item::Data {
+        Ok(if self.item == Item::Data {
             count - 1
         } else {
             count
-        }
+        })
     }
 
-    fn start_section(&mut self, func: u32) -> u32 {
+    fn start_section(&mut self, func: u32) -> ReencodeResult<u32> {
         self.function_reference_action = Funcref::Skip;
         let func = self.remap(Item::Function, func);
         self.function_reference_action = Funcref::Save;

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

@@ -13,7 +13,7 @@ use std::{
 };
 use wasm_encoder::{
     self,
-    reencode::{utils::parse_custom_section, Reencode},
+    reencode::{self, utils::parse_custom_section, Reencode},
 };
 use wasmparser::{self, BinaryReaderError, Import, KnownCustom, Parser, TypeRef};
 
@@ -138,7 +138,7 @@ struct DedupingReencoder {
     remappings: Remappings,
 }
 
-type DeduperError = wasm_encoder::reencode::Error<NoDuplicatesFound>;
+type DeduperError = reencode::Error<NoDuplicatesFound>;
 
 impl Reencode for DedupingReencoder {
     type Error = NoDuplicatesFound;
@@ -149,7 +149,7 @@ impl Reencode for DedupingReencoder {
         &mut self,
         imports: &mut wasm_encoder::ImportSection,
         section: wasmparser::ImportSectionReader<'_>,
-    ) -> Result<(), wasm_encoder::reencode::Error<Self::Error>> {
+    ) -> Result<(), reencode::Error<Self::Error>> {
         self.remappings.imports(section.clone())?;
 
         // If no duplicates, take the fast path out.
@@ -165,8 +165,8 @@ impl Reencode for DedupingReencoder {
         Ok(())
     }
 
-    fn function_index(&mut self, func_idx: u32) -> u32 {
-        self.remappings.new_index_for(func_idx)
+    fn function_index(&mut self, func_idx: u32) -> Result<u32, reencode::Error<Self::Error>> {
+        Ok(self.remappings.new_index_for(func_idx))
     }
 
     // Strip all but known-safe custom sections.
@@ -181,7 +181,7 @@ impl Reencode for DedupingReencoder {
         &mut self,
         module: &mut wasm_encoder::Module,
         section: wasmparser::CustomSectionReader<'_>,
-    ) -> Result<(), wasm_encoder::reencode::Error<Self::Error>> {
+    ) -> Result<(), reencode::Error<Self::Error>> {
         match section.as_known() {
             KnownCustom::Name(_) | KnownCustom::Producers(_) => {
                 // Keep this section verbatim:

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

@@ -538,10 +538,10 @@ impl<'a> Module<'a> {
             let ty = map.memory_type(mem.ty);
             match &mem.def {
                 Definition::Import(m, n) => {
-                    imports.import(m, n, ty);
+                    imports.import(m, n, ty?);
                 }
                 Definition::Local(()) => {
-                    memories.memory(ty);
+                    memories.memory(ty?);
                 }
             }
             num_memories += 1;
@@ -1030,23 +1030,25 @@ struct Encoder {
     tables: Remap,
 }
 
+type ReencodeResult<T> = Result<T, wasm_encoder::reencode::Error<Infallible>>;
+
 impl Reencode for Encoder {
     type Error = Infallible;
 
-    fn type_index(&mut self, i: u32) -> u32 {
-        self.types.remap(i)
+    fn type_index(&mut self, i: u32) -> ReencodeResult<u32> {
+        Ok(self.types.remap(i))
     }
-    fn function_index(&mut self, i: u32) -> u32 {
-        self.funcs.remap(i)
+    fn function_index(&mut self, i: u32) -> ReencodeResult<u32> {
+        Ok(self.funcs.remap(i))
     }
-    fn memory_index(&mut self, i: u32) -> u32 {
-        self.memories.remap(i)
+    fn memory_index(&mut self, i: u32) -> ReencodeResult<u32> {
+        Ok(self.memories.remap(i))
     }
-    fn global_index(&mut self, i: u32) -> u32 {
-        self.globals.remap(i)
+    fn global_index(&mut self, i: u32) -> ReencodeResult<u32> {
+        Ok(self.globals.remap(i))
     }
-    fn table_index(&mut self, i: u32) -> u32 {
-        self.tables.remap(i)
+    fn table_index(&mut self, i: u32) -> ReencodeResult<u32> {
+        Ok(self.tables.remap(i))
     }
 }
 

crates/wit-component/src/encoding/dedupe.rs

@@ -171,8 +171,8 @@ struct RemoveImports<'a, 'b> {
 impl Reencode for RemoveImports<'_, '_> {
     type Error = std::convert::Infallible;
 
-    fn function_index(&mut self, idx: u32) -> u32 {
-        idx - self.removed_funcs
+    fn function_index(&mut self, idx: u32) -> Result<u32, reencode::Error<Self::Error>> {
+        Ok(idx - self.removed_funcs)
     }
 
     fn parse_import(

crates/wit-component/src/gc.rs

@@ -1273,12 +1273,12 @@ struct ModuleExtractor<'a> {
 impl Reencode for ModuleExtractor<'_> {
     type Error = std::convert::Infallible;
 
-    fn type_index(&mut self, index: u32) -> u32 {
-        if self.depth == 0 {
+    fn type_index(&mut self, index: u32) -> Result<u32, Error<Self::Error>> {
+        Ok(if self.depth == 0 {
             index + self.extra_core_types
         } else {
             index
-        }
+        })
     }
 }
 
@@ -1291,12 +1291,12 @@ impl ReencodeComponent for ModuleExtractor<'_> {
         }
     }
 
-    fn outer_type_index(&mut self, count: u32, index: u32) -> u32 {
-        if self.depth == count {
+    fn outer_type_index(&mut self, count: u32, index: u32) -> Result<u32, Error<Self::Error>> {
+        Ok(if self.depth == count {
             index + self.extra_core_types
         } else {
             index
-        }
+        })
     }
 
     fn outer_module_index(&mut self, count: u32, index: u32) -> u32 {