implement host function calls in return_call[_indirect]

Author: Robbepop

crates/wasmi/src/engine/executor/handler/exec.rs

@@ -34,6 +34,7 @@ use crate::{
                 resolve_memory,
                 resolve_table,
                 resolve_table_mut,
+                return_call_host,
                 return_call_wasm,
                 set_global,
                 update_instance,
@@ -285,22 +286,23 @@ pub fn return_call_imported(
 ) -> Done {
     let (_, crate::ir::decode::ReturnCallImported { params, func }) = unsafe { decode_op(ip) };
     let func = fetch_func(instance, func);
-    let func = resolve_func(state.store, &func);
-    let (callee_ip, sp, mem0, mem0_len, instance) = match func {
+    let func_entity = resolve_func(state.store, &func);
+    let (callee_ip, sp, new_instance) = match func_entity {
         FuncEntity::Wasm(func) => {
             let wasm_func = func.func_body();
             let callee_instance = *func.instance();
             let callee_instance = resolve_instance(state.store, &callee_instance).into();
             let (callee_ip, callee_sp) =
                 return_call_wasm(state, params, wasm_func, Some(instance))?;
-            let (instance, mem0, mem0_len) =
-                update_instance(state.store, instance, callee_instance, mem0, mem0_len);
-            (callee_ip, callee_sp, mem0, mem0_len, instance)
+            (callee_ip, callee_sp, callee_instance)
         }
-        FuncEntity::Host(_host_func) => {
-            todo!()
+        FuncEntity::Host(host_func) => {
+            let host_func = *host_func;
+            return_call_host(state, func, host_func, params, instance)?
         }
     };
+    let (instance, mem0, mem0_len) =
+        update_instance(state.store, instance, new_instance, mem0, mem0_len);
     dispatch!(state, callee_ip, sp, mem0, mem0_len, instance)
 }
 
@@ -323,22 +325,23 @@ pub fn return_call_indirect(
     ) = unsafe { decode_op(ip) };
     let func =
         resolve_indirect_func(index, table, func_type, state, sp, instance).into_control()?;
-    let func = resolve_func(state.store, &func);
-    let (callee_ip, sp, mem0, mem0_len, instance) = match func {
+    let func_entity = resolve_func(state.store, &func);
+    let (callee_ip, sp, callee_instance) = match func_entity {
         FuncEntity::Wasm(func) => {
             let wasm_func = func.func_body();
             let callee_instance = *func.instance();
             let callee_instance: Inst = resolve_instance(state.store, &callee_instance).into();
             let (callee_ip, callee_sp) =
                 return_call_wasm(state, params, wasm_func, Some(instance))?;
-            let (instance, mem0, mem0_len) =
-                update_instance(state.store, instance, callee_instance, mem0, mem0_len);
-            (callee_ip, callee_sp, mem0, mem0_len, instance)
+            (callee_ip, callee_sp, callee_instance)
         }
-        FuncEntity::Host(_func) => {
-            todo!()
+        FuncEntity::Host(host_func) => {
+            let host_func = *host_func;
+            return_call_host(state, func, host_func, params, instance)?
         }
     };
+    let (instance, mem0, mem0_len) =
+        update_instance(state.store, instance, callee_instance, mem0, mem0_len);
     dispatch!(state, callee_ip, sp, mem0, mem0_len, instance)
 }
 

crates/wasmi/src/engine/executor/handler/state.rs

@@ -2,7 +2,10 @@ use crate::{
     core::{ReadAs, UntypedVal, WriteAs},
     engine::{
         executor::{
-            handler::{dispatch::ExecutionOutcome, utils::extract_mem0},
+            handler::{
+                dispatch::{Control, ExecutionOutcome},
+                utils::extract_mem0,
+            },
             CodeMap,
         },
         utils::unreachable_unchecked,
@@ -252,6 +255,8 @@ pub struct Stack {
     frames: CallStack,
 }
 
+type ReturnCallHost = Control<(Ip, Sp, Inst), Sp>;
+
 impl Stack {
     pub fn new(config: &StackConfig) -> Self {
         Self {
@@ -276,6 +281,17 @@ impl Stack {
         self.values.capacity()
     }
 
+    pub fn return_prepare_host_frame<'a>(
+        &'a mut self,
+        callee_params: BoundedSlotSpan,
+        results_len: u16,
+        caller_instance: Inst,
+    ) -> Result<(ReturnCallHost, FuncInOut<'a>), TrapCode> {
+        let (callee_start, caller) = self.frames.return_prepare_host_frame(caller_instance);
+        self.values
+            .return_prepare_host_frame(caller, callee_start, callee_params, results_len)
+    }
+
     pub fn prepare_host_frame<'a>(
         &'a mut self,
         caller_ip: Option<Ip>,
@@ -376,6 +392,57 @@ impl ValueStack {
         }
     }
 
+    /// # Note
+    ///
+    /// In the following code, `callee` represents the called host function frame
+    /// and `caller` represents the caller of the caller of the host function, a.k.a.
+    /// the caller's caller.
+    fn return_prepare_host_frame<'a>(
+        &'a mut self,
+        caller: Option<(Ip, usize, Inst)>,
+        callee_start: usize,
+        callee_params: BoundedSlotSpan,
+        results_len: u16,
+    ) -> Result<(ReturnCallHost, FuncInOut<'a>), TrapCode> {
+        let caller_start = caller.map(|(_, start, _)| start).unwrap_or(0);
+        let params_offset = usize::from(u16::from(callee_params.span().head()));
+        let params_len = usize::from(callee_params.len());
+        let results_len = usize::from(results_len);
+        let callee_size = params_len.max(results_len);
+        if callee_size == 0 {
+            let sp = match caller {
+                Some(_) if caller_start != callee_start => self.sp(caller_start),
+                _ => Sp::dangling(),
+            };
+            let inout = FuncInOut::new(&mut [], 0, 0);
+            let control = match caller {
+                Some((ip, _, instance)) => ReturnCallHost::Continue((ip, sp, instance)),
+                None => ReturnCallHost::Break(sp),
+            };
+            return Ok((control, inout));
+        }
+        let Some(params_start) = callee_start.checked_add(params_offset) else {
+            return Err(TrapCode::StackOverflow);
+        };
+        let Some(params_end) = params_start.checked_add(params_len) else {
+            return Err(TrapCode::StackOverflow);
+        };
+        self.cells
+            .copy_within(params_start..params_end, callee_start);
+        let Some(callee_end) = callee_start.checked_add(callee_size) else {
+            return Err(TrapCode::StackOverflow);
+        };
+        self.cells.resize(callee_end, UntypedVal::default());
+        let caller_sp = self.sp(caller_start);
+        let cells = &mut self.cells[callee_start..];
+        let inout = FuncInOut::new(cells, params_len, results_len);
+        let control = match caller {
+            Some((ip, _, instance)) => ReturnCallHost::Continue((ip, caller_sp, instance)),
+            None => ReturnCallHost::Break(caller_sp),
+        };
+        Ok((control, inout))
+    }
+
     fn prepare_host_frame<'a>(
         &'a mut self,
         caller_start: usize,
@@ -504,6 +571,26 @@ impl CallStack {
         self.top_start()
     }
 
+    /// # Note
+    ///
+    /// In the following code, `callee` represents the called host function frame
+    /// and `caller` represents the caller of the caller of the host function, a.k.a.
+    /// the caller's caller.
+    pub fn return_prepare_host_frame(
+        &mut self,
+        callee_instance: Inst,
+    ) -> (usize, Option<(Ip, usize, Inst)>) {
+        let callee_start = self.top_start();
+        let caller = match self.pop() {
+            Some((ip, start, instance)) => {
+                let instance = instance.unwrap_or(callee_instance);
+                Some((ip, start, instance))
+            }
+            None => None,
+        };
+        (callee_start, caller)
+    }
+
     #[inline(always)]
     fn push(
         &mut self,

crates/wasmi/src/engine/executor/handler/utils.rs

@@ -429,6 +429,39 @@ pub fn call_host(
     Control::Continue(sp)
 }
 
+pub fn return_call_host(
+    state: &mut VmState,
+    func: Func,
+    host_func: HostFuncEntity,
+    params: BoundedSlotSpan,
+    instance: Inst,
+) -> Control<(Ip, Sp, Inst), Break> {
+    debug_assert_eq!(params.len(), host_func.len_params());
+    let trampoline = *host_func.trampoline();
+    let (control, params_results) = state
+        .stack
+        .return_prepare_host_frame(params, host_func.len_results(), instance)
+        .into_control()?;
+    match state
+        .store
+        .call_host_func(trampoline, Some(instance), params_results, CallHooks::Call)
+    {
+        Ok(()) => {}
+        Err(StoreError::External(error)) => {
+            done!(state, DoneReason::host_error(error, func, params.span()))
+        }
+        Err(StoreError::Internal(error)) => unsafe {
+            unreachable_unchecked!(
+                "internal interpreter error while executing host function: {error}"
+            )
+        },
+    }
+    match control {
+        Control::Continue((ip, sp, instance)) => Control::Continue((ip, sp, instance)),
+        Control::Break(sp) => done!(state, DoneReason::Return(sp)),
+    }
+}
+
 pub fn call_wasm_or_host(
     state: &mut VmState,
     caller_ip: Ip,