Improve table support for components (#10499)

* Fix documentation, add debug check

This change cleans up some missing items from #10497: a copy-paste typo
and some missing debug assertions about the initial state of the table
list in the `VMComponentContext`.

* Add missing `VMContext` field for component table use

When components access tables, these may lazily initialized. Eventually,
we will need to call `Instance::get_defined_table_with_lazy_init` to set
up the table immediately prior to use. To do so, this change stores an
additional `*mut VMContext` (for a core instance, not a component) in
the `VMComponentContext` along with each `*mut VMTableDefinition` so
that we can recover the core instance needed for lazy initialization.

* Use `VMTableImport` to record component tables

To avoid extra structures, this change uses `VMTableImport` everywhere a
component would need to keep track of a `VMTableDefinition` and
`VMContext` (recall the `VMContext` is needed for lazy initialization).
This is a bit awkward since this usage is not truly an import, but
perhaps we can change the name later.

Author: abrown

crates/environ/src/component/vmcomponent_offsets.rs

@@ -8,7 +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],
+//      tables: [VMTableImport; 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],
@@ -151,7 +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(tables) = cmul(ret.num_runtime_tables, ret.size_of_vmtable_import()),
             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()),
@@ -275,12 +275,18 @@ impl<P: PtrSize> VMComponentOffsets<P> {
         self.tables
     }
 
-    /// The offset of the `*mut VMTableDefinition` for the runtime index
-    /// provided.
+    /// The offset of the table 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())
+        self.runtime_tables() + index.as_u32() * u32::from(self.size_of_vmtable_import())
+    }
+
+    /// Return the size of `VMTableImport`, used here to hold the pointers to
+    /// the `VMTableDefinition` and `VMContext` (not because this is an import).
+    #[inline]
+    pub fn size_of_vmtable_import(&self) -> u8 {
+        2 * self.pointer_size()
     }
 
     /// The offset of the base of the `runtime_reallocs` field

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

@@ -687,7 +687,7 @@ impl<'a> Instantiator<'a> {
         };
         self.data
             .state
-            .set_runtime_table(table.index, export.definition);
+            .set_runtime_table(table.index, export.definition, export.vmctx);
     }
 
     fn build_imports<'b>(

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

@@ -8,9 +8,9 @@
 
 use crate::prelude::*;
 use crate::runtime::vm::{
-    SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
-    VMOpaqueContext, VMStore, VMStoreRawPtr, VMTableDefinition, VMWasmCallFunction, ValRaw, VmPtr,
-    VmSafe,
+    SendSyncPtr, VMArrayCallFunction, VMContext, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
+    VMOpaqueContext, VMStore, VMStoreRawPtr, VMTableDefinition, VMTableImport, VMWasmCallFunction,
+    ValRaw, VmPtr, VmSafe,
 };
 use alloc::alloc::Layout;
 use alloc::sync::Arc;
@@ -290,16 +290,20 @@ impl ComponentInstance {
         }
     }
 
-    /// Returns the runtime table definition corresponding to the index of the
-    /// table provided.
+    /// Returns the runtime table definition and associated instance `VMContext`
+    /// 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 {
+    ///
+    /// Note that we use a `VMTableImport` here because of its structure, not
+    /// because these tables are actually imports.
+    pub fn runtime_table(&self, idx: RuntimeTableIndex) -> VMTableImport {
         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()
+            let ret = *self.vmctx_plus_offset::<VMTableImport>(self.offsets.runtime_table(idx));
+            debug_assert!(ret.from.as_ptr() as usize != INVALID_PTR);
+            debug_assert!(ret.vmctx.as_ptr() as usize != INVALID_PTR);
+            ret
         }
     }
 
@@ -418,20 +422,31 @@ impl ComponentInstance {
         }
     }
 
-    /// Stores the runtime memory pointer at the index specified.
+    /// Stores the runtime table 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
+    /// component once a table 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>) {
+    pub fn set_runtime_table(
+        &mut self,
+        idx: RuntimeTableIndex,
+        ptr: NonNull<VMTableDefinition>,
+        vmctx: NonNull<VMContext>,
+    ) {
+        // Note that we use a `VMTableImport` here because of its structure, not
+        // because these tables are actually imports.
         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();
+            let storage =
+                self.vmctx_plus_offset_mut::<VMTableImport>(self.offsets.runtime_table(idx));
+            debug_assert!((*storage).vmctx.as_ptr() as usize == INVALID_PTR);
+            debug_assert!((*storage).from.as_ptr() as usize == INVALID_PTR);
+            *storage = VMTableImport {
+                vmctx: vmctx.into(),
+                from: ptr.into(),
+            };
         }
     }
 
@@ -554,6 +569,11 @@ impl ComponentInstance {
                 let offset = self.offsets.resource_destructor(i);
                 *self.vmctx_plus_offset_mut(offset) = INVALID_PTR;
             }
+            for i in 0..self.offsets.num_runtime_tables {
+                let i = RuntimeTableIndex::from_u32(i);
+                let offset = self.offsets.runtime_table(i);
+                *self.vmctx_plus_offset_mut(offset) = INVALID_PTR;
+            }
         }
     }
 
@@ -833,8 +853,13 @@ impl OwnedComponentInstance {
     }
 
     /// 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) }
+    pub fn set_runtime_table(
+        &mut self,
+        idx: RuntimeTableIndex,
+        ptr: NonNull<VMTableDefinition>,
+        vmctx: NonNull<VMContext>,
+    ) {
+        unsafe { self.instance_mut().set_runtime_table(idx, ptr, vmctx) }
     }
 
     /// See `ComponentInstance::set_lowering`

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

@@ -160,6 +160,7 @@ mod test_vmtable_import {
     use super::VMTableImport;
     use core::mem::offset_of;
     use std::mem::size_of;
+    use wasmtime_environ::component::{Component, VMComponentOffsets};
     use wasmtime_environ::{HostPtr, Module, VMOffsets};
 
     #[test]
@@ -179,6 +180,21 @@ mod test_vmtable_import {
             usize::from(offsets.vmtable_import_vmctx())
         );
     }
+
+    #[test]
+    fn ensure_sizes_match() {
+        // Because we use `VMTableImport` for recording tables used by
+        // components, we want to make sure that the size calculations between
+        // `VMOffsets` and `VMComponentOffsets` stay the same.
+        let module = Module::new();
+        let vm_offsets = VMOffsets::new(HostPtr, &module);
+        let component = Component::default();
+        let vm_component_offsets = VMComponentOffsets::new(HostPtr, &component);
+        assert_eq!(
+            vm_offsets.size_of_vmtable_import(),
+            vm_component_offsets.size_of_vmtable_import()
+        );
+    }
 }
 
 /// The fields compiled code needs to access to utilize a WebAssembly linear