Start removing custom translation in wasm-mutate

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 {