Add a VMComponentContext type and create it on instantiation (#4215)

* Add a `VMComponentContext` type and create it on instantiation

This commit fills out the `wasmtime-runtime` crate's support for
`VMComponentContext` and creates it as part of the instantiation
process. This moves a few maps that were temporarily allocated in an
`InstanceData` into the `VMComponentContext` and additionally reads the
canonical options data from there instead.

This type still won't be used in its "full glory" until the lowering of
host functions is completely implemented, however, which will be coming
in a future commit.

* Remove `DerefMut` implementation

* Rebase conflicts

crates/environ/src/component/info.rs

@@ -154,11 +154,21 @@ pub enum Initializer {
     /// 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`.
-    ExtractMemory(CoreExport<MemoryIndex>),
+    ExtractMemory {
+        /// The index of the memory being defined.
+        index: RuntimeMemoryIndex,
+        /// Where this memory is being extracted from.
+        export: CoreExport<MemoryIndex>,
+    },
 
     /// Same as `ExtractMemory`, except it's extracting a function pointer to be
     /// used as a `realloc` function.
-    ExtractRealloc(CoreDef),
+    ExtractRealloc {
+        /// The index of the realloc being defined.
+        index: RuntimeReallocIndex,
+        /// Where this realloc is being extracted from.
+        def: CoreDef,
+    },
 
     /// The `module` specified is saved into the runtime state at the next
     /// `RuntimeModuleIndex`, referred to later by `Export` definitions.

crates/environ/src/component/translate.rs

@@ -1027,7 +1027,7 @@ impl<'a, 'data> Translator<'a, 'data> {
         self.result
             .component
             .initializers
-            .push(Initializer::ExtractMemory(export));
+            .push(Initializer::ExtractMemory { index, export });
         index
     }
 
@@ -1041,7 +1041,7 @@ impl<'a, 'data> Translator<'a, 'data> {
         self.result
             .component
             .initializers
-            .push(Initializer::ExtractRealloc(def));
+            .push(Initializer::ExtractRealloc { index, def });
         index
     }
 }

crates/runtime/Cargo.toml

@@ -54,3 +54,5 @@ pooling-allocator = []
 # It is useful for applications that do not bind their own exception ports and
 # need portable signal handling.
 posix-signals-on-macos = []
+
+component-model = ["wasmtime-environ/component-model"]

crates/runtime/src/component.rs

@@ -0,0 +1,461 @@
+//! Runtime support for the component model in Wasmtime
+//!
+//! Currently this runtime support includes a `VMComponentContext` which is
+//! similar in purpose to `VMContext`. The context is read from
+//! cranelift-generated trampolines when entering the host from a wasm module.
+//! Eventually it's intended that module-to-module calls, which would be
+//! cranelift-compiled adapters, will use this `VMComponentContext` as well.
+
+use crate::{
+    Store, VMCallerCheckedAnyfunc, VMFunctionBody, VMMemoryDefinition, VMOpaqueContext,
+    VMSharedSignatureIndex, ValRaw,
+};
+use memoffset::offset_of;
+use std::alloc::{self, Layout};
+use std::marker;
+use std::mem;
+use std::ops::Deref;
+use std::ptr::{self, NonNull};
+use wasmtime_environ::component::{
+    Component, LoweredIndex, RuntimeMemoryIndex, RuntimeReallocIndex, StringEncoding,
+    VMComponentOffsets, VMCOMPONENT_MAGIC,
+};
+use wasmtime_environ::HostPtr;
+
+const INVALID_PTR: usize = 0xdead_dead_beef_beef_u64 as usize;
+
+/// Runtime representation of a component instance and all state necessary for
+/// the instance itself.
+///
+/// This type never exists by-value, but rather it's always behind a pointer.
+/// The size of the allocation for `ComponentInstance` includes the trailing
+/// `VMComponentContext` which is variably sized based on the `offsets`
+/// contained within.
+#[repr(C)]
+pub struct ComponentInstance {
+    /// Size and offset information for the trailing `VMComponentContext`.
+    offsets: VMComponentOffsets<HostPtr>,
+
+    /// A zero-sized field which represents the end of the struct for the actual
+    /// `VMComponentContext` to be allocated behind.
+    vmctx: VMComponentContext,
+}
+
+/// Type signature for host-defined trampolines that are called from
+/// WebAssembly.
+///
+/// This function signature is invoked from a cranelift-compiled trampoline that
+/// adapts from the core wasm System-V ABI into the ABI provided here:
+///
+/// * `vmctx` - this is the first argument to the wasm import, and should always
+///   end up being a `VMComponentContext`.
+/// * `data` - this is the data pointer associated with the `VMLowering` for
+///   which this function pointer was registered.
+/// * `opt_memory` - this nullable pointer represents the memory configuration
+///   option for the canonical ABI options.
+/// * `opt_realloc` - this nullable pointer represents the realloc configuration
+///   option for the canonical ABI options.
+/// * `string_encoding` - this is the configured string encoding for the
+///   canonical ABI this lowering corresponds to.
+/// * `args_and_results` - pointer to stack-allocated space in the caller where
+///   all the arguments are stored as well as where the results will be written
+///   to. The size and initialized bytes of this depends on the core wasm type
+///   signature that this callee corresponds to.
+/// * `nargs_and_results` - the size, in units of `ValRaw`, of
+///   `args_and_results`.
+pub type VMLoweringCallee = extern "C" fn(
+    vmctx: *mut VMOpaqueContext,
+    data: *mut u8,
+    opt_memory: *mut VMMemoryDefinition,
+    opt_realloc: *mut VMCallerCheckedAnyfunc,
+    string_encoding: StringEncoding,
+    args_and_results: *mut ValRaw,
+    nargs_and_results: usize,
+);
+
+/// Structure describing a lowered host function stored within a
+/// `VMComponentContext` per-lowering.
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub struct VMLowering {
+    /// The host function pointer that is invoked when this lowering is
+    /// invoked.
+    pub callee: VMLoweringCallee,
+    /// The host data pointer (think void* pointer) to get passed to `callee`.
+    pub data: *mut u8,
+}
+
+/// This is a marker type to represent the underlying allocation of a
+/// `VMComponentContext`.
+///
+/// This type is similar to `VMContext` for core wasm and is allocated once per
+/// component instance in Wasmtime. While the static size of this type is 0 the
+/// actual runtime size is variable depending on the shape of the component that
+/// this corresponds to. This structure always trails a `ComponentInstance`
+/// allocation and the allocation/liftetime of this allocation is managed by
+/// `ComponentInstance`.
+#[repr(C)]
+// Set an appropriate alignment for this structure where the most-aligned value
+// internally right now is a pointer.
+#[cfg_attr(target_pointer_width = "32", repr(align(4)))]
+#[cfg_attr(target_pointer_width = "64", repr(align(8)))]
+pub struct VMComponentContext {
+    /// For more information about this see the equivalent field in `VMContext`
+    _marker: marker::PhantomPinned,
+}
+
+impl ComponentInstance {
+    /// Returns the layout corresponding to what would be an allocation of a
+    /// `ComponentInstance` for the `offsets` provided.
+    ///
+    /// The returned layout has space for both the `ComponentInstance` and the
+    /// trailing `VMComponentContext`.
+    fn alloc_layout(offsets: &VMComponentOffsets<HostPtr>) -> Layout {
+        let size = mem::size_of::<Self>()
+            .checked_add(usize::try_from(offsets.size_of_vmctx()).unwrap())
+            .unwrap();
+        let align = mem::align_of::<Self>();
+        Layout::from_size_align(size, align).unwrap()
+    }
+
+    /// Initializes an uninitialized pointer to a `ComponentInstance` in
+    /// addition to its trailing `VMComponentContext`.
+    ///
+    /// The `ptr` provided must be valid for `alloc_size` bytes and will be
+    /// entirely overwritten by this function call. The `offsets` correspond to
+    /// the shape of the component being instantiated and `store` is a pointer
+    /// back to the Wasmtime store for host functions to have access to.
+    unsafe fn new_at(
+        ptr: *mut ComponentInstance,
+        alloc_size: usize,
+        offsets: VMComponentOffsets<HostPtr>,
+        store: *mut dyn Store,
+    ) {
+        assert!(alloc_size >= Self::alloc_layout(&offsets).size());
+
+        ptr::write(
+            ptr,
+            ComponentInstance {
+                offsets,
+                vmctx: VMComponentContext {
+                    _marker: marker::PhantomPinned,
+                },
+            },
+        );
+
+        (*ptr).initialize_vmctx(store);
+    }
+
+    fn vmctx(&self) -> *mut VMComponentContext {
+        &self.vmctx as *const VMComponentContext as *mut VMComponentContext
+    }
+
+    unsafe fn vmctx_plus_offset<T>(&self, offset: u32) -> *mut T {
+        self.vmctx()
+            .cast::<u8>()
+            .add(usize::try_from(offset).unwrap())
+            .cast()
+    }
+
+    /// Returns a pointer to the "may leave" flag for this instance specified
+    /// for canonical lowering and lifting operations.
+    pub fn may_leave(&self) -> *mut bool {
+        unsafe { self.vmctx_plus_offset(self.offsets.may_leave()) }
+    }
+
+    /// Returns a pointer to the "may enter" flag for this instance specified
+    /// for canonical lowering and lifting operations.
+    pub fn may_enter(&self) -> *mut bool {
+        unsafe { self.vmctx_plus_offset(self.offsets.may_enter()) }
+    }
+
+    /// Returns the store that this component was created with.
+    pub fn store(&self) -> *mut dyn Store {
+        unsafe {
+            let ret = *self.vmctx_plus_offset::<*mut dyn Store>(self.offsets.store());
+            assert!(!ret.is_null());
+            ret
+        }
+    }
+
+    /// Returns the runtime memory definition corresponding to the index of the
+    /// memory provided.
+    ///
+    /// This can only be called after `idx` has been initialized at runtime
+    /// during the instantiation process of a component.
+    pub fn runtime_memory(&self, idx: RuntimeMemoryIndex) -> *mut VMMemoryDefinition {
+        unsafe {
+            let ret = *self.vmctx_plus_offset(self.offsets.runtime_memory(idx));
+            debug_assert!(ret as usize != INVALID_PTR);
+            ret
+        }
+    }
+
+    /// Returns the realloc pointer corresponding to the index provided.
+    ///
+    /// This can only be called after `idx` has been initialized at runtime
+    /// during the instantiation process of a component.
+    pub fn runtime_realloc(&self, idx: RuntimeReallocIndex) -> NonNull<VMCallerCheckedAnyfunc> {
+        unsafe {
+            let ret = *self.vmctx_plus_offset::<NonNull<_>>(self.offsets.runtime_realloc(idx));
+            debug_assert!(ret.as_ptr() as usize != INVALID_PTR);
+            ret
+        }
+    }
+    /// Returns the host information for the lowered function at the index
+    /// specified.
+    ///
+    /// This can only be called after `idx` has been initialized at runtime
+    /// during the instantiation process of a component.
+    pub fn lowering(&self, idx: LoweredIndex) -> VMLowering {
+        unsafe {
+            let ret = *self.vmctx_plus_offset::<VMLowering>(self.offsets.lowering(idx));
+            debug_assert!(ret.callee as usize != INVALID_PTR);
+            debug_assert!(ret.data as usize != INVALID_PTR);
+            ret
+        }
+    }
+
+    /// Returns the core wasm function pointer corresponding to the lowering
+    /// index specified.
+    ///
+    /// The returned function is suitable to pass directly to a wasm module
+    /// instantiation and the function is a cranelift-compiled trampoline.
+    ///
+    /// This can only be called after `idx` has been initialized at runtime
+    /// during the instantiation process of a component.
+    pub fn lowering_anyfunc(&self, idx: LoweredIndex) -> NonNull<VMCallerCheckedAnyfunc> {
+        unsafe {
+            let ret = self
+                .vmctx_plus_offset::<VMCallerCheckedAnyfunc>(self.offsets.lowering_anyfunc(idx));
+            debug_assert!((*ret).func_ptr.as_ptr() as usize != INVALID_PTR);
+            debug_assert!((*ret).vmctx as usize != INVALID_PTR);
+            NonNull::new(ret).unwrap()
+        }
+    }
+
+    /// 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_memory(&mut self, idx: RuntimeMemoryIndex, ptr: *mut VMMemoryDefinition) {
+        unsafe {
+            debug_assert!(!ptr.is_null());
+            let storage = self.vmctx_plus_offset(self.offsets.runtime_memory(idx));
+            debug_assert!(*storage as usize == INVALID_PTR);
+            *storage = ptr;
+        }
+    }
+
+    /// Same as `set_runtime_memory` but for realloc function pointers.
+    pub fn set_runtime_realloc(
+        &mut self,
+        idx: RuntimeReallocIndex,
+        ptr: NonNull<VMCallerCheckedAnyfunc>,
+    ) {
+        unsafe {
+            let storage = self.vmctx_plus_offset(self.offsets.runtime_realloc(idx));
+            debug_assert!(*storage as usize == INVALID_PTR);
+            *storage = ptr.as_ptr();
+        }
+    }
+
+    /// Configures a lowered host function with all the pieces necessary.
+    ///
+    /// * `idx` - the index that's being configured
+    /// * `lowering` - the host-related closure information to get invoked when
+    ///   the lowering is called.
+    /// * `anyfunc_func_ptr` - the cranelift-compiled trampoline which will
+    ///   read the `VMComponentContext` and invoke `lowering` provided. This
+    ///   function pointer will be passed to wasm if wasm needs to instantiate
+    ///   something.
+    /// * `anyfunc_type_index` - the signature index for the core wasm type
+    ///   registered within the engine already.
+    pub fn set_lowering(
+        &mut self,
+        idx: LoweredIndex,
+        lowering: VMLowering,
+        anyfunc_func_ptr: NonNull<VMFunctionBody>,
+        anyfunc_type_index: VMSharedSignatureIndex,
+    ) {
+        unsafe {
+            debug_assert!(
+                *self.vmctx_plus_offset::<usize>(self.offsets.lowering_callee(idx)) == INVALID_PTR
+            );
+            debug_assert!(
+                *self.vmctx_plus_offset::<usize>(self.offsets.lowering_data(idx)) == INVALID_PTR
+            );
+            debug_assert!(
+                *self.vmctx_plus_offset::<usize>(self.offsets.lowering_anyfunc(idx)) == INVALID_PTR
+            );
+            *self.vmctx_plus_offset(self.offsets.lowering(idx)) = lowering;
+            let vmctx = self.vmctx();
+            *self.vmctx_plus_offset(self.offsets.lowering_anyfunc(idx)) = VMCallerCheckedAnyfunc {
+                func_ptr: anyfunc_func_ptr,
+                type_index: anyfunc_type_index,
+                vmctx: VMOpaqueContext::from_vmcomponent(vmctx),
+            };
+        }
+    }
+
+    unsafe fn initialize_vmctx(&mut self, store: *mut dyn Store) {
+        *self.vmctx_plus_offset(self.offsets.magic()) = VMCOMPONENT_MAGIC;
+        *self.may_leave() = true;
+        *self.may_enter() = true;
+        *self.vmctx_plus_offset(self.offsets.store()) = store;
+
+        // 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
+        // erroneous usage and enable debug assertions above as well to prevent
+        // loading these before they're configured or setting them twice.
+        if cfg!(debug_assertions) {
+            for i in 0..self.offsets.num_lowerings {
+                let i = LoweredIndex::from_u32(i);
+                let offset = self.offsets.lowering_callee(i);
+                *self.vmctx_plus_offset(offset) = INVALID_PTR;
+                let offset = self.offsets.lowering_data(i);
+                *self.vmctx_plus_offset(offset) = INVALID_PTR;
+                let offset = self.offsets.lowering_anyfunc(i);
+                *self.vmctx_plus_offset(offset) = INVALID_PTR;
+            }
+            for i in 0..self.offsets.num_runtime_memories {
+                let i = RuntimeMemoryIndex::from_u32(i);
+                let offset = self.offsets.runtime_memory(i);
+                *self.vmctx_plus_offset(offset) = INVALID_PTR;
+            }
+            for i in 0..self.offsets.num_runtime_reallocs {
+                let i = RuntimeReallocIndex::from_u32(i);
+                let offset = self.offsets.runtime_realloc(i);
+                *self.vmctx_plus_offset(offset) = INVALID_PTR;
+            }
+        }
+    }
+}
+
+impl VMComponentContext {
+    /// Moves the `self` pointer backwards to the `ComponentInstance` pointer
+    /// that this `VMComponentContext` trails.
+    pub fn instance(&self) -> *mut ComponentInstance {
+        unsafe {
+            (self as *const Self as *mut u8)
+                .offset(-(offset_of!(ComponentInstance, vmctx) as isize))
+                as *mut ComponentInstance
+        }
+    }
+}
+
+/// An owned version of `ComponentInstance` which is akin to
+/// `Box<ComponentInstance>`.
+///
+/// This type can be dereferenced to `ComponentInstance` to access the
+/// underlying methods.
+pub struct OwnedComponentInstance {
+    ptr: ptr::NonNull<ComponentInstance>,
+}
+
+// Using `NonNull` turns off auto-derivation of these traits but the owned usage
+// here enables these trait impls so long as `ComponentInstance` itself
+// implements these traits.
+unsafe impl Send for OwnedComponentInstance where ComponentInstance: Send {}
+unsafe impl Sync for OwnedComponentInstance where ComponentInstance: Sync {}
+
+impl OwnedComponentInstance {
+    /// Allocates a new `ComponentInstance + VMComponentContext` pair on the
+    /// heap with `malloc` and configures it for the `component` specified.
+    pub fn new(component: &Component, store: *mut dyn Store) -> OwnedComponentInstance {
+        let offsets = VMComponentOffsets::new(HostPtr, component);
+        let layout = ComponentInstance::alloc_layout(&offsets);
+        unsafe {
+            // Technically it is not required to `alloc_zeroed` here. The
+            // primary reason for doing this is because a component context
+            // start is a "partly initialized" state where pointers and such are
+            // configured as the instantiation process continues. The component
+            // model should guarantee that we never access uninitialized memory
+            // in the context, but to help protect against possible bugs a
+            // zeroed allocation is done here to try to contain
+            // use-before-initialized issues.
+            let ptr = alloc::alloc_zeroed(layout) as *mut ComponentInstance;
+            let ptr = ptr::NonNull::new(ptr).unwrap();
+
+            ComponentInstance::new_at(ptr.as_ptr(), layout.size(), offsets, store);
+
+            OwnedComponentInstance { ptr }
+        }
+    }
+
+    // Note that this is technically unsafe due to the fact that it enables
+    // `mem::swap`-ing two component instances which would get all the offsets
+    // mixed up and cause issues. This is scoped to just this module though as a
+    // convenience to forward to `&mut` methods on `ComponentInstance`.
+    unsafe fn instance_mut(&mut self) -> &mut ComponentInstance {
+        &mut *self.ptr.as_ptr()
+    }
+
+    /// See `ComponentInstance::set_runtime_memory`
+    pub fn set_runtime_memory(&mut self, idx: RuntimeMemoryIndex, ptr: *mut VMMemoryDefinition) {
+        unsafe { self.instance_mut().set_runtime_memory(idx, ptr) }
+    }
+
+    /// See `ComponentInstance::set_runtime_realloc`
+    pub fn set_runtime_realloc(
+        &mut self,
+        idx: RuntimeReallocIndex,
+        ptr: NonNull<VMCallerCheckedAnyfunc>,
+    ) {
+        unsafe { self.instance_mut().set_runtime_realloc(idx, ptr) }
+    }
+
+    /// See `ComponentInstance::set_lowering`
+    pub fn set_lowering(
+        &mut self,
+        idx: LoweredIndex,
+        lowering: VMLowering,
+        anyfunc_func_ptr: NonNull<VMFunctionBody>,
+        anyfunc_type_index: VMSharedSignatureIndex,
+    ) {
+        unsafe {
+            self.instance_mut()
+                .set_lowering(idx, lowering, anyfunc_func_ptr, anyfunc_type_index)
+        }
+    }
+}
+
+impl Deref for OwnedComponentInstance {
+    type Target = ComponentInstance;
+    fn deref(&self) -> &ComponentInstance {
+        unsafe { &*self.ptr.as_ptr() }
+    }
+}
+
+impl Drop for OwnedComponentInstance {
+    fn drop(&mut self) {
+        let layout = ComponentInstance::alloc_layout(&self.offsets);
+        unsafe {
+            ptr::drop_in_place(self.ptr.as_ptr());
+            alloc::dealloc(self.ptr.as_ptr().cast(), layout);
+        }
+    }
+}
+
+impl VMComponentContext {
+    /// Helper function to cast between context types using a debug assertion to
+    /// protect against some mistakes.
+    #[inline]
+    pub unsafe fn from_opaque(opaque: *mut VMOpaqueContext) -> *mut VMComponentContext {
+        // See comments in `VMContext::from_opaque` for this debug assert
+        debug_assert_eq!((*opaque).magic, VMCOMPONENT_MAGIC);
+        opaque.cast()
+    }
+}
+
+impl VMOpaqueContext {
+    /// Helper function to clearly indicate the cast desired
+    #[inline]
+    pub fn from_vmcomponent(ptr: *mut VMComponentContext) -> *mut VMOpaqueContext {
+        ptr.cast()
+    }
+}

crates/runtime/src/lib.rs

@@ -30,6 +30,8 @@ use wasmtime_environ::DefinedMemoryIndex;
 use wasmtime_environ::FunctionInfo;
 use wasmtime_environ::SignatureIndex;
 
+#[cfg(feature = "component-model")]
+pub mod component;
 mod export;
 mod externref;
 mod imports;

crates/runtime/src/vmcontext.rs

@@ -1058,7 +1058,7 @@ pub type VMTrampoline =
 /// mode to debug-assert that the casts here are correct and have at least a
 /// little protection against incorrect casts.
 pub struct VMOpaqueContext {
-    magic: u32,
+    pub(crate) magic: u32,
     _marker: marker::PhantomPinned,
 }
 

crates/wasmtime/Cargo.toml

@@ -111,4 +111,5 @@ memory-init-cow = ["wasmtime-runtime/memory-init-cow"]
 component-model = [
   "wasmtime-environ/component-model",
   "wasmtime-cranelift?/component-model",
+  "wasmtime-runtime/component-model",
 ]

crates/wasmtime/src/component/func.rs

@@ -96,8 +96,8 @@ impl Func {
         let trampoline = store.lookup_trampoline(unsafe { export.anyfunc.as_ref() });
         let memory = options
             .memory
-            .map(|i| NonNull::new(data.runtime_memory(i).definition).unwrap());
+            .map(|i| NonNull::new(data.instance().runtime_memory(i)).unwrap());
-        let realloc = options.realloc.map(|i| data.runtime_realloc(i).anyfunc);
+        let realloc = options.realloc.map(|i| data.instance().runtime_realloc(i));
         let options = unsafe { Options::new(store.id(), memory, realloc, options.string_encoding) };
         Func(store.store_data_mut().insert(FuncData {
             trampoline,

crates/wasmtime/src/component/instance.rs

@@ -7,10 +7,11 @@ use std::marker;
 use std::sync::Arc;
 use wasmtime_environ::component::{
     ComponentTypes, CoreDef, CoreExport, Export, ExportItem, Initializer, InstantiateModule,
-    RuntimeImportIndex, RuntimeInstanceIndex, RuntimeMemoryIndex, RuntimeModuleIndex,
+    LowerImport, RuntimeImportIndex, RuntimeInstanceIndex, RuntimeMemoryIndex, RuntimeModuleIndex,
     RuntimeReallocIndex,
 };
-use wasmtime_environ::{EntityIndex, PrimaryMap};
+use wasmtime_environ::{EntityIndex, MemoryIndex, PrimaryMap};
+use wasmtime_runtime::component::{ComponentInstance, OwnedComponentInstance};
 
 /// An instantiated component.
 ///
@@ -30,9 +31,7 @@ pub(crate) struct InstanceData {
     component: Component,
     exported_modules: PrimaryMap<RuntimeModuleIndex, Module>,
 
-    // TODO: move these to `VMComponentContext`
+    state: OwnedComponentInstance,
-    memories: PrimaryMap<RuntimeMemoryIndex, wasmtime_runtime::ExportMemory>,
-    reallocs: PrimaryMap<RuntimeReallocIndex, wasmtime_runtime::ExportFunction>,
 }
 
 impl Instance {
@@ -161,20 +160,13 @@ impl InstanceData {
         instance.get_export_by_index(idx)
     }
 
+    pub fn instance(&self) -> &ComponentInstance {
+        &self.state
+    }
+
     pub fn component_types(&self) -> &Arc<ComponentTypes> {
         self.component.types()
     }
-
-    pub fn runtime_memory(&self, memory: RuntimeMemoryIndex) -> wasmtime_runtime::ExportMemory {
-        self.memories[memory].clone()
-    }
-
-    pub fn runtime_realloc(
-        &self,
-        realloc: RuntimeReallocIndex,
-    ) -> wasmtime_runtime::ExportFunction {
-        self.reallocs[realloc].clone()
-    }
 }
 
 struct Instantiator<'a> {
@@ -191,6 +183,7 @@ pub enum RuntimeImport {
 impl<'a> Instantiator<'a> {
     fn new(
         component: &'a Component,
+        store: &mut StoreOpaque,
         imports: &'a PrimaryMap<RuntimeImportIndex, RuntimeImport>,
     ) -> Instantiator<'a> {
         let env_component = component.env_component();
@@ -204,8 +197,7 @@ impl<'a> Instantiator<'a> {
                 exported_modules: PrimaryMap::with_capacity(
                     env_component.num_runtime_modules as usize,
                 ),
-                memories: Default::default(),
+                state: OwnedComponentInstance::new(env_component, store.traitobj()),
-                reallocs: Default::default(),
             },
         }
     }
@@ -246,31 +238,14 @@ impl<'a> Instantiator<'a> {
                         unsafe { crate::Instance::new_started(store, module, imports.as_ref())? };
                     self.data.instances.push(i);
                 }
-                Initializer::LowerImport(i) => {
+                Initializer::LowerImport(import) => self.lower_import(import),
-                    drop(self.component.trampoline_ptr(i.index));
+
-                    drop(
+                Initializer::ExtractMemory { index, export } => {
-                        self.component
+                    self.extract_memory(store.0, *index, export)
-                            .signatures()
-                            .shared_signature(i.canonical_abi)
-                            .unwrap(),
-                    );
-                    unimplemented!()
                 }
 
-                Initializer::ExtractMemory(export) => {
+                Initializer::ExtractRealloc { index, def } => {
-                    let memory = match self.data.lookup_export(store.0, export) {
+                    self.extract_realloc(store.0, *index, def)
-                        wasmtime_runtime::Export::Memory(m) => m,
-                        _ => unreachable!(),
-                    };
-                    self.data.memories.push(memory);
-                }
-
-                Initializer::ExtractRealloc(def) => {
-                    let func = match self.data.lookup_def(store.0, def) {
-                        wasmtime_runtime::Export::Function(f) => f,
-                        _ => unreachable!(),
-                    };
-                    self.data.reallocs.push(func);
                 }
 
                 Initializer::SaveModuleUpvar(idx) => {
@@ -278,6 +253,7 @@ impl<'a> Instantiator<'a> {
                         .exported_modules
                         .push(self.component.upvar(*idx).clone());
                 }
+
                 Initializer::SaveModuleImport(idx) => {
                     self.data.exported_modules.push(match &self.imports[*idx] {
                         RuntimeImport::Module(m) => m.clone(),
@@ -288,6 +264,43 @@ impl<'a> Instantiator<'a> {
         Ok(())
     }
 
+    fn lower_import(&mut self, import: &LowerImport) {
+        drop(self.component.trampoline_ptr(import.index));
+        drop(
+            self.component
+                .signatures()
+                .shared_signature(import.canonical_abi)
+                .unwrap(),
+        );
+        unimplemented!()
+    }
+
+    fn extract_memory(
+        &mut self,
+        store: &mut StoreOpaque,
+        index: RuntimeMemoryIndex,
+        export: &CoreExport<MemoryIndex>,
+    ) {
+        let memory = match self.data.lookup_export(store, export) {
+            wasmtime_runtime::Export::Memory(m) => m,
+            _ => unreachable!(),
+        };
+        self.data.state.set_runtime_memory(index, memory.definition);
+    }
+
+    fn extract_realloc(
+        &mut self,
+        store: &mut StoreOpaque,
+        index: RuntimeReallocIndex,
+        def: &CoreDef,
+    ) {
+        let anyfunc = match self.data.lookup_def(store, def) {
+            wasmtime_runtime::Export::Function(f) => f.anyfunc,
+            _ => unreachable!(),
+        };
+        self.data.state.set_runtime_realloc(index, anyfunc);
+    }
+
     fn build_imports<'b>(
         &mut self,
         store: &mut StoreOpaque,
@@ -353,7 +366,7 @@ impl<T> InstancePre<T> {
     // TODO: needs more docs
     pub fn instantiate(&self, mut store: impl AsContextMut<Data = T>) -> Result<Instance> {
         let mut store = store.as_context_mut();
-        let mut i = Instantiator::new(&self.component, &self.imports);
+        let mut i = Instantiator::new(&self.component, store.0, &self.imports);
         i.run(&mut store)?;
         let data = Box::new(i.data);
         Ok(Instance(store.0.store_data_mut().insert(Some(data))))