Recognize _initialize in wasm-tools component new (#1747)

This commit implements recognition of the `_initialize` function from
WASIp1 in the componentization process of `wasm-tools component new`.
This additionally corresponds to the same function in the proposed
[BuildTargets.md](WebAssembly/component-model#378).
This is implemented by having a small core wasm module which is just an
import and a `start` section get instantiated at the end of a component
to run `_initialize` before all other exports.

Author: alexcrichton

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

@@ -206,6 +206,20 @@ impl ComponentBuilder {
         })
     }
 
+    /// Creates an alias to a previous core instance's exported item.
+    ///
+    /// The `instance` provided is the instance to access and the `name` is the
+    /// item to access.
+    ///
+    /// Returns the index of the new item defined.
+    pub fn alias_core_export(&mut self, instance: u32, name: &str, kind: ExportKind) -> u32 {
+        self.alias(Alias::CoreInstanceExport {
+            instance,
+            kind,
+            name,
+        })
+    }
+
     fn inc_kind(&mut self, kind: ComponentExportKind) -> u32 {
         match kind {
             ComponentExportKind::Func => inc(&mut self.funcs),

crates/wit-component/src/encoding.rs

@@ -631,6 +631,8 @@ impl<'a> EncodingState<'a> {
             self.instantiate_adapter_module(&shims, name, adapter);
         }
 
+        self.encode_initialize_with_start()?;
+
         Ok(())
     }
 
@@ -1689,6 +1691,60 @@ impl<'a> EncodingState<'a> {
         });
         self.adapter_import_reallocs.insert(name, realloc);
     }
+
+    /// Generates component bits that are responsible for executing
+    /// `_initialize`, if found, in the original component.
+    ///
+    /// The `_initialize` function was a part of WASIp1 where it generally is
+    /// intended to run after imports and memory and such are all "hooked up"
+    /// and performs other various initialization tasks. This is additionally
+    /// specified in https://github.com/WebAssembly/component-model/pull/378
+    /// to be part of the component model lowerings as well.
+    ///
+    /// This implements this functionality by encoding a core module that
+    /// imports a function and then registers a `start` section with that
+    /// imported function. This is all encoded after the
+    /// imports/lowerings/tables/etc are all filled in above meaning that this
+    /// is the last piece to run. That means that when this is running
+    /// everything should be hooked up for all imported functions to work.
+    ///
+    /// Note that at this time `_initialize` is only detected in the "main
+    /// module", not adapters/libraries.
+    fn encode_initialize_with_start(&mut self) -> Result<()> {
+        let initialize = match self.info.info.initialize {
+            Some(name) => name,
+            // If this core module didn't have `_initialize` or similar, then
+            // there's nothing to do here.
+            None => return Ok(()),
+        };
+        let initialize_index = self.component.alias_core_export(
+            self.instance_index.unwrap(),
+            initialize,
+            ExportKind::Func,
+        );
+        let mut shim = Module::default();
+        let mut section = TypeSection::new();
+        section.function([], []);
+        shim.section(&section);
+        let mut section = ImportSection::new();
+        section.import("", "", EntityType::Function(0));
+        shim.section(&section);
+        shim.section(&StartSection { function_index: 0 });
+
+        // Declare the core module within the component, create a dummy core
+        // instance with one export of our `_initialize` function, and then use
+        // that to instantiate the module we emit to run the `start` function in
+        // core wasm to run `_initialize`.
+        let shim_module_index = self.component.core_module(&shim);
+        let shim_args_instance_index =
+            self.component
+                .core_instantiate_exports([("", ExportKind::Func, initialize_index)]);
+        self.component.core_instantiate(
+            shim_module_index,
+            [("", ModuleArg::Instance(shim_args_instance_index))],
+        );
+        Ok(())
+    }
 }
 
 /// A list of "shims" which start out during the component instantiation process

crates/wit-component/src/validation.rs

@@ -100,6 +100,10 @@ pub struct ValidatedModule<'a> {
     /// Post-return functions annotated with `cabi_post_*` in their function
     /// name.
     pub post_returns: IndexSet<String>,
+
+    /// Exported function like `_initialize` which needs to be run after
+    /// everything else has been instantiated.
+    pub initialize: Option<&'a str>,
 }
 
 #[derive(Default)]
@@ -146,6 +150,7 @@ pub fn validate_module<'a>(
         metadata: &metadata.metadata,
         required_resource_funcs: Default::default(),
         post_returns: Default::default(),
+        initialize: None,
     };
 
     for payload in Parser::new(0).parse_all(bytes) {
@@ -194,7 +199,11 @@ pub fn validate_module<'a>(
                                 }
                             }
 
-                            assert!(export_funcs.insert(export.name, export.index).is_none())
+                            if export.name == "_initialize" {
+                                ret.initialize = Some(export.name);
+                            } else {
+                                assert!(export_funcs.insert(export.name, export.index).is_none())
+                            }
                         }
                         ExternalKind::Memory => {
                             if export.name == "memory" {

crates/wit-component/tests/components/initialize/component.wat

@@ -0,0 +1,38 @@
+(component
+  (core module (;0;)
+    (type (;0;) (func))
+    (func (;0;) (type 0)
+      unreachable
+    )
+    (func (;1;) (type 0)
+      unreachable
+    )
+    (export "a" (func 0))
+    (export "_initialize" (func 1))
+    (@producers
+      (processed-by "wit-component" "$CARGO_PKG_VERSION")
+      (processed-by "my-fake-bindgen" "123.45")
+    )
+  )
+  (core instance (;0;) (instantiate 0))
+  (alias core export 0 "_initialize" (core func (;0;)))
+  (core module (;1;)
+    (type (;0;) (func))
+    (import "" "" (func (;0;) (type 0)))
+    (start 0)
+  )
+  (core instance (;1;)
+    (export "" (func 0))
+  )
+  (core instance (;2;) (instantiate 1
+      (with "" (instance 1))
+    )
+  )
+  (type (;0;) (func))
+  (alias core export 0 "a" (core func (;1;)))
+  (func (;0;) (type 0) (canon lift (core func 1)))
+  (export (;1;) "a" (func 0))
+  (@producers
+    (processed-by "wit-component" "$CARGO_PKG_VERSION")
+  )
+)

crates/wit-component/tests/components/initialize/component.wit.print

@@ -0,0 +1,5 @@
+package root:component;
+
+world root {
+  export a: func();
+}

crates/wit-component/tests/components/initialize/module.wat

@@ -0,0 +1,6 @@
+;; This test ensures that the `_initialize` function is hooked up to execute
+;; when the component is instantiated.
+(module
+  (func (export "a") unreachable)
+  (func (export "_initialize") unreachable)
+)

crates/wit-component/tests/components/initialize/module.wit

@@ -0,0 +1,5 @@
+package foo:foo;
+
+world module {
+  export a: func();
+}