Implement lowered-then-lifted functions (#4327)

* Implement lowered-then-lifted functions

This commit is a few features bundled into one, culminating in the
implementation of lowered-then-lifted functions for the component model.
It's probably not going to be used all that often but this is possible
within a valid component so Wasmtime needs to do something relatively
reasonable. The main things implemented in this commit are:

* Component instances are now assigned a `RuntimeComponentInstanceIndex`
  to differentiate each one. This will be used in the future to detect
  fusion (one instance lowering a function from another instance). For
  now it's used to allocate separate `VMComponentFlags` for each
  internal component instance.

* The `CoreExport<FuncIndex>` of lowered functions was changed to a
  `CoreDef` since technically a lowered function can use another lowered
  function as the callee. This ended up being not too difficult to plumb
  through as everything else was already in place.

* A need arose to compile host-to-wasm trampolines which weren't already
  present. Currently wasm in a component is always entered through a
  host-to-wasm trampoline but core wasm modules are the source of all
  the trampolines. In the case of a lowered-then-lifted function there
  may not actually be any core wasm modules, so component objects now
  contain necessary trampolines not otherwise provided by the core wasm
  objects. This feature required splitting a new function into the
  `Compiler` trait for creating a host-to-wasm trampoline. After doing
  this core wasm compilation was also updated to leverage this which
  further enabled compiling trampolines in parallel as opposed to the
  previous synchronous compilation.

* Review comments

crates/cranelift/src/compiler.rs

@@ -288,11 +288,19 @@ impl wasmtime_environ::Compiler for Compiler {
         }))
     }
 
+    fn compile_host_to_wasm_trampoline(
+        &self,
+        ty: &WasmFuncType,
+    ) -> Result<Box<dyn Any + Send>, CompileError> {
+        self.host_to_wasm_trampoline(ty)
+            .map(|x| Box::new(x) as Box<_>)
+    }
+
     fn emit_obj(
         &self,
         translation: &ModuleTranslation,
-        types: &ModuleTypes,
         funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
+        compiled_trampolines: Vec<Box<dyn Any + Send>>,
         tunables: &Tunables,
         obj: &mut Object<'static>,
     ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)> {
@@ -300,6 +308,10 @@ impl wasmtime_environ::Compiler for Compiler {
             .into_iter()
             .map(|(_i, f)| *f.downcast().unwrap())
             .collect();
+        let compiled_trampolines: Vec<CompiledFunction> = compiled_trampolines
+            .into_iter()
+            .map(|f| *f.downcast().unwrap())
+            .collect();
 
         let mut builder = ModuleTextBuilder::new(obj, &translation.module, &*self.isa);
         if self.linkopts.force_jump_veneers {
@@ -307,11 +319,6 @@ impl wasmtime_environ::Compiler for Compiler {
         }
         let mut addrs = AddressMapSection::default();
         let mut traps = TrapEncodingBuilder::default();
-        let compiled_trampolines = translation
-            .exported_signatures
-            .iter()
-            .map(|i| self.host_to_wasm_trampoline(&types[*i]))
-            .collect::<Result<Vec<_>, _>>()?;
 
         let mut func_starts = Vec::with_capacity(funcs.len());
         for (i, func) in funcs.iter() {
@@ -325,6 +332,10 @@ impl wasmtime_environ::Compiler for Compiler {
         }
 
         // Build trampolines for every signature that can be used by this module.
+        assert_eq!(
+            translation.exported_signatures.len(),
+            compiled_trampolines.len()
+        );
         let mut trampolines = Vec::with_capacity(translation.exported_signatures.len());
         for (i, func) in translation
             .exported_signatures

crates/cranelift/src/compiler/component.rs

@@ -10,9 +10,9 @@ use object::write::Object;
 use std::any::Any;
 use wasmtime_environ::component::{
     CanonicalOptions, Component, ComponentCompiler, ComponentTypes, LowerImport, LoweredIndex,
-    TrampolineInfo, VMComponentOffsets,
+    LoweringInfo, VMComponentOffsets,
 };
-use wasmtime_environ::PrimaryMap;
+use wasmtime_environ::{PrimaryMap, SignatureIndex, Trampoline};
 
 impl ComponentCompiler for Compiler {
     fn compile_lowered_trampoline(
@@ -52,6 +52,7 @@ impl ComponentCompiler for Compiler {
         let mut host_sig = ir::Signature::new(crate::wasmtime_call_conv(isa));
 
         let CanonicalOptions {
+            instance,
             memory,
             realloc,
             post_return,
@@ -71,6 +72,14 @@ impl ComponentCompiler for Compiler {
             i32::try_from(offsets.lowering_data(lowering.index)).unwrap(),
         ));
 
+        // flags: *mut VMComponentFlags
+        host_sig.params.push(ir::AbiParam::new(pointer_type));
+        callee_args.push(
+            builder
+                .ins()
+                .iadd_imm(vmctx, i64::from(offsets.flags(instance))),
+        );
+
         // memory: *mut VMMemoryDefinition
         host_sig.params.push(ir::AbiParam::new(pointer_type));
         callee_args.push(match memory {
@@ -145,32 +154,42 @@ impl ComponentCompiler for Compiler {
 
     fn emit_obj(
         &self,
-        trampolines: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        lowerings: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        trampolines: Vec<(SignatureIndex, Box<dyn Any + Send>)>,
         obj: &mut Object<'static>,
-    ) -> Result<PrimaryMap<LoweredIndex, TrampolineInfo>> {
+    ) -> Result<(PrimaryMap<LoweredIndex, LoweringInfo>, Vec<Trampoline>)> {
-        let trampolines: PrimaryMap<LoweredIndex, CompiledFunction> = trampolines
+        let lowerings: PrimaryMap<LoweredIndex, CompiledFunction> = lowerings
             .into_iter()
             .map(|(_, f)| *f.downcast().unwrap())
             .collect();
+        let trampolines: Vec<(SignatureIndex, CompiledFunction)> = trampolines
+            .into_iter()
+            .map(|(i, f)| (i, *f.downcast().unwrap()))
+            .collect();
+
         let module = Default::default();
         let mut text = ModuleTextBuilder::new(obj, &module, &*self.isa);
         let mut ret = PrimaryMap::new();
-        for (idx, trampoline) in trampolines.iter() {
+        for (idx, lowering) in lowerings.iter() {
             let (_symbol, range) = text.append_func(
                 false,
-                format!("_wasm_component_host_trampoline{}", idx.as_u32()).into_bytes(),
+                format!("_wasm_component_lowering_trampoline{}", idx.as_u32()).into_bytes(),
-                &trampoline,
+                &lowering,
             );
 
-            let i = ret.push(TrampolineInfo {
+            let i = ret.push(LoweringInfo {
                 start: u32::try_from(range.start).unwrap(),
                 length: u32::try_from(range.end - range.start).unwrap(),
             });
             assert_eq!(i, idx);
         }
+        let ret_trampolines = trampolines
+            .iter()
+            .map(|(i, func)| text.trampoline(*i, func))
+            .collect();
 
         text.finish()?;
 
-        Ok(ret)
+        Ok((ret, ret_trampolines))
     }
 }

crates/environ/src/compilation.rs

@@ -148,12 +148,21 @@ pub trait Compiler: Send + Sync {
         types: &ModuleTypes,
     ) -> Result<Box<dyn Any + Send>, CompileError>;
 
+    /// Creates a function of type `VMTrampoline` which will then call the
+    /// function pointer argument which has the `ty` type provided.
+    fn compile_host_to_wasm_trampoline(
+        &self,
+        ty: &WasmFuncType,
+    ) -> Result<Box<dyn Any + Send>, CompileError>;
+
     /// Collects the results of compilation into an in-memory object.
     ///
     /// This function will receive the same `Box<dyn Ayn>` produced as part of
     /// `compile_function`, as well as the general compilation environment with
-    /// the translation/types. This method is expected to populate information
+    /// the translation. THe `trampolines` argument is generated by
-    /// in the object file such as:
+    /// `compile_host_to_wasm_trampoline` for each of
+    /// `module.exported_signatures`. This method is expected to populate
+    /// information in the object file such as:
     ///
     /// * Compiled code in a `.text` section
     /// * Unwind information in Wasmtime-specific sections
@@ -163,11 +172,14 @@ pub trait Compiler: Send + Sync {
     ///
     /// The final result of compilation will contain more sections inserted by
     /// the compiler-agnostic runtime.
+    ///
+    /// This function returns information about the compiled functions (where
+    /// they are in the text section) along with where trampolines are located.
     fn emit_obj(
         &self,
         module: &ModuleTranslation,
-        types: &ModuleTypes,
         funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
+        trampolines: Vec<Box<dyn Any + Send>>,
         tunables: &Tunables,
         obj: &mut Object<'static>,
     ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)>;

crates/environ/src/component/compiler.rs

@@ -1,5 +1,5 @@
 use crate::component::{Component, ComponentTypes, LowerImport, LoweredIndex};
-use crate::PrimaryMap;
+use crate::{PrimaryMap, SignatureIndex, Trampoline};
 use anyhow::Result;
 use object::write::Object;
 use serde::{Deserialize, Serialize};
@@ -8,7 +8,7 @@ use std::any::Any;
 /// Description of where a trampoline is located in the text section of a
 /// compiled image.
 #[derive(Serialize, Deserialize)]
-pub struct TrampolineInfo {
+pub struct LoweringInfo {
     /// The byte offset from the start of the text section where this trampoline
     /// starts.
     pub start: u32,
@@ -42,8 +42,8 @@ pub trait ComponentCompiler: Send + Sync {
         types: &ComponentTypes,
     ) -> Result<Box<dyn Any + Send>>;
 
-    /// Emits the `trampolines` specified into the in-progress ELF object
+    /// Emits the `lowerings` and `trampolines` specified into the in-progress
-    /// specified by `obj`.
+    /// ELF object specified by `obj`.
     ///
     /// Returns a map of trampoline information for where to find them all in
     /// the text section.
@@ -52,7 +52,8 @@ pub trait ComponentCompiler: Send + Sync {
     /// trampolines as necessary.
     fn emit_obj(
         &self,
-        trampolines: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        lowerings: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        tramplines: Vec<(SignatureIndex, Box<dyn Any + Send>)>,
         obj: &mut Object<'static>,
-    ) -> Result<PrimaryMap<LoweredIndex, TrampolineInfo>>;
+    ) -> Result<(PrimaryMap<LoweredIndex, LoweringInfo>, Vec<Trampoline>)>;
 }

crates/environ/src/component/info.rs

@@ -114,6 +114,10 @@ pub struct Component {
     /// when instantiating this component.
     pub num_runtime_instances: u32,
 
+    /// Same as `num_runtime_instances`, but for `RuntimeComponentInstanceIndex`
+    /// instead.
+    pub num_runtime_component_instances: u32,
+
     /// The number of runtime memories (maximum `RuntimeMemoryIndex`) needed to
     /// instantiate this component.
     ///
@@ -355,7 +359,7 @@ pub enum Export {
         /// The component function type of the function being created.
         ty: TypeFuncIndex,
         /// Which core WebAssembly export is being lifted.
-        func: CoreExport<FuncIndex>,
+        func: CoreDef,
         /// Any options, if present, associated with this lifting.
         options: CanonicalOptions,
     },
@@ -369,6 +373,9 @@ pub enum Export {
 /// Canonical ABI options associated with a lifted or lowered function.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CanonicalOptions {
+    /// The component instance that this bundle was associated with.
+    pub instance: RuntimeComponentInstanceIndex,
+
     /// The encoding used for strings.
     pub string_encoding: StringEncoding,
 
@@ -382,17 +389,6 @@ pub struct CanonicalOptions {
     pub post_return: Option<RuntimePostReturnIndex>,
 }
 
-impl Default for CanonicalOptions {
-    fn default() -> CanonicalOptions {
-        CanonicalOptions {
-            string_encoding: StringEncoding::Utf8,
-            memory: None,
-            realloc: None,
-            post_return: None,
-        }
-    }
-}
-
 /// Possible encodings of strings within the component model.
 //
 // Note that the `repr(u8)` is load-bearing here since this is used in an

crates/environ/src/component/translate/inline.rs

@@ -104,7 +104,14 @@ pub(super) fn run(
     // initial frame. When the inliner finishes it will return the exports of
     // the root frame which are then used for recording the exports of the
     // component.
-    let mut frames = vec![InlinerFrame::new(result, ComponentClosure::default(), args)];
+    let index = RuntimeComponentInstanceIndex::from_u32(0);
+    inliner.result.num_runtime_component_instances += 1;
+    let mut frames = vec![InlinerFrame::new(
+        index,
+        result,
+        ComponentClosure::default(),
+        args,
+    )];
     let exports = inliner.run(&mut frames)?;
     assert!(frames.is_empty());
 
@@ -195,6 +202,8 @@ struct Inliner<'a> {
 /// inliner frames are stored on the heap to avoid recursion based on user
 /// input.
 struct InlinerFrame<'a> {
+    instance: RuntimeComponentInstanceIndex,
+
     /// The remaining initializers to process when instantiating this component.
     initializers: std::slice::Iter<'a, LocalInitializer<'a>>,
 
@@ -312,7 +321,7 @@ enum ComponentFuncDef<'a> {
     /// A core wasm function was lifted into a component function.
     Lifted {
         ty: TypeFuncIndex,
-        func: CoreExport<FuncIndex>,
+        func: CoreDef,
         options: CanonicalOptions,
     },
 }
@@ -509,19 +518,7 @@ impl<'a> Inliner<'a> {
                 let options = self.canonical_options(frame, options);
                 frame.component_funcs.push(ComponentFuncDef::Lifted {
                     ty: *ty,
-                    func: match frame.funcs[*func].clone() {
+                    func: frame.funcs[*func].clone(),
-                        CoreDef::Export(e) => e.map_index(|i| match i {
-                            EntityIndex::Function(i) => i,
-                            _ => unreachable!("not possible in valid components"),
-                        }),
-
-                        // TODO: lifting a lowered function only happens within
-                        // one component so this runs afoul of "someone needs to
-                        // really closely interpret the may_{enter,leave} flags"
-                        // in the component model spec. That has not currently
-                        // been done so this is left to panic.
-                        CoreDef::Lowered(_) => unimplemented!("lifting a lowered function"),
-                    },
                     options,
                 });
             }
@@ -618,7 +615,12 @@ impl<'a> Inliner<'a> {
             // stack.
             ComponentInstantiate(component, args) => {
                 let component: &ComponentDef<'a> = &frame.components[*component];
+                let index = RuntimeComponentInstanceIndex::from_u32(
+                    self.result.num_runtime_component_instances,
+                );
+                self.result.num_runtime_component_instances += 1;
                 let frame = InlinerFrame::new(
+                    index,
                     &self.nested_components[component.index],
                     component.closure.clone(),
                     args.iter()
@@ -872,6 +874,7 @@ impl<'a> Inliner<'a> {
                 })
         });
         CanonicalOptions {
+            instance: frame.instance,
             string_encoding: options.string_encoding,
             memory,
             realloc,
@@ -882,6 +885,7 @@ impl<'a> Inliner<'a> {
 
 impl<'a> InlinerFrame<'a> {
     fn new(
+        instance: RuntimeComponentInstanceIndex,
         translation: &'a Translation<'a>,
         closure: ComponentClosure<'a>,
         args: HashMap<&'a str, ComponentItemDef<'a>>,
@@ -891,6 +895,7 @@ impl<'a> InlinerFrame<'a> {
         // all the maps below. Given that doing such would be wordy and compile
         // time is otherwise not super crucial it's not done at this time.
         InlinerFrame {
+            instance,
             translation,
             closure,
             args,

crates/environ/src/component/types.rs

@@ -122,6 +122,9 @@ indices! {
     /// refer back to previously created instances for exports and such.
     pub struct RuntimeInstanceIndex(u32);
 
+    /// Same as `RuntimeInstanceIndex` but tracks component instances instead.
+    pub struct RuntimeComponentInstanceIndex(u32);
+
     /// Used to index imports into a `Component`
     ///
     /// This does not correspond to anything in the binary format for the

crates/environ/src/component/vmcomponent_offsets.rs

@@ -2,8 +2,8 @@
 //
 // struct VMComponentContext {
 //      magic: u32,
-//      flags: u8,
 //      store: *mut dyn Store,
+//      flags: [VMComponentFlags; component.num_runtime_component_instances],
 //      lowering_anyfuncs: [VMCallerCheckedAnyfunc; component.num_lowerings],
 //      lowerings: [VMLowering; component.num_lowerings],
 //      memories: [*mut VMMemoryDefinition; component.num_memories],
@@ -12,7 +12,8 @@
 // }
 
 use crate::component::{
-    Component, LoweredIndex, RuntimeMemoryIndex, RuntimePostReturnIndex, RuntimeReallocIndex,
+    Component, LoweredIndex, RuntimeComponentInstanceIndex, RuntimeMemoryIndex,
+    RuntimePostReturnIndex, RuntimeReallocIndex,
 };
 use crate::PtrSize;
 
@@ -48,11 +49,14 @@ pub struct VMComponentOffsets<P> {
     pub num_runtime_reallocs: u32,
     /// The number of post-returns which are recorded in this component for options.
     pub num_runtime_post_returns: u32,
+    /// Number of component instances internally in the component (always at
+    /// least 1).
+    pub num_runtime_component_instances: u32,
 
     // precalculated offsets of various member fields
     magic: u32,
-    flags: u32,
     store: u32,
+    flags: u32,
     lowering_anyfuncs: u32,
     lowerings: u32,
     memories: u32,
@@ -77,9 +81,13 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             num_runtime_memories: component.num_runtime_memories.try_into().unwrap(),
             num_runtime_reallocs: component.num_runtime_reallocs.try_into().unwrap(),
             num_runtime_post_returns: component.num_runtime_post_returns.try_into().unwrap(),
+            num_runtime_component_instances: component
+                .num_runtime_component_instances
+                .try_into()
+                .unwrap(),
             magic: 0,
-            flags: 0,
             store: 0,
+            flags: 0,
             lowering_anyfuncs: 0,
             lowerings: 0,
             memories: 0,
@@ -114,9 +122,10 @@ impl<P: PtrSize> VMComponentOffsets<P> {
 
         fields! {
             size(magic) = 4u32,
-            size(flags) = 1u32,
             align(u32::from(ret.ptr.size())),
             size(store) = cmul(2, ret.ptr.size()),
+            size(flags) = cmul(ret.num_runtime_component_instances, ret.size_of_vmcomponent_flags()),
+            align(u32::from(ret.ptr.size())),
             size(lowering_anyfuncs) = cmul(ret.num_lowerings, ret.ptr.size_of_vmcaller_checked_anyfunc()),
             size(lowerings) = cmul(ret.num_lowerings, ret.ptr.size() * 2),
             size(memories) = cmul(ret.num_runtime_memories, ret.ptr.size()),
@@ -146,10 +155,17 @@ impl<P: PtrSize> VMComponentOffsets<P> {
         self.magic
     }
 
+    /// The size of the `VMComponentFlags` type.
+    #[inline]
+    pub fn size_of_vmcomponent_flags(&self) -> u8 {
+        1
+    }
+
     /// The offset of the `flags` field.
     #[inline]
-    pub fn flags(&self) -> u32 {
+    pub fn flags(&self, index: RuntimeComponentInstanceIndex) -> u32 {
-        self.flags
+        assert!(index.as_u32() < self.num_runtime_component_instances);
+        self.flags + index.as_u32()
     }
 
     /// The offset of the `store` field.

crates/runtime/src/component.rs

@@ -17,9 +17,10 @@ use std::mem;
 use std::ops::Deref;
 use std::ptr::{self, NonNull};
 use wasmtime_environ::component::{
-    Component, LoweredIndex, RuntimeMemoryIndex, RuntimePostReturnIndex, RuntimeReallocIndex,
+    Component, LoweredIndex, RuntimeComponentInstanceIndex, RuntimeMemoryIndex,
-    StringEncoding, VMComponentOffsets, VMCOMPONENT_FLAG_MAY_ENTER, VMCOMPONENT_FLAG_MAY_LEAVE,
+    RuntimePostReturnIndex, RuntimeReallocIndex, StringEncoding, VMComponentOffsets,
-    VMCOMPONENT_FLAG_NEEDS_POST_RETURN, VMCOMPONENT_MAGIC,
+    VMCOMPONENT_FLAG_MAY_ENTER, VMCOMPONENT_FLAG_MAY_LEAVE, VMCOMPONENT_FLAG_NEEDS_POST_RETURN,
+    VMCOMPONENT_MAGIC,
 };
 use wasmtime_environ::HostPtr;
 
@@ -52,6 +53,8 @@ pub struct ComponentInstance {
 ///   end up being a `VMComponentContext`.
 /// * `data` - this is the data pointer associated with the `VMLowering` for
 ///   which this function pointer was registered.
+/// * `flags` - the component flags for may_enter/leave corresponding to the
+///   component instance that the lowering happened within.
 /// * `opt_memory` - this nullable pointer represents the memory configuration
 ///   option for the canonical ABI options.
 /// * `opt_realloc` - this nullable pointer represents the realloc configuration
@@ -65,13 +68,14 @@ pub struct ComponentInstance {
 /// * `nargs_and_results` - the size, in units of `ValRaw`, of
 ///   `args_and_results`.
 //
-// FIXME: 7 arguments is probably too many. The `data` through `string-encoding`
+// FIXME: 8 arguments is probably too many. The `data` through `string-encoding`
 // parameters should probably get packaged up into the `VMComponentContext`.
 // Needs benchmarking one way or another though to figure out what the best
 // balance is here.
 pub type VMLoweringCallee = extern "C" fn(
     vmctx: *mut VMOpaqueContext,
     data: *mut u8,
+    flags: *mut VMComponentFlags,
     opt_memory: *mut VMMemoryDefinition,
     opt_realloc: *mut VMCallerCheckedAnyfunc,
     string_encoding: StringEncoding,
@@ -170,8 +174,8 @@ impl ComponentInstance {
 
     /// Returns a pointer to the "may leave" flag for this instance specified
     /// for canonical lowering and lifting operations.
-    pub fn flags(&self) -> *mut VMComponentFlags {
+    pub fn flags(&self, instance: RuntimeComponentInstanceIndex) -> *mut VMComponentFlags {
-        unsafe { self.vmctx_plus_offset(self.offsets.flags()) }
+        unsafe { self.vmctx_plus_offset(self.offsets.flags(instance)) }
     }
 
     /// Returns the store that this component was created with.
@@ -338,8 +342,11 @@ impl ComponentInstance {
 
     unsafe fn initialize_vmctx(&mut self, store: *mut dyn Store) {
         *self.vmctx_plus_offset(self.offsets.magic()) = VMCOMPONENT_MAGIC;
-        *self.flags() = VMComponentFlags::new();
         *self.vmctx_plus_offset(self.offsets.store()) = store;
+        for i in 0..self.offsets.num_runtime_component_instances {
+            let i = RuntimeComponentInstanceIndex::from_u32(i);
+            *self.flags(i) = VMComponentFlags::new();
+        }
 
         // In debug mode set non-null bad values to all "pointer looking" bits
         // and pices related to lowering and such. This'll help detect any
@@ -555,3 +562,19 @@ impl VMComponentFlags {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::mem::size_of;
+
+    #[test]
+    fn size_of_vmcomponent_flags() {
+        let component = Component::default();
+        let offsets = VMComponentOffsets::new(size_of::<*mut u8>() as u8, &component);
+        assert_eq!(
+            size_of::<VMComponentFlags>(),
+            usize::from(offsets.size_of_vmcomponent_flags())
+        );
+    }
+}

crates/wasmtime/src/component/component.rs

@@ -2,13 +2,14 @@ use crate::signatures::SignatureCollection;
 use crate::{Engine, Module};
 use anyhow::{bail, Context, Result};
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::fs;
 use std::ops::Range;
 use std::path::Path;
 use std::ptr::NonNull;
 use std::sync::Arc;
 use wasmtime_environ::component::{
-    ComponentTypes, GlobalInitializer, LoweredIndex, StaticModuleIndex, TrampolineInfo, Translator,
+    ComponentTypes, GlobalInitializer, LoweredIndex, LoweringInfo, StaticModuleIndex, Translator,
 };
 use wasmtime_environ::PrimaryMap;
 use wasmtime_jit::CodeMemory;
@@ -52,7 +53,7 @@ struct ComponentInner {
 
     /// Where trampolines are located within the `text` section of
     /// `trampoline_obj`.
-    trampolines: PrimaryMap<LoweredIndex, TrampolineInfo>,
+    trampolines: PrimaryMap<LoweredIndex, LoweringInfo>,
 }
 
 impl Component {
@@ -116,6 +117,40 @@ impl Component {
             .context("failed to parse WebAssembly module")?;
         let types = Arc::new(types.finish());
 
+        // All lowered functions will require a trampoline to be available in
+        // case they're used when entering wasm. For example a lowered function
+        // could be immediately lifted in which case we'll need a trampoline to
+        // call that lowered function.
+        //
+        // Most of the time trampolines can come from the core wasm modules
+        // since lifted functions come from core wasm. For these esoteric cases
+        // though we may have to compile trampolines specifically into the
+        // component object as well in case core wasm doesn't provide the
+        // necessary trampoline.
+        let lowerings = component
+            .initializers
+            .iter()
+            .filter_map(|init| match init {
+                GlobalInitializer::LowerImport(i) => Some(i),
+                _ => None,
+            })
+            .collect::<Vec<_>>();
+        let required_trampolines = lowerings
+            .iter()
+            .map(|l| l.canonical_abi)
+            .collect::<HashSet<_>>();
+        let provided_trampolines = modules
+            .iter()
+            .flat_map(|(_, m)| m.exported_signatures.iter().copied())
+            .collect::<HashSet<_>>();
+        let mut trampolines_to_compile = required_trampolines
+            .difference(&provided_trampolines)
+            .collect::<Vec<_>>();
+        // Ensure a deterministically compiled artifact by sorting this list
+        // which was otherwise created with nondeterministically ordered hash
+        // tables.
+        trampolines_to_compile.sort();
+
         let (static_modules, trampolines) = engine.join_maybe_parallel(
             // In one (possibly) parallel task all the modules found within this
             // component are compiled. Note that this will further parallelize
@@ -139,28 +174,40 @@ impl Component {
             // In another (possibly) parallel task we compile lowering
             // trampolines necessary found in the component.
             || -> Result<_> {
-                let lowerings = component
+                let compiler = engine.compiler();
-                    .initializers
+                let (lowered_trampolines, core_trampolines) = engine.join_maybe_parallel(
-                    .iter()
+                    // Compile all the lowered trampolines here which implement
-                    .filter_map(|init| match init {
+                    // `canon lower` and are used to exit wasm into the host.
-                        GlobalInitializer::LowerImport(i) => Some(i),
+                    || -> Result<_> {
-                        _ => None,
+                        Ok(engine
-                    })
-                    .collect::<Vec<_>>();
-                let compiler = engine.compiler().component_compiler();
-                let trampolines = engine
                             .run_maybe_parallel(lowerings, |lowering| {
-                        compiler.compile_lowered_trampoline(&component, lowering, &types)
+                                compiler
+                                    .component_compiler()
+                                    .compile_lowered_trampoline(&component, lowering, &types)
                             })?
                             .into_iter()
-                    .collect();
+                            .collect())
+                    },
+                    // Compile all entry host-to-wasm trampolines here that
+                    // aren't otherwise provided by core wasm modules.
+                    || -> Result<_> {
+                        engine.run_maybe_parallel(trampolines_to_compile.clone(), |i| {
+                            let ty = &types[*i];
+                            Ok((*i, compiler.compile_host_to_wasm_trampoline(ty)?))
+                        })
+                    },
+                );
                 let mut obj = engine.compiler().object()?;
-                let trampolines = compiler.emit_obj(trampolines, &mut obj)?;
+                let trampolines = compiler.component_compiler().emit_obj(
+                    lowered_trampolines?,
+                    core_trampolines?,
+                    &mut obj,
+                )?;
                 Ok((trampolines, wasmtime_jit::mmap_vec_from_obj(obj)?))
             },
         );
         let static_modules = static_modules?;
-        let (trampolines, trampoline_obj) = trampolines?;
+        let ((lowering_trampolines, core_trampolines), trampoline_obj) = trampolines?;
         let mut trampoline_obj = CodeMemory::new(trampoline_obj);
         let code = trampoline_obj.publish()?;
         let text = wasmtime_jit::subslice_range(code.text, code.mmap);
@@ -184,6 +231,13 @@ impl Component {
                 vmtrampolines.insert(idx, trampoline);
             }
         }
+        for (signature, trampoline) in trampolines_to_compile.iter().zip(core_trampolines) {
+            vmtrampolines.insert(**signature, unsafe {
+                let ptr =
+                    code.text[trampoline.start as usize..][..trampoline.length as usize].as_ptr();
+                std::mem::transmute::<*const u8, wasmtime_runtime::VMTrampoline>(ptr)
+            });
+        }
 
         // FIXME: for the same reason as above where each module is
         // re-registering everything this should only be registered once. This
@@ -202,7 +256,7 @@ impl Component {
                 signatures,
                 trampoline_obj,
                 text,
-                trampolines,
+                trampolines: lowering_trampolines,
             }),
         })
     }

crates/wasmtime/src/component/func.rs

@@ -5,8 +5,9 @@ use anyhow::{Context, Result};
 use std::mem::MaybeUninit;
 use std::ptr::NonNull;
 use std::sync::Arc;
-use wasmtime_environ::component::{CanonicalOptions, ComponentTypes, CoreExport, TypeFuncIndex};
+use wasmtime_environ::component::{
-use wasmtime_environ::FuncIndex;
+    CanonicalOptions, ComponentTypes, CoreDef, RuntimeComponentInstanceIndex, TypeFuncIndex,
+};
 use wasmtime_runtime::{Export, ExportFunction, VMTrampoline};
 
 const MAX_STACK_PARAMS: usize = 16;
@@ -82,6 +83,7 @@ pub struct FuncData {
     types: Arc<ComponentTypes>,
     options: Options,
     instance: Instance,
+    component_instance: RuntimeComponentInstanceIndex,
     post_return: Option<(ExportFunction, VMTrampoline)>,
     post_return_arg: Option<ValRaw>,
 }
@@ -92,10 +94,10 @@ impl Func {
         instance: &Instance,
         data: &InstanceData,
         ty: TypeFuncIndex,
-        func: &CoreExport<FuncIndex>,
+        func: &CoreDef,
         options: &CanonicalOptions,
     ) -> Func {
-        let export = match data.lookup_export(store, func) {
+        let export = match data.lookup_def(store, func) {
             Export::Function(f) => f,
             _ => unreachable!(),
         };
@@ -109,6 +111,7 @@ impl Func {
             let trampoline = store.lookup_trampoline(unsafe { anyfunc.as_ref() });
             (ExportFunction { anyfunc }, trampoline)
         });
+        let component_instance = options.instance;
         let options = unsafe { Options::new(store.id(), memory, realloc, options.string_encoding) };
         Func(store.store_data_mut().insert(FuncData {
             trampoline,
@@ -117,6 +120,7 @@ impl Func {
             ty,
             types: data.component_types().clone(),
             instance: *instance,
+            component_instance,
             post_return,
             post_return_arg: None,
         }))

crates/wasmtime/src/component/func/host.rs

@@ -27,6 +27,7 @@ pub trait IntoComponentFunc<T, Params, Return> {
     extern "C" fn entrypoint(
         cx: *mut VMOpaqueContext,
         data: *mut u8,
+        flags: *mut VMComponentFlags,
         memory: *mut VMMemoryDefinition,
         realloc: *mut VMCallerCheckedAnyfunc,
         string_encoding: StringEncoding,
@@ -105,6 +106,7 @@ where
 /// the select few places it's intended to be called from.
 unsafe fn call_host<T, Params, Return, F>(
     cx: *mut VMOpaqueContext,
+    flags: *mut VMComponentFlags,
     memory: *mut VMMemoryDefinition,
     realloc: *mut VMCallerCheckedAnyfunc,
     string_encoding: StringEncoding,
@@ -136,7 +138,6 @@ where
 
     let cx = VMComponentContext::from_opaque(cx);
     let instance = (*cx).instance();
-    let flags = (*instance).flags();
     let mut cx = StoreContextMut::from_raw((*instance).store());
 
     let options = Options::new(
@@ -282,6 +283,7 @@ macro_rules! impl_into_component_func {
             extern "C" fn entrypoint(
                 cx: *mut VMOpaqueContext,
                 data: *mut u8,
+                flags: *mut VMComponentFlags,
                 memory: *mut VMMemoryDefinition,
                 realloc: *mut VMCallerCheckedAnyfunc,
                 string_encoding: StringEncoding,
@@ -292,6 +294,7 @@ macro_rules! impl_into_component_func {
                 unsafe {
                     handle_result(|| call_host::<T, _, _, _>(
                         cx,
+                        flags,
                         memory,
                         realloc,
                         string_encoding,
@@ -318,6 +321,7 @@ macro_rules! impl_into_component_func {
             extern "C" fn entrypoint(
                 cx: *mut VMOpaqueContext,
                 data: *mut u8,
+                flags: *mut VMComponentFlags,
                 memory: *mut VMMemoryDefinition,
                 realloc: *mut VMCallerCheckedAnyfunc,
                 string_encoding: StringEncoding,
@@ -328,6 +332,7 @@ macro_rules! impl_into_component_func {
                 unsafe {
                     handle_result(|| call_host::<T, _, _, _>(
                         cx,
+                        flags,
                         memory,
                         realloc,
                         string_encoding,

crates/wasmtime/src/component/func/typed.rs

@@ -308,6 +308,7 @@ where
             export,
             options,
             instance,
+            component_instance,
             ..
         } = store.0[self.func.0];
 
@@ -329,7 +330,7 @@ where
         assert!(mem::align_of_val(map_maybe_uninit!(space.ret)) == val_align);
 
         let instance = store.0[instance.0].as_ref().unwrap().instance();
-        let flags = instance.flags();
+        let flags = instance.flags(component_instance);
 
         unsafe {
             if !(*flags).may_enter() {
@@ -448,9 +449,10 @@ where
         let data = &mut store.0[self.func.0];
         let instance = data.instance;
         let post_return = data.post_return;
+        let component_instance = data.component_instance;
         let post_return_arg = data.post_return_arg.take();
         let instance = store.0[instance.0].as_ref().unwrap().instance();
-        let flags = instance.flags();
+        let flags = instance.flags(component_instance);
 
         unsafe {
             // First assert that the instance is in a "needs post return" state.

crates/wasmtime/src/component/instance.rs

@@ -143,7 +143,7 @@ impl InstanceData {
         }
     }
 
-    fn lookup_def(&self, store: &mut StoreOpaque, def: &CoreDef) -> wasmtime_runtime::Export {
+    pub fn lookup_def(&self, store: &mut StoreOpaque, def: &CoreDef) -> wasmtime_runtime::Export {
         match def {
             CoreDef::Lowered(idx) => {
                 wasmtime_runtime::Export::Function(wasmtime_runtime::ExportFunction {

crates/wasmtime/src/module.rs

@@ -367,18 +367,20 @@ impl Module {
         mut translation: ModuleTranslation<'_>,
         types: &ModuleTypes,
     ) -> Result<(MmapVec, Option<CompiledModuleInfo>)> {
-        // Compile all functions in parallel using rayon. This will also perform
-        // validation of function bodies.
         let tunables = &engine.config().tunables;
         let functions = mem::take(&mut translation.function_body_inputs);
         let functions = functions.into_iter().collect::<Vec<_>>();
-        let funcs = engine
+        let compiler = engine.compiler();
-            .run_maybe_parallel(functions, |(index, func)| {
+        let (funcs, trampolines) = engine.join_maybe_parallel(
+            // In one (possibly) parallel task all wasm functions are compiled
+            // in parallel. Note that this is also where the actual validation
+            // of all function bodies happens as well.
+            || -> Result<_> {
+                let funcs = engine.run_maybe_parallel(functions, |(index, func)| {
                     let offset = func.body.range().start;
-                engine
+                    let result =
-                    .compiler()
+                        compiler.compile_function(&translation, index, func, tunables, types);
-                    .compile_function(&translation, index, func, tunables, types)
+                    result.with_context(|| {
-                    .with_context(|| {
                         let index = translation.module.func_index(index);
                         let name = match translation.debuginfo.name_section.func_names.get(&index) {
                             Some(name) => format!(" (`{}`)", name),
@@ -389,16 +391,28 @@ impl Module {
                             "failed to compile wasm function {index}{name} at offset {offset:#x}"
                         )
                     })
-            })?
+                })?;
-            .into_iter()
+
-            .collect();
+                Ok(funcs.into_iter().collect())
+            },
+            // In another (possibly) parallel task all trampolines necessary
+            // for untyped host-to-wasm entry are compiled. Note that this
+            // isn't really expected to take all that long, it's moreso "well
+            // if we're using rayon why not use it here too".
+            || -> Result<_> {
+                engine.run_maybe_parallel(translation.exported_signatures.clone(), |sig| {
+                    let ty = &types[sig];
+                    Ok(compiler.compile_host_to_wasm_trampoline(ty)?)
+                })
+            },
+        );
 
         // Collect all the function results into a final ELF object.
         let mut obj = engine.compiler().object()?;
         let (funcs, trampolines) =
             engine
                 .compiler()
-                .emit_obj(&translation, types, funcs, tunables, &mut obj)?;
+                .emit_obj(&translation, funcs?, trampolines?, tunables, &mut obj)?;
 
         // If configured attempt to use static memory initialization which
         // can either at runtime be implemented as a single memcpy to

tests/all/component_model/func.rs

@@ -2152,5 +2152,120 @@ fn raw_slice_of_various_types() -> Result<()> {
             i64::to_le(0x0f_0e_0d_0c_0b_0a_09_08),
         ]
     );
+
+    Ok(())
+}
+
+#[test]
+fn lower_then_lift() -> Result<()> {
+    // First test simple integers when the import/export ABI happen to line up
+    let component = r#"
+(component $c
+  (import "f" (func $f (result u32)))
+
+  (core func $f_lower
+    (canon lower (func $f))
+  )
+  (func $f2 (result s32)
+    (canon lift (core func $f_lower))
+  )
+  (export "f" (func $f2))
+)
+    "#;
+
+    let engine = super::engine();
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    let mut linker = Linker::new(&engine);
+    linker.root().func_wrap("f", || Ok(2u32))?;
+    let instance = linker.instantiate(&mut store, &component)?;
+
+    let f = instance.get_typed_func::<(), i32, _>(&mut store, "f")?;
+    assert_eq!(f.call(&mut store, ())?, 2);
+
+    // First test strings when the import/export ABI happen to line up
+    let component = format!(
+        r#"
+(component $c
+  (import "s" (func $f (param string)))
+
+  (core module $libc
+    (memory (export "memory") 1)
+    {REALLOC_AND_FREE}
+  )
+  (core instance $libc (instantiate $libc))
+
+  (core func $f_lower
+    (canon lower (func $f) (memory $libc "memory"))
+  )
+  (func $f2 (param string)
+    (canon lift (core func $f_lower)
+        (memory $libc "memory")
+        (realloc (func $libc "realloc"))
+    )
+  )
+  (export "f" (func $f2))
+)
+    "#
+    );
+
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    linker
+        .root()
+        .func_wrap("s", |store: StoreContextMut<'_, ()>, x: WasmStr| {
+            assert_eq!(x.to_str(&store)?, "hello");
+            Ok(())
+        })?;
+    let instance = linker.instantiate(&mut store, &component)?;
+
+    let f = instance.get_typed_func::<(&str,), (), _>(&mut store, "f")?;
+    f.call(&mut store, ("hello",))?;
+
+    // Next test "type punning" where return values are reinterpreted just
+    // because the return ABI happens to line up.
+    let component = format!(
+        r#"
+(component $c
+  (import "s2" (func $f (param string) (result u32)))
+
+  (core module $libc
+    (memory (export "memory") 1)
+    {REALLOC_AND_FREE}
+  )
+  (core instance $libc (instantiate $libc))
+
+  (core func $f_lower
+    (canon lower (func $f) (memory $libc "memory"))
+  )
+  (func $f2 (param string) (result string)
+    (canon lift (core func $f_lower)
+        (memory $libc "memory")
+        (realloc (func $libc "realloc"))
+    )
+  )
+  (export "f" (func $f2))
+)
+    "#
+    );
+
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    linker
+        .root()
+        .func_wrap("s2", |store: StoreContextMut<'_, ()>, x: WasmStr| {
+            assert_eq!(x.to_str(&store)?, "hello");
+            Ok(u32::MAX)
+        })?;
+    let instance = linker.instantiate(&mut store, &component)?;
+
+    let f = instance.get_typed_func::<(&str,), WasmStr, _>(&mut store, "f")?;
+    let err = f.call(&mut store, ("hello",)).err().unwrap();
+    assert!(
+        err.to_string().contains("return pointer not aligned"),
+        "{}",
+        err
+    );
+
     Ok(())
 }

tests/misc_testsuite/component-model/adapter.wast

@@ -70,3 +70,16 @@
     )
   )
 )
+
+;; lower something then immediately lift it
+(component $c
+  (import "host-return-two" (func $f (result u32)))
+
+  (core func $f_lower
+    (canon lower (func $f))
+  )
+  (func $f2 (result s32)
+    (canon lift (core func $f_lower))
+  )
+  (export "f" (func $f2))
+)