Refactor Instance management within the CallStack (take 2)

crates/collections/src/head_vec.rs

@@ -0,0 +1,77 @@
+use core::mem;
+use std::vec::Vec;
+
+/// A [`Vec`]-like data structure with fast access to the last item.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct HeadVec<T> {
+    /// The top (or last) item in the [`HeadVec`].
+    head: Option<T>,
+    /// The rest of the items in the [`HeadVec`] excluding the last item.
+    rest: Vec<T>,
+}
+
+impl<T> Default for HeadVec<T> {
+    #[inline]
+    fn default() -> Self {
+        Self {
+            head: None,
+            rest: Vec::new(),
+        }
+    }
+}
+
+impl<T> HeadVec<T> {
+    /// Removes all items from the [`HeadVec`].
+    #[inline]
+    pub fn clear(&mut self) {
+        self.head = None;
+        self.rest.clear();
+    }
+
+    /// Returns the number of items stored in the [`HeadVec`].
+    #[inline]
+    pub fn len(&self) -> usize {
+        match self.head {
+            Some(_) => 1 + self.rest.len(),
+            None => 0,
+        }
+    }
+
+    /// Returns `true` if the [`HeadVec`] contains no items.
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Returns a shared reference to the last item in the [`HeadVec`] if any.
+    ///
+    /// Returns `None` if the [`HeadVec`] is empty.
+    #[inline]
+    pub fn last(&self) -> Option<&T> {
+        self.head.as_ref()
+    }
+
+    /// Returns an exclusive reference to the last item in the [`HeadVec`] if any.
+    ///
+    /// Returns `None` if the [`HeadVec`] is empty.
+    #[inline]
+    pub fn last_mut(&mut self) -> Option<&mut T> {
+        self.head.as_mut()
+    }
+
+    /// Pushes a new `value` onto the [`HeadVec`].
+    #[inline]
+    pub fn push(&mut self, value: T) {
+        let prev_head = self.head.replace(value);
+        if let Some(prev_head) = prev_head {
+            self.rest.push(prev_head);
+        }
+    }
+
+    /// Pops the last `value` from the [`HeadVec`] if any.
+    #[inline]
+    pub fn pop(&mut self) -> Option<T> {
+        let new_top = self.rest.pop();
+        mem::replace(&mut self.head, new_top)
+    }
+}

crates/collections/src/lib.rs

@@ -38,6 +38,7 @@ extern crate std;
 
 pub mod arena;
 pub mod hash;
+mod head_vec;
 pub mod map;
 pub mod set;
 pub mod string_interner;
@@ -48,6 +49,7 @@ mod tests;
 #[doc(inline)]
 pub use self::{
     arena::{Arena, ComponentVec, DedupArena},
+    head_vec::HeadVec,
     map::Map,
     set::Set,
     string_interner::StringInterner,

crates/wasmi/src/engine/executor/instrs.rs

@@ -152,8 +152,7 @@ impl<'engine> Executor<'engine> {
     #[inline(always)]
     fn instance(call_stack: &CallStack) -> &Instance {
         call_stack
-            .peek()
-            .map(CallFrame::instance)
+            .instance()
             .expect("missing instance for non-empty call stack")
     }
 

crates/wasmi/src/engine/executor/instrs/call.rs

@@ -214,7 +214,6 @@
         &mut self,
         results: RegisterSpan,
         func: &CompiledFuncEntity,
-        instance: &Instance,
     ) -> Result<CallFrame, Error> {
         // We have to reinstantiate the `self.sp` [`FrameRegisters`] since we just called
         // [`ValueStack::alloc_call_frame`] which might invalidate all live [`FrameRegisters`].
@@ -227,7 +226,7 @@
             self.sp = unsafe { this.stack_ptr_at(caller.base_offset()) };
         })?;
         let instr_ptr = InstructionPtr::new(func.instrs().as_ptr());
-        let frame = CallFrame::new(instr_ptr, offsets, results, *instance);
+        let frame = CallFrame::new(instr_ptr, offsets, results);
         if <C as CallContext>::HAS_PARAMS {
             self.copy_call_params(&mut uninit_params);
         }
@@ -300,8 +299,7 @@
         mut instance: Option<Instance>,
     ) -> Result<(), Error> {
         let func = self.code_map.get(Some(store.fuel_mut()), func)?;
-        let instance = instance.get_or_insert_with(|| *Self::instance(self.call_stack));
-        let mut called = self.dispatch_compiled_func::<C>(results, func, instance)?;
+        let mut called = self.dispatch_compiled_func::<C>(results, func)?;
         match <C as CallContext>::KIND {
             CallKind::Nested => {
                 // We need to update the instruction pointer of the caller call frame.
@@ -318,11 +316,16 @@
                 // on the value stack which is what the function expects. After this operation we ensure
                 // that `self.sp` is adjusted via a call to `init_call_frame` since it may have been
                 // invalidated by this method.
-                unsafe { Stack::merge_call_frames(self.call_stack, self.value_stack, &mut called) };
+                let caller_instance = unsafe {
+                    Stack::merge_call_frames(self.call_stack, self.value_stack, &mut called)
+                };
+                if let Some(caller_instance) = caller_instance {
+                    instance.get_or_insert(caller_instance);
+                }
             }
         }
         self.init_call_frame(&called);
-        self.call_stack.push(called)?;
+        self.call_stack.push(called, instance)?;
         Ok(())
     }
 
@@ -495,11 +498,15 @@
     ) -> Result<(), Error> {
         let (len_params, len_results) = self.func_types.len_params_results(host_func.ty_dedup());
         let max_inout = len_params.max(len_results);
+        let instance = *self
+            .call_stack
+            .instance()
+            .expect("need to have an instance on the call stack");
         // We have to reinstantiate the `self.sp` [`FrameRegisters`] since we just called
         // [`ValueStack::reserve`] which might invalidate all live [`FrameRegisters`].
         let caller = match <C as CallContext>::KIND {
             CallKind::Nested => self.call_stack.peek().copied(),
-            CallKind::Tail => self.call_stack.pop(),
+            CallKind::Tail => self.call_stack.pop().map(|(frame, _instance)| frame),
         }
         .expect("need to have a caller on the call stack");
         let buffer = self.value_stack.extend_by(max_inout, |this| {
@@ -513,13 +520,13 @@
         if matches!(<C as CallContext>::KIND, CallKind::Nested) {
             self.update_instr_ptr_at(1);
         }
-        self.dispatch_host_func::<T>(store, host_func, caller)
+        self.dispatch_host_func::<T>(store, host_func, &instance)
             .map_err(|error| match self.call_stack.is_empty() {
                 true => error,
                 false => ResumableHostError::new(error, *func, results).into(),
             })?;
-        self.memory.update(&mut store.inner, caller.instance());
-        self.global.update(&mut store.inner, caller.instance());
+        self.memory.update(&mut store.inner, &instance);
+        self.global.update(&mut store.inner, &instance);
         let results = results.iter(len_results);
         let returned = self.value_stack.drop_return(max_inout);
         for (result, value) in results.zip(returned) {
@@ -542,14 +549,14 @@
         &mut self,
         store: &mut Store<T>,
         host_func: HostFuncEntity,
-        caller: CallFrame,
+        instance: &Instance,
     ) -> Result<(usize, usize), Error> {
         dispatch_host_func(
             store,
             self.func_types,
             self.value_stack,
             host_func,
-            Some(caller.instance()),
+            Some(instance),
         )
     }
 

crates/wasmi/src/engine/executor/instrs/return_.rs

@@ -25,17 +25,19 @@ impl<'engine> Executor<'engine> {
     /// from the returning callee to the caller.
     #[inline(always)]
     fn return_impl(&mut self, store: &mut StoreInner) -> ReturnOutcome {
-        let returned = self
+        let (returned, popped_instance) = self
             .call_stack
             .pop()
             .expect("the executing call frame is always on the stack");
         self.value_stack.truncate(returned.frame_offset());
+        let new_instance = popped_instance.and_then(|_| self.call_stack.instance());
+        if let Some(new_instance) = new_instance {
+            self.global.update(store, new_instance);
+            self.memory.update(store, new_instance);
+        }
         match self.call_stack.peek() {
             Some(caller) => {
                 Self::init_call_frame_impl(self.value_stack, &mut self.sp, &mut self.ip, caller);
-                let instance = caller.instance();
-                self.memory.update(store, instance);
-                self.global.update(store, instance);
                 ReturnOutcome::Wasm
             }
             None => ReturnOutcome::Host,

crates/wasmi/src/engine/executor/mod.rs

@@ -215,12 +215,14 @@
                     unsafe { uninit_params.init_next(value) };
                 }
                 uninit_params.init_zeroes();
-                self.stack.calls.push(CallFrame::new(
-                    InstructionPtr::new(compiled_func.instrs().as_ptr()),
-                    offsets,
-                    RegisterSpan::new(Register::from_i16(0)),
-                    instance,
-                ))?;
+                self.stack.calls.push(
+                    CallFrame::new(
+                        InstructionPtr::new(compiled_func.instrs().as_ptr()),
+                        offsets,
+                        RegisterSpan::new(Register::from_i16(0)),
+                    ),
+                    Some(instance),
+                )?;
                 self.execute_func(store)?;
             }
             FuncEntity::Host(host_func) => {