Add a new wasm-tools component unbundle command (#1725)

* Add a new `wasm-tools component unbundle` command

This commit adds a new subcommand to `wasm-tools component` named
`unbundle` which is tasked with extracting core wasm modules out of a
component. This is intended to showcase an example implementation of
doing this and the main motivation is to extract core wasm modules out
of a component for runtimes such as Wasmtime. This enables sharing
machine code for WebAssembly modules between components that use the
same core wasm module.

This new subcommand is primarily built on `wasm_encoder::reencode` with
a simple size-based heuristic for which modules to extract.

* Test module output as well

Author: alexcrichton

README.md

@@ -148,6 +148,7 @@ that can be use programmatically as well:
 | `wasm-tools component embed` |  |  | Embed a `component-type` custom section in a core wasm binary |
 | `wasm-tools metadata show` |  [wasm-metadata] |  | Show name and producer metadata in a component or module |
 | `wasm-tools metadata add` |  |  | Add name or producer metadata to a component or module |
+| `wasm-tools metadata unbundle` |  |  | Extract core wasm modules from a component |
 | `wasm-tools addr2line` |  |  | Translate wasm offsets to filename/line numbers with DWARF |
 | `wasm-tools completion` |  |  | Generate shell completion scripts for `wasm-tools` |
 | `wasm-tools json-from-wast` |  |  | Convert a `*.wast` file into JSON commands |

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

@@ -104,6 +104,12 @@ impl ComponentNameSection {
         self.component_decls(INSTANCE_SORT, names)
     }
 
+    /// Appends a raw subsection with the given id and data.
+    pub fn raw(&mut self, id: u8, data: &[u8]) {
+        self.bytes.push(id);
+        data.encode(&mut self.bytes);
+    }
+
     fn component_decls(&mut self, kind: u8, names: &NameMap) {
         self.subsection_header(Subsection::Decls, 1 + names.size());
         self.bytes.push(kind);

crates/wasm-encoder/src/core/names.rs

@@ -180,6 +180,12 @@ impl NameSection {
         names.encode(&mut self.bytes);
     }
 
+    /// Appends a raw subsection with the given id and data.
+    pub fn raw(&mut self, id: u8, data: &[u8]) {
+        self.bytes.push(id);
+        data.encode(&mut self.bytes);
+    }
+
     fn subsection_header(&mut self, id: Subsection, len: usize) {
         self.bytes.push(id as u8);
         len.encode(&mut self.bytes);

crates/wasm-encoder/src/reencode.rs

@@ -45,6 +45,13 @@ pub trait Reencode {
         utils::type_index(self, ty)
     }
 
+    fn type_index_unpacked(
+        &mut self,
+        ty: wasmparser::UnpackedIndex,
+    ) -> Result<u32, Error<Self::Error>> {
+        utils::type_index_unpacked(self, ty)
+    }
+
     fn external_index(&mut self, kind: wasmparser::ExternalKind, index: u32) -> u32 {
         match kind {
             wasmparser::ExternalKind::Func => self.function_index(index),
@@ -463,6 +470,21 @@ pub trait Reencode {
     ) -> Result<(), Error<Self::Error>> {
         utils::parse_core_module(self, module, parser, data)
     }
+
+    fn custom_name_section(
+        &mut self,
+        section: wasmparser::NameSectionReader<'_>,
+    ) -> Result<crate::NameSection, Error<Self::Error>> {
+        utils::custom_name_section(self, section)
+    }
+
+    fn parse_custom_name_subsection(
+        &mut self,
+        names: &mut crate::NameSection,
+        section: wasmparser::Name<'_>,
+    ) -> Result<(), Error<Self::Error>> {
+        utils::parse_custom_name_subsection(self, names, section)
+    }
 }
 
 /// An error when re-encoding from `wasmparser` to `wasm-encoder`.
@@ -860,7 +882,14 @@ pub mod utils {
         module: &mut crate::Module,
         section: wasmparser::CustomSectionReader<'_>,
     ) -> Result<(), Error<T::Error>> {
-        module.section(&reencoder.custom_section(section));
+        match section.as_known() {
+            wasmparser::KnownCustom::Name(name) => {
+                module.section(&reencoder.custom_name_section(name)?);
+            }
+            _ => {
+                module.section(&reencoder.custom_section(section));
+            }
+        }
         Ok(())
     }
 
@@ -915,6 +944,15 @@ pub mod utils {
         ty
     }
 
+    pub fn type_index_unpacked<T: ?Sized + Reencode>(
+        reencoder: &mut T,
+        ty: wasmparser::UnpackedIndex,
+    ) -> Result<u32, Error<T::Error>> {
+        ty.as_module_index()
+            .map(|ty| reencoder.type_index(ty))
+            .ok_or(Error::CanonicalizedHeapTypeReference)
+    }
+
     pub fn tag_type<T: ?Sized + Reencode>(
         reencoder: &mut T,
         tag_ty: wasmparser::TagType,
@@ -986,11 +1024,7 @@ pub mod utils {
             is_final: sub_ty.is_final,
             supertype_idx: sub_ty
                 .supertype_idx
-                .map(|i| {
-                    i.as_module_index()
-                        .map(|ty| reencoder.type_index(ty))
-                        .ok_or(Error::CanonicalizedHeapTypeReference)
-                })
+                .map(|i| reencoder.type_index_unpacked(i.unpack()))
                 .transpose()?,
             composite_type: reencoder.composite_type(sub_ty.composite_type)?,
         })
@@ -1101,11 +1135,9 @@ pub mod utils {
         heap_type: wasmparser::HeapType,
     ) -> Result<crate::HeapType, Error<T::Error>> {
         Ok(match heap_type {
-            wasmparser::HeapType::Concrete(i) => crate::HeapType::Concrete(
-                i.as_module_index()
-                    .map(|ty| reencoder.type_index(ty))
-                    .ok_or(Error::CanonicalizedHeapTypeReference)?,
-            ),
+            wasmparser::HeapType::Concrete(i) => {
+                crate::HeapType::Concrete(reencoder.type_index_unpacked(i)?)
+            }
             wasmparser::HeapType::Abstract { shared, ty } => crate::HeapType::Abstract {
                 shared,
                 ty: reencoder.abstract_heap_type(ty),
@@ -1580,6 +1612,90 @@ pub mod utils {
         module.section(&crate::RawSection { id, data: contents });
         Ok(())
     }
+
+    pub fn custom_name_section<T: ?Sized + Reencode>(
+        reencoder: &mut T,
+        section: wasmparser::NameSectionReader<'_>,
+    ) -> Result<crate::NameSection, Error<T::Error>> {
+        let mut ret = crate::NameSection::new();
+        for subsection in section {
+            reencoder.parse_custom_name_subsection(&mut ret, subsection?)?;
+        }
+        Ok(ret)
+    }
+
+    pub fn parse_custom_name_subsection<T: ?Sized + Reencode>(
+        reencoder: &mut T,
+        names: &mut crate::NameSection,
+        section: wasmparser::Name<'_>,
+    ) -> Result<(), Error<T::Error>> {
+        match section {
+            wasmparser::Name::Module { name, .. } => {
+                names.module(name);
+            }
+            wasmparser::Name::Function(map) => {
+                names.functions(&name_map(map, |i| reencoder.function_index(i))?);
+            }
+            wasmparser::Name::Type(map) => {
+                names.types(&name_map(map, |i| reencoder.type_index(i))?);
+            }
+            wasmparser::Name::Local(map) => {
+                names.locals(&indirect_name_map(map, |i| reencoder.function_index(i))?);
+            }
+            wasmparser::Name::Label(map) => {
+                names.labels(&indirect_name_map(map, |i| reencoder.function_index(i))?);
+            }
+            wasmparser::Name::Table(map) => {
+                names.tables(&name_map(map, |i| reencoder.table_index(i))?);
+            }
+            wasmparser::Name::Memory(map) => {
+                names.memories(&name_map(map, |i| reencoder.memory_index(i))?);
+            }
+            wasmparser::Name::Global(map) => {
+                names.globals(&name_map(map, |i| reencoder.global_index(i))?);
+            }
+            wasmparser::Name::Element(map) => {
+                names.elements(&name_map(map, |i| reencoder.element_index(i))?);
+            }
+            wasmparser::Name::Data(map) => {
+                names.data(&name_map(map, |i| reencoder.data_index(i))?);
+            }
+            wasmparser::Name::Tag(map) => {
+                names.tags(&name_map(map, |i| reencoder.tag_index(i))?);
+            }
+            wasmparser::Name::Field(map) => {
+                names.fields(&indirect_name_map(map, |i| reencoder.type_index(i))?);
+            }
+            wasmparser::Name::Unknown { ty, data, .. } => {
+                names.raw(ty, data);
+            }
+        }
+        Ok(())
+    }
+
+    pub fn name_map(
+        map: wasmparser::NameMap<'_>,
+        mut map_index: impl FnMut(u32) -> u32,
+    ) -> wasmparser::Result<crate::NameMap> {
+        let mut ret = crate::NameMap::new();
+        for naming in map {
+            let naming = naming?;
+            ret.append(map_index(naming.index), naming.name);
+        }
+        Ok(ret)
+    }
+
+    pub fn indirect_name_map(
+        map: wasmparser::IndirectNameMap<'_>,
+        mut map_index: impl FnMut(u32) -> u32,
+    ) -> wasmparser::Result<crate::IndirectNameMap> {
+        let mut ret = crate::IndirectNameMap::new();
+        for naming in map {
+            let naming = naming?;
+            ret.append(map_index(naming.index), &name_map(naming.names, |i| i)?);
+        }
+        Ok(ret)
+    }
 }
 
 impl From<wasmparser::MemArg> for crate::MemArg {