wizer: Partially support nested components (#11898)

* wizer: Partially support nested components

Support nested components insofar as they don't have access to modules
at this time. This helps `wasmtime wizer` support wasm-tools-generated
components which use nested components for type-ascribing exports.
Otherwise though nested components are not allowed to defined or
instantiate core modules as that would need more implementation in Wizer
to support.

* Two-level test too

* Add test.wasm to wizer's gitignore

* Fix test

Author: alexcrichton

crates/wizer/.gitignore

@@ -5,4 +5,5 @@ npm/wizer-darwin-x64/
 npm/wizer-linux-x64/
 npm/wizer-win32-x64/
 npm/wizer-linux-s390x/
-npm/wizer-linux-arm64/
+npm/wizer-linux-arm64/
+test.wasm

crates/wizer/src/component/parse.rs

@@ -13,29 +13,45 @@ use wasmparser::{
 /// parsing phase which in theory could be lifted in the future but serve as
 /// simplifying assumptions for now:
 ///
-/// * Nested components are not supported.
+/// * Nested components with modules are not supported.
 /// * Imported modules or components are not supported.
 /// * Instantiating a module twice is not supported.
 /// * Component-level start functions are not supported.
 ///
 /// Some of these restrictions are likely to be loosened over time, however.
 pub(crate) fn parse<'a>(full_wasm: &'a [u8]) -> anyhow::Result<ComponentContext<'a>> {
     let mut component = ComponentContext::default();
-    let mut parser = Parser::new(0).parse_all(full_wasm);
+    let parser = Parser::new(0).parse_all(full_wasm);
+    parse_into(Some(&mut component), full_wasm, parser)?;
+    Ok(component)
+}
 
-    while let Some(payload) = parser.next() {
+fn parse_into<'a>(
+    mut cx: Option<&mut ComponentContext<'a>>,
+    full_wasm: &'a [u8],
+    mut iter: impl Iterator<Item = wasmparser::Result<Payload<'a>>>,
+) -> anyhow::Result<()> {
+    let mut stack = Vec::new();
+    while let Some(payload) = iter.next() {
         let payload = payload?;
 
         match &payload {
             // Module sections get parsed with wizer's core wasm support.
-            Payload::ModuleSection { .. } => {
-                let info = crate::parse::parse_with(&full_wasm, &mut parser)?;
-                component.push_module_section(info);
-            }
+            Payload::ModuleSection { .. } => match &mut cx {
+                Some(component) => {
+                    let info = crate::parse::parse_with(&full_wasm, &mut iter)?;
+                    component.push_module_section(info);
+                }
+                None => {
+                    bail!("nested components with modules not currently supported");
+                }
+            },
 
             // All other sections get pushed raw as-is into the component.
             _ => {
-                if let Some((id, range)) = payload.as_section() {
+                if let Some((id, range)) = payload.as_section()
+                    && let Some(component) = &mut cx
+                {
                     component.push_raw_section(wasm_encoder::RawSection {
                         id,
                         data: &full_wasm[range],
@@ -56,43 +72,61 @@ pub(crate) fn parse<'a>(full_wasm: &'a [u8]) -> anyhow::Result<ComponentContext<
             }
 
             Payload::ComponentSection { .. } => {
-                bail!("wizer does not currently support nested components");
+                stack.push(cx.take());
             }
 
-            Payload::InstanceSection(reader) => {
-                for instance in reader {
-                    let instance_index = component.inc_core_instances();
+            Payload::End(_) => {
+                if stack.len() > 0 {
+                    cx = stack.pop().unwrap();
+                }
+            }
 
-                    if let Instance::Instantiate { module_index, .. } = instance? {
-                        match component.core_instantiations.entry(module_index) {
-                            Entry::Vacant(entry) => {
-                                entry.insert(instance_index);
+            Payload::InstanceSection(reader) => {
+                if let Some(component) = &mut cx {
+                    for instance in reader {
+                        let instance_index = component.inc_core_instances();
+
+                        if let Instance::Instantiate { module_index, .. } = instance? {
+                            match component.core_instantiations.entry(module_index) {
+                                Entry::Vacant(entry) => {
+                                    entry.insert(instance_index);
+                                }
+                                Entry::Occupied(_) => {
+                                    bail!("modules may be instantiated at most once")
+                                }
                             }
-                            Entry::Occupied(_) => bail!("modules may be instantiated at most once"),
                         }
                     }
                 }
             }
             Payload::ComponentInstanceSection(reader) => {
-                for _ in reader {
-                    component.inc_instances();
+                if let Some(component) = &mut cx {
+                    for _ in reader {
+                        component.inc_instances();
+                    }
                 }
             }
 
             Payload::ComponentAliasSection(reader) => {
                 for alias in reader {
                     match alias? {
                         ComponentAlias::CoreInstanceExport { kind, .. } => {
-                            component.inc_core(kind);
+                            if let Some(component) = &mut cx {
+                                component.inc_core(kind);
+                            }
                         }
                         ComponentAlias::InstanceExport { kind, .. } => {
                             validate_item_kind(kind, "aliases")?;
-                            component.inc(kind);
+                            if let Some(component) = &mut cx {
+                                component.inc(kind);
+                            }
                         }
                         ComponentAlias::Outer { kind, .. } => match kind {
                             ComponentOuterAliasKind::CoreType => {}
                             ComponentOuterAliasKind::Type => {
-                                component.inc_types();
+                                if let Some(component) = &mut cx {
+                                    component.inc_types();
+                                }
                             }
                             ComponentOuterAliasKind::CoreModule => {
                                 bail!("wizer does not currently support module aliases");
@@ -109,10 +143,14 @@ pub(crate) fn parse<'a>(full_wasm: &'a [u8]) -> anyhow::Result<ComponentContext<
                 for function in reader {
                     match function? {
                         CanonicalFunction::Lift { .. } => {
-                            component.inc_funcs();
+                            if let Some(component) = &mut cx {
+                                component.inc_funcs();
+                            }
                         }
                         _ => {
-                            component.inc_core_funcs();
+                            if let Some(component) = &mut cx {
+                                component.inc_core_funcs();
+                            }
                         }
                     }
                 }
@@ -122,21 +160,27 @@ pub(crate) fn parse<'a>(full_wasm: &'a [u8]) -> anyhow::Result<ComponentContext<
                 for import in reader {
                     let kind = import?.ty.kind();
                     validate_item_kind(kind, "imports")?;
-                    component.inc(kind);
+                    if let Some(component) = &mut cx {
+                        component.inc(kind);
+                    }
                 }
             }
 
             Payload::ComponentExportSection(reader) => {
                 for export in reader {
                     let kind = export?.kind;
                     validate_item_kind(kind, "exports")?;
-                    component.inc(kind);
+                    if let Some(component) = &mut cx {
+                        component.inc(kind);
+                    }
                 }
             }
 
             Payload::ComponentTypeSection(reader) => {
                 for _ in reader {
-                    component.inc_types();
+                    if let Some(component) = &mut cx {
+                        component.inc_types();
+                    }
                 }
             }
 
@@ -151,7 +195,7 @@ pub(crate) fn parse<'a>(full_wasm: &'a [u8]) -> anyhow::Result<ComponentContext<
         }
     }
 
-    Ok(component)
+    Ok(())
 }
 
 fn validate_item_kind(kind: ComponentExternalKind, msg: &str) -> Result<()> {

crates/wizer/src/component/rewrite.rs

@@ -107,7 +107,11 @@ impl Reencoder<'_> {
             // indices in case the index is higher than the one that we're
             // removing.
             id if id == wasm_encoder::ComponentSectionId::Instance as u8 => {
-                self.rewrite(encoder, section.data, Self::parse_instance_section);
+                self.rewrite(
+                    encoder,
+                    section.data,
+                    Self::parse_component_instance_section,
+                );
             }
             id if id == wasm_encoder::ComponentSectionId::Alias as u8 => {
                 self.rewrite(encoder, section.data, Self::parse_component_alias_section);

crates/wizer/test.wasm

@@ -7,6 +7,10 @@ fn fail_wizening(msg: &str, wasm: &[u8]) -> Result<()> {
     let _ = env_logger::try_init();
 
     let wasm = wat::parse_bytes(wasm)?;
+    log::debug!(
+        "testing wizening failure for wasm:\n{}",
+        wasmprinter::print_bytes(&wasm)?
+    );
     match Wizer::new().instrument_component(&wasm) {
         Ok(_) => bail!("expected wizening to fail"),
         Err(e) => {
@@ -29,9 +33,40 @@ fn unsupported_constructs() -> Result<()> {
     )?;
 
     fail_wizening(
-        "does not currently support nested components",
+        "nested components with modules not currently supported",
+        br#"(component
+            (component (core module))
+        )"#,
+    )?;
+    fail_wizening(
+        "nested components with modules not currently supported",
         br#"(component
             (component)
+            (component (core module))
+        )"#,
+    )?;
+    fail_wizening(
+        "wizer does not currently support module imports",
+        br#"(component
+            (component (import "x" (core module)))
+        )"#,
+    )?;
+    fail_wizening(
+        "wizer does not currently support module aliases",
+        br#"(component
+            (core module $a)
+            (component
+                (core instance (instantiate $a))
+            )
+        )"#,
+    )?;
+    fail_wizening(
+        "wizer does not currently support component aliases",
+        br#"(component
+            (component $a)
+            (component
+                (instance (instantiate $a))
+            )
         )"#,
     )?;
 
@@ -308,6 +343,44 @@ async fn snapshot_memory() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn nested_components() -> Result<()> {
+    wizen_and_run_wasm(
+        42,
+        r#"(component
+            (component $a)
+            (instance (instantiate $a))
+            (instance (export "hi") (instantiate $a))
+
+            (component $b
+                (type $t string)
+                (import "x" (type (eq $t)))
+                (component $a)
+                (instance (instantiate $a))
+                (instance (export "hi") (instantiate $a))
+            )
+            (type $x string)
+            (instance (instantiate $b
+                (with "x" (type $x))
+            ))
+            (instance (export "hi2") (instantiate $b
+                (with "x" (type $x))
+            ))
+
+            (core module $m
+                (func (export "init"))
+                (func (export "run") (result i32) i32.const 42)
+            )
+            (core instance $i (instantiate $m))
+            (func (export "run") (result u32) (canon lift (core func $i "run")))
+            (func (export "wizer-initialize") (canon lift (core func $i "init")))
+        )"#,
+    )
+    .await?;
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn multiple_modules() -> Result<()> {
     wizen_and_run_wasm(