wasmtime-environ: add initial table support (#10497)

* wasmtime-environ: add initial table support

In order to lower the [`thread.spawn_indirect`] canonical builtin from
the component model, we need to be able to refer to the tables exported
from instance. This change adds the initial plumbing to access such
tables.

[`thread.spawn_indirect`]: https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md#-canon-threadspawn_indirect

* review: add missing doc comment

* review: add missing field calculation

* review: fix a return type

* fix: gate `pub use` behind component-model feature

Author: abrown

crates/environ/src/component/dfg.rs

@@ -175,6 +175,7 @@ macro_rules! id {
 id! {
     pub struct InstanceId(u32);
     pub struct MemoryId(u32);
+    pub struct TableId(u32);
     pub struct ReallocId(u32);
     pub struct CallbackId(u32);
     pub struct AdapterId(u32);
@@ -497,6 +498,7 @@ impl ComponentDfg {
             dfg: &self,
             initializers: Vec::new(),
             runtime_memories: Default::default(),
+            runtime_tables: Default::default(),
             runtime_post_return: Default::default(),
             runtime_reallocs: Default::default(),
             runtime_callbacks: Default::default(),
@@ -538,6 +540,7 @@ impl ComponentDfg {
                 num_lowerings: linearize.num_lowerings,
 
                 num_runtime_memories: linearize.runtime_memories.len() as u32,
+                num_runtime_tables: linearize.runtime_tables.len() as u32,
                 num_runtime_post_returns: linearize.runtime_post_return.len() as u32,
                 num_runtime_reallocs: linearize.runtime_reallocs.len() as u32,
                 num_runtime_callbacks: linearize.runtime_callbacks.len() as u32,
@@ -574,6 +577,7 @@ struct LinearizeDfg<'a> {
     trampoline_defs: PrimaryMap<TrampolineIndex, info::Trampoline>,
     trampoline_map: HashMap<TrampolineIndex, TrampolineIndex>,
     runtime_memories: HashMap<MemoryId, RuntimeMemoryIndex>,
+    runtime_tables: HashMap<TableId, RuntimeTableIndex>,
     runtime_reallocs: HashMap<ReallocId, RuntimeReallocIndex>,
     runtime_callbacks: HashMap<CallbackId, RuntimeCallbackIndex>,
     runtime_post_return: HashMap<PostReturnId, RuntimePostReturnIndex>,

crates/environ/src/component/info.rs

@@ -141,6 +141,10 @@ pub struct Component {
     /// stored in two different locations).
     pub num_runtime_memories: u32,
 
+    /// The number of runtime tables (maximum `RuntimeTableIndex`) needed to
+    /// instantiate this component. See notes on `num_runtime_memories`.
+    pub num_runtime_tables: u32,
+
     /// The number of runtime reallocs (maximum `RuntimeReallocIndex`) needed to
     /// instantiate this component.
     ///
@@ -215,10 +219,11 @@ impl Component {
     }
 }
 
-/// GlobalInitializer instructions to get processed when instantiating a component
+/// GlobalInitializer instructions to get processed when instantiating a
+/// component.
 ///
-/// The variants of this enum are processed during the instantiation phase of
-/// a component in-order from front-to-back. These are otherwise emitted as a
+/// The variants of this enum are processed during the instantiation phase of a
+/// component in-order from front-to-back. These are otherwise emitted as a
 /// component is parsed and read and translated.
 //
 // FIXME(#2639) if processing this list is ever a bottleneck we could
@@ -257,11 +262,10 @@ pub enum GlobalInitializer {
     /// A core wasm linear memory is going to be saved into the
     /// `VMComponentContext`.
     ///
-    /// This instruction indicates that the `index`th core wasm linear memory
-    /// needs to be extracted from the `export` specified, a pointer to a
-    /// previously created module instance, and stored into the
-    /// `VMComponentContext` at the `index` specified. This lowering is then
-    /// used in the future by pointers from `CanonicalOptions`.
+    /// This instruction indicates that a core wasm linear memory needs to be
+    /// extracted from the `export` and stored into the `VMComponentContext` at
+    /// the `index` specified. This lowering is then used in the future by
+    /// pointers from `CanonicalOptions`.
     ExtractMemory(ExtractMemory),
 
     /// Same as `ExtractMemory`, except it's extracting a function pointer to be
@@ -276,14 +280,24 @@ pub enum GlobalInitializer {
     /// used as a `post-return` function.
     ExtractPostReturn(ExtractPostReturn),
 
+    /// A core wasm table is going to be saved into the `VMComponentContext`.
+    ///
+    /// This instruction indicates that s core wasm table needs to be extracted
+    /// from its `export` and stored into the `VMComponentContext` at the
+    /// `index` specified. During this extraction, we will also capture the
+    /// table's containing instance pointer to access the table at runtime. This
+    /// extraction is useful for `thread.spawn_indirect`.
+    ExtractTable(ExtractTable),
+
     /// Declares a new defined resource within this component.
     ///
     /// Contains information about the destructor, for example.
     Resource(Resource),
 }
 
-/// Metadata for extraction of a memory of what's being extracted and where it's
-/// going.
+/// Metadata for extraction of a memory; contains what's being extracted (the
+/// memory at `export`) and where it's going (the `index` within a
+/// `VMComponentContext`).
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ExtractMemory {
     /// The index of the memory being defined.
@@ -319,6 +333,15 @@ pub struct ExtractPostReturn {
     pub def: CoreDef,
 }
 
+/// Metadata for extraction of a table.
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ExtractTable {
+    /// The index of the table being defined in a `VMComponentContext`.
+    pub index: RuntimeTableIndex,
+    /// Where this table is being extracted from.
+    pub export: CoreExport<TableIndex>,
+}
+
 /// Different methods of instantiating a core wasm module.
 #[derive(Debug, Serialize, Deserialize)]
 pub enum InstantiateModule {

crates/environ/src/component/types.rs

@@ -222,6 +222,14 @@ indices! {
     /// Same as `RuntimeMemoryIndex` except for the `post-return` function.
     pub struct RuntimePostReturnIndex(u32);
 
+    /// Index representing a table extracted from a wasm instance which is
+    /// stored in a `VMComponentContext`. This is used to deduplicate references
+    /// to the same table when it's only stored once in a `VMComponentContext`.
+    ///
+    /// This does not correspond to anything in the binary format for the
+    /// component model.
+    pub struct RuntimeTableIndex(u32);
+
     /// Index for all trampolines that are compiled in Cranelift for a
     /// component.
     ///

crates/environ/src/component/vmcomponent_offsets.rs

@@ -8,6 +8,7 @@
 //      trampoline_func_refs: [VMFuncRef; component.num_trampolines],
 //      lowerings: [VMLowering; component.num_lowerings],
 //      memories: [*mut VMMemoryDefinition; component.num_runtime_memories],
+//      tables: [*mut VMTableDefinition; component.num_runtime_tables],
 //      reallocs: [*mut VMFuncRef; component.num_runtime_reallocs],
 //      post_returns: [*mut VMFuncRef; component.num_runtime_post_returns],
 //      resource_destructors: [*mut VMFuncRef; component.num_resources],
@@ -44,6 +45,8 @@ pub struct VMComponentOffsets<P> {
     pub num_lowerings: u32,
     /// The number of memories which are recorded in this component for options.
     pub num_runtime_memories: u32,
+    /// The number of tables which are recorded in this component for options.
+    pub num_runtime_tables: u32,
     /// The number of reallocs which are recorded in this component for options.
     pub num_runtime_reallocs: u32,
     /// The number of callbacks which are recorded in this component for options.
@@ -66,6 +69,7 @@ pub struct VMComponentOffsets<P> {
     trampoline_func_refs: u32,
     lowerings: u32,
     memories: u32,
+    tables: u32,
     reallocs: u32,
     callbacks: u32,
     post_returns: u32,
@@ -87,6 +91,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             ptr,
             num_lowerings: component.num_lowerings,
             num_runtime_memories: component.num_runtime_memories.try_into().unwrap(),
+            num_runtime_tables: component.num_runtime_tables.try_into().unwrap(),
             num_runtime_reallocs: component.num_runtime_reallocs.try_into().unwrap(),
             num_runtime_callbacks: component.num_runtime_callbacks.try_into().unwrap(),
             num_runtime_post_returns: component.num_runtime_post_returns.try_into().unwrap(),
@@ -103,6 +108,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             trampoline_func_refs: 0,
             lowerings: 0,
             memories: 0,
+            tables: 0,
             reallocs: 0,
             callbacks: 0,
             post_returns: 0,
@@ -145,6 +151,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             size(trampoline_func_refs) = cmul(ret.num_trampolines, ret.ptr.size_of_vm_func_ref()),
             size(lowerings) = cmul(ret.num_lowerings, ret.ptr.size() * 2),
             size(memories) = cmul(ret.num_runtime_memories, ret.ptr.size()),
+            size(tables) = cmul(ret.num_runtime_tables, ret.ptr.size()),
             size(reallocs) = cmul(ret.num_runtime_reallocs, ret.ptr.size()),
             size(callbacks) = cmul(ret.num_runtime_callbacks, ret.ptr.size()),
             size(post_returns) = cmul(ret.num_runtime_post_returns, ret.ptr.size()),
@@ -262,6 +269,20 @@ impl<P: PtrSize> VMComponentOffsets<P> {
         self.runtime_memories() + index.as_u32() * u32::from(self.ptr.size())
     }
 
+    /// The offset of the base of the `runtime_tables` field
+    #[inline]
+    pub fn runtime_tables(&self) -> u32 {
+        self.tables
+    }
+
+    /// The offset of the `*mut VMTableDefinition` for the runtime index
+    /// provided.
+    #[inline]
+    pub fn runtime_table(&self, index: RuntimeTableIndex) -> u32 {
+        assert!(index.as_u32() < self.num_runtime_tables);
+        self.runtime_tables() + index.as_u32() * u32::from(self.ptr.size())
+    }
+
     /// The offset of the base of the `runtime_reallocs` field
     #[inline]
     pub fn runtime_reallocs(&self) -> u32 {

crates/wasmtime/src/runtime/component/component.rs

@@ -623,6 +623,7 @@ impl Component {
                 },
                 GlobalInitializer::LowerImport { .. }
                 | GlobalInitializer::ExtractMemory(_)
+                | GlobalInitializer::ExtractTable(_)
                 | GlobalInitializer::ExtractRealloc(_)
                 | GlobalInitializer::ExtractCallback(_)
                 | GlobalInitializer::ExtractPostReturn(_)

crates/wasmtime/src/runtime/component/instance.rs

@@ -601,6 +601,8 @@ impl<'a> Instantiator<'a> {
                     self.data.state.set_lowering(*index, func.lowering());
                 }
 
+                GlobalInitializer::ExtractTable(table) => self.extract_table(store.0, table),
+
                 GlobalInitializer::ExtractMemory(mem) => self.extract_memory(store.0, mem),
 
                 GlobalInitializer::ExtractRealloc(realloc) => {
@@ -678,6 +680,16 @@ impl<'a> Instantiator<'a> {
             .set_runtime_post_return(post_return.index, func_ref);
     }
 
+    fn extract_table(&mut self, store: &mut StoreOpaque, table: &ExtractTable) {
+        let export = match self.data.lookup_export(store, &table.export) {
+            crate::runtime::vm::Export::Table(t) => t,
+            _ => unreachable!(),
+        };
+        self.data
+            .state
+            .set_runtime_table(table.index, export.definition);
+    }
+
     fn build_imports<'b>(
         &mut self,
         store: &mut StoreOpaque,

crates/wasmtime/src/runtime/vm.rs

@@ -90,6 +90,8 @@ pub use crate::runtime::vm::sys::unwind::UnwindRegistration;
 pub use crate::runtime::vm::table::{Table, TableElement};
 pub use crate::runtime::vm::traphandlers::*;
 pub use crate::runtime::vm::unwind::*;
+#[cfg(feature = "component-model")]
+pub use crate::runtime::vm::vmcontext::VMTableDefinition;
 pub use crate::runtime::vm::vmcontext::{
     VMArrayCallFunction, VMArrayCallHostFuncContext, VMContext, VMFuncRef, VMFunctionBody,
     VMFunctionImport, VMGlobalDefinition, VMGlobalImport, VMMemoryDefinition, VMMemoryImport,

crates/wasmtime/src/runtime/vm/component.rs

@@ -9,7 +9,8 @@
 use crate::prelude::*;
 use crate::runtime::vm::{
     SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
-    VMOpaqueContext, VMStore, VMStoreRawPtr, VMWasmCallFunction, ValRaw, VmPtr, VmSafe,
+    VMOpaqueContext, VMStore, VMStoreRawPtr, VMTableDefinition, VMWasmCallFunction, ValRaw, VmPtr,
+    VmSafe,
 };
 use alloc::alloc::Layout;
 use alloc::sync::Arc;
@@ -289,6 +290,19 @@ impl ComponentInstance {
         }
     }
 
+    /// Returns the runtime table definition corresponding to the index of the
+    /// table provided.
+    ///
+    /// This can only be called after `idx` has been initialized at runtime
+    /// during the instantiation process of a component.
+    pub fn runtime_table(&self, idx: RuntimeTableIndex) -> *mut VMTableDefinition {
+        unsafe {
+            let ret = *self.vmctx_plus_offset::<VmPtr<_>>(self.offsets.runtime_table(idx));
+            debug_assert!(ret.as_ptr() as usize != INVALID_PTR);
+            ret.as_ptr()
+        }
+    }
+
     /// Returns the realloc pointer corresponding to the index provided.
     ///
     /// This can only be called after `idx` has been initialized at runtime
@@ -404,6 +418,23 @@ impl ComponentInstance {
         }
     }
 
+    /// Stores the runtime memory pointer at the index specified.
+    ///
+    /// This is intended to be called during the instantiation process of a
+    /// component once a memory is available, which may not be until part-way
+    /// through component instantiation.
+    ///
+    /// Note that it should be a property of the component model that the `ptr`
+    /// here is never needed prior to it being configured here in the instance.
+    pub fn set_runtime_table(&mut self, idx: RuntimeTableIndex, ptr: NonNull<VMTableDefinition>) {
+        unsafe {
+            let storage = self
+                .vmctx_plus_offset_mut::<VmPtr<VMTableDefinition>>(self.offsets.runtime_table(idx));
+            debug_assert!((*storage).as_ptr() as usize == INVALID_PTR);
+            *storage = ptr.into();
+        }
+    }
+
     /// Configures host runtime lowering information associated with imported f
     /// functions for the `idx` specified.
     pub fn set_lowering(&mut self, idx: LoweredIndex, lowering: VMLowering) {
@@ -801,6 +832,11 @@ impl OwnedComponentInstance {
         unsafe { self.instance_mut().set_runtime_post_return(idx, ptr) }
     }
 
+    /// See `ComponentInstance::set_runtime_table`
+    pub fn set_runtime_table(&mut self, idx: RuntimeTableIndex, ptr: NonNull<VMTableDefinition>) {
+        unsafe { self.instance_mut().set_runtime_table(idx, ptr) }
+    }
+
     /// See `ComponentInstance::set_lowering`
     pub fn set_lowering(&mut self, idx: LoweredIndex, lowering: VMLowering) {
         unsafe { self.instance_mut().set_lowering(idx, lowering) }

crates/wasmtime/src/runtime/vm/instance/allocator/pooling.rs

@@ -527,6 +527,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
                 }
                 LowerImport { .. }
                 | ExtractMemory(_)
+                | ExtractTable(_)
                 | ExtractRealloc(_)
                 | ExtractCallback(_)
                 | ExtractPostReturn(_)