Start removing custom translation in wasm-mutate (#1794)

Now that `wasm_encoder::reencode` exists that's the better option to use
for reencoding modules. This commit removes one of the major users of
the equivalent `Translator` trait in `wasm_mutate` to start down the
path of removing that trait entirely and instead relying on the support
in `wasm_encoder`.

Author: alexcrichton

crates/wasm-encoder/src/reencode.rs

@@ -486,6 +486,14 @@ pub trait Reencode {
     ) -> Result<(), Error<Self::Error>> {
         utils::parse_custom_name_subsection(self, names, section)
     }
+
+    fn data_count(&mut self, count: u32) -> u32 {
+        count
+    }
+
+    fn start_section(&mut self, start: u32) -> u32 {
+        self.function_index(start)
+    }
 }
 
 /// An error when re-encoding from `wasmparser` to `wasm-encoder`.
@@ -498,6 +506,8 @@ pub enum Error<E = Infallible> {
     /// The const expression is invalid: not actually constant or something like
     /// that.
     InvalidConstExpr,
+    /// The code section size listed was not valid for the wasm binary provided.
+    InvalidCodeSectionSize,
     /// There was a section that does not belong into a core wasm module.
     UnexpectedNonCoreModuleSection,
     /// There was a section that does not belong into a compoennt module.
@@ -539,6 +549,7 @@ impl<E: std::fmt::Display> std::fmt::Display for Error<E> {
             Self::UnsupportedCoreTypeInComponent => {
                 fmt.write_str("unsupported core type in a component")
             }
+            Self::InvalidCodeSectionSize => fmt.write_str("invalid code section size"),
         }
     }
 }
@@ -552,7 +563,8 @@ impl<E: 'static + std::error::Error> std::error::Error for Error<E> {
             | Self::CanonicalizedHeapTypeReference
             | Self::UnexpectedNonCoreModuleSection
             | Self::UnexpectedNonComponentSection
-            | Self::UnsupportedCoreTypeInComponent => None,
+            | Self::UnsupportedCoreTypeInComponent
+            | Self::InvalidCodeSectionSize => None,
         }
     }
 }
@@ -588,11 +600,10 @@ pub mod utils {
             reencoder.intersperse_section_hook(module, after, before)
         }
 
-        let mut sections = parser.parse_all(data);
-        let mut next_section = sections.next();
+        let orig_offset = parser.offset() as usize;
         let mut last_section = None;
 
-        'outer: while let Some(section) = next_section {
+        for section in parser.parse_all(data) {
             match section? {
                 wasmparser::Payload::Version {
                     encoding: wasmparser::Encoding::Module,
@@ -697,7 +708,7 @@ pub mod utils {
                         Some(crate::SectionId::Start),
                     )?;
                     module.section(&crate::StartSection {
-                        function_index: reencoder.function_index(func),
+                        function_index: reencoder.start_section(func),
                     });
                 }
                 wasmparser::Payload::ElementSection(section) => {
@@ -718,6 +729,7 @@ pub mod utils {
                         &mut last_section,
                         Some(crate::SectionId::DataCount),
                     )?;
+                    let count = reencoder.data_count(count);
                     module.section(&crate::DataCountSection { count });
                 }
                 wasmparser::Payload::DataSection(section) => {
@@ -731,47 +743,37 @@ pub mod utils {
                     reencoder.parse_data_section(&mut data, section)?;
                     module.section(&data);
                 }
-                wasmparser::Payload::CodeSectionStart { count, .. } => {
+                wasmparser::Payload::CodeSectionStart { range, .. } => {
                     handle_intersperse_section_hook(
                         reencoder,
                         module,
                         &mut last_section,
                         Some(crate::SectionId::Code),
                     )?;
                     let mut codes = crate::CodeSection::new();
-                    for _ in 0..count {
-                        if let Some(Ok(wasmparser::Payload::CodeSectionEntry(section))) =
-                            sections.next()
-                        {
-                            reencoder.parse_function_body(&mut codes, section)?;
-                        } else {
-                            return Err(Error::UnexpectedNonCoreModuleSection);
-                        }
-                    }
-                    module.section(&codes);
-                }
-                wasmparser::Payload::CodeSectionEntry(section) => {
-                    handle_intersperse_section_hook(
-                        reencoder,
-                        module,
-                        &mut last_section,
-                        Some(crate::SectionId::Code),
-                    )?;
-                    // we can't do better than start a new code section here
-                    let mut codes = crate::CodeSection::new();
-                    reencoder.parse_function_body(&mut codes, section)?;
-                    while let Some(section) = sections.next() {
-                        let section = section?;
-                        if let wasmparser::Payload::CodeSectionEntry(section) = section {
-                            reencoder.parse_function_body(&mut codes, section)?;
-                        } else {
-                            module.section(&codes);
-                            next_section = Some(Ok(section));
-                            continue 'outer;
-                        }
-                    }
+
+                    // Convert from `range` to a byte range within `data` while
+                    // accounting for various offsets. Then create a
+                    // `CodeSectionReader` (which notably the payload does not
+                    // give us here) and recurse with that. This means that
+                    // users overridding `parse_code_section` always get that
+                    // function called.
+                    let section = match data.get(range.start - orig_offset..range.end - orig_offset)
+                    {
+                        Some(section) => section,
+                        None => return Err(Error::InvalidCodeSectionSize),
+                    };
+                    let reader = wasmparser::BinaryReader::new(section, range.start);
+                    let section = wasmparser::CodeSectionReader::new(reader)?;
+                    reencoder.parse_code_section(&mut codes, section)?;
                     module.section(&codes);
                 }
+
+                // Parsing of code section entries (function bodies) happens
+                // above during the handling of the code section. That means
+                // that we just skip all these payloads.
+                wasmparser::Payload::CodeSectionEntry(_) => {}
+
                 wasmparser::Payload::ModuleSection { .. }
                 | wasmparser::Payload::InstanceSection(_)
                 | wasmparser::Payload::CoreTypeSection(_)
@@ -795,8 +797,6 @@ pub mod utils {
                     handle_intersperse_section_hook(reencoder, module, &mut last_section, None)?;
                 }
             }
-
-            next_section = sections.next();
         }
 
         Ok(())
@@ -1413,7 +1413,20 @@ pub mod utils {
                 table_index,
                 offset_expr,
             } => elements.active(
-                table_index.map(|t| reencoder.table_index(t)),
+                // Inform the reencoder that a table index is being used even if
+                // it's not actually present here. That helps wasm-mutate for
+                // example which wants to track uses to know when it's ok to
+                // remove a table.
+                //
+                // If the table index is still zero though go ahead and pass
+                // `None` here since that implicitly references table 0. Note
+                // that this means that this does not round-trip the encoding of
+                // `Some(0)` since that reencodes to `None`, but that's seen as
+                // hopefully ok.
+                match reencoder.table_index(table_index.unwrap_or(0)) {
+                    0 => None,
+                    i => Some(i),
+                },
                 &reencoder.const_expr(offset_expr)?,
                 elems,
             ),

crates/wasm-mutate/Cargo.toml

@@ -24,3 +24,4 @@ anyhow = { workspace = true }
 wat = { path = "../wat" }
 wasmprinter = { path = "../wasmprinter" }
 env_logger = { workspace = true }
+wasmparser = { workspace = true, features = ['validate', 'features'] }

crates/wasm-mutate/src/error.rs

@@ -56,6 +56,16 @@ impl From<wasmparser::BinaryReaderError> for Error {
     }
 }
 
+impl From<wasm_encoder::reencode::Error<Error>> for Error {
+    fn from(e: wasm_encoder::reencode::Error<Error>) -> Self {
+        match e {
+            wasm_encoder::reencode::Error::ParseError(e) => Error::parse(e),
+            wasm_encoder::reencode::Error::UserError(e) => e,
+            other => Error::other(other.to_string()),
+        }
+    }
+}
+
 /// The kind of error.
 #[derive(thiserror::Error, Debug)]
 pub enum ErrorKind {