Consolidate translation of wasmfx instructions in FuncEnvironment (by…

…tecodealliance#152)

This is just a minor cleanup, influencing where different parts of
translating wasmfx instructions happens.

For `resume` instructions, at some point we moved all of the translation
logic out of `code_translator.rs` and into `func_environ.rs`. This means
that the `Operator::Resume` case in `translate_operator` only
manipulates the value stack, and then delegates the actual translation
to `translate_resume` in `FuncEnvironment` (which then in turn delegates
to the corresponding translation functions in the baseline or optimized
implementation).

This PR does the same for `cont.bind`, `cont.new`, and `suspend`, where
some of the translation logic was scattered between `code_translator.rs`
and `func_environ.rs`. The previous situation had the downside of
requiring a bunch of `typed_continuations_*` functions (e.g.,
`typed_continuations_store_payloads`) to be part of the interface of
`FuncEnvironment`, and then both the baseline and optimized
implementation had to implement that same interface. This created some
unnecessary coupling between the two implementations.

After these changes, the case for any wasmfx instruction `foo` in
`code_translator.rs` only manipulates the value stack, and then calls
`translate_foo` in `FuncEnvironment`, which then calls `translate_foo`
in the baseline or optimized implementation. In particular, no
`typed_continuations_*` functions are required in the `FuncEnvironment`
trait itself anymore.

cranelift/wasm/src/code_translator.rs

@@ -2548,9 +2548,8 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let arg_types = environ.continuation_arguments(*type_index).to_vec();
             let result_types = environ.continuation_returns(*type_index).to_vec();
             let r = state.pop1();
-            let contref =
+            let contobj =
                 environ.translate_cont_new(builder, state, r, &arg_types, &result_types)?;
-            let contobj = environ.typed_continuations_new_cont_obj(builder, contref);
             state.push1(contobj);
         }
         Operator::Resume {
@@ -2583,22 +2582,16 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
         }
         Operator::Suspend { tag_index } => {
             let param_types = environ.tag_params(*tag_index).to_vec();
+            let return_types = environ.tag_returns(*tag_index).to_vec();
 
             let params = state.peekn(param_types.len());
             let param_count = params.len();
 
-            environ.typed_continuations_store_payloads(builder, &param_types, params);
-            state.popn(param_count);
-
             let tag_index_val = builder.ins().iconst(I32, *tag_index as i64);
-            environ.translate_suspend(builder, tag_index_val);
-
-            let contref = environ.typed_continuations_load_continuation_reference(builder);
-
-            let return_types = environ.tag_returns(*tag_index).to_vec();
             let return_values =
-                environ.typed_continuations_load_tag_return_values(builder, contref, &return_types);
+                environ.translate_suspend(builder, tag_index_val, params, &return_types);
 
+            state.popn(param_count);
             state.pushn(&return_values);
         }
         Operator::ContBind {
@@ -2610,13 +2603,9 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let arg_count = src_arity - dst_arity;
 
             let (original_contobj, args) = state.peekn(arg_count + 1).split_last().unwrap();
-            let contref =
-                environ.typed_continuations_cont_obj_get_cont_ref(builder, *original_contobj);
-
-            let src_arity_value = builder.ins().iconst(I32, src_arity as i64);
-            environ.typed_continuations_store_resume_args(builder, args, src_arity_value, contref);
 
-            let new_contobj = environ.typed_continuations_new_cont_obj(builder, contref);
+            let new_contobj =
+                environ.translate_cont_bind(builder, *original_contobj, args, src_arity);
 
             state.popn(arg_count + 1);
             state.push1(new_contobj);

cranelift/wasm/src/environ/spec.rs

@@ -616,6 +616,18 @@ pub trait FuncEnvironment: TargetEnvironment {
         false
     }
 
+    /// Translates cont.bind instructions.
+    /// `remaining_arg_count` is the *overall* number of remaining arguments of
+    /// the continuation before supplying `args` (i.e., it is the number of
+    /// parameters of the continuation's type before cont.bind was executed).
+    fn translate_cont_bind(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        contobj: ir::Value,
+        args: &[ir::Value],
+        remaining_arg_count: usize,
+    ) -> ir::Value;
+
     /// TODO(dhil): write documentation.
     fn translate_cont_new(
         &mut self,
@@ -652,66 +664,22 @@ pub trait FuncEnvironment: TargetEnvironment {
         &mut self,
         builder: &mut FunctionBuilder,
         tag_index: ir::Value,
-    ) -> ir::Value;
+        suspend_args: &[ir::Value],
+        tag_return_types: &[wasmtime_types::WasmValType],
+    ) -> Vec<ir::Value>;
 
     /// TODO
     fn continuation_arguments(&self, type_index: u32) -> &[wasmtime_types::WasmValType];
 
     /// TODO
     fn continuation_returns(&self, type_index: u32) -> &[wasmtime_types::WasmValType];
 
-    /// TODO
-    fn typed_continuations_load_tag_return_values(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        contref: ir::Value,
-        valtypes: &[wasmtime_types::WasmValType],
-    ) -> std::vec::Vec<ir::Value>;
-
-    /// TODO
-    fn typed_continuations_store_payloads(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        valtypes: &[wasmtime_types::WasmValType],
-        values: &[ir::Value],
-    );
-
-    /// TODO
-    fn typed_continuations_store_resume_args(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        values: &[ir::Value],
-        remaining_arg_count: ir::Value,
-        contref: ir::Value,
-    );
-
     /// TODO
     fn tag_params(&self, tag_index: u32) -> &[wasmtime_types::WasmValType];
 
     /// TODO
     fn tag_returns(&self, tag_index: u32) -> &[wasmtime_types::WasmValType];
 
-    /// Returns a pointer to the currently running continuation reference.
-    /// Traps if not currently running inside a continuation.
-    fn typed_continuations_load_continuation_reference(
-        &mut self,
-        builder: &mut FunctionBuilder,
-    ) -> ir::Value;
-
-    /// TODO
-    fn typed_continuations_new_cont_obj(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        contref_addr: ir::Value,
-    ) -> ir::Value;
-
-    /// TODO
-    fn typed_continuations_cont_obj_get_cont_ref(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        contobj: ir::Value,
-    ) -> ir::Value;
-
     /// Returns whether the CLIF `x86_blendv` instruction should be used for the
     /// relaxed simd `*.relaxed_laneselect` instruction for the specified type.
     fn use_x86_blendv_for_relaxed_laneselect(&self, ty: Type) -> bool {

crates/cranelift/src/func_environ.rs

@@ -2584,6 +2584,16 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         self.isa.triple().architecture == target_lexicon::Architecture::X86_64
     }
 
+    fn translate_cont_bind(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        contobj: ir::Value,
+        args: &[ir::Value],
+        remaining_arg_count: usize,
+    ) -> ir::Value {
+        wasmfx_impl::translate_cont_bind(self, builder, contobj, args, remaining_arg_count)
+    }
+
     fn translate_cont_new(
         &mut self,
         builder: &mut FunctionBuilder,
@@ -2624,8 +2634,10 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         &mut self,
         builder: &mut FunctionBuilder,
         tag_index: ir::Value,
-    ) -> ir::Value {
-        wasmfx_impl::translate_suspend(self, builder, tag_index)
+        suspend_args: &[ir::Value],
+        tag_return_types: &[WasmValType],
+    ) -> Vec<ir::Value> {
+        wasmfx_impl::translate_suspend(self, builder, tag_index, suspend_args, tag_return_types)
     }
 
     fn continuation_arguments(&self, index: u32) -> &[WasmValType] {
@@ -2648,66 +2660,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         self.types[idx].returns()
     }
 
-    fn typed_continuations_load_tag_return_values(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        contref: ir::Value,
-        valtypes: &[WasmValType],
-    ) -> Vec<ir::Value> {
-        wasmfx_impl::typed_continuations_load_tag_return_values(self, builder, contref, valtypes)
-    }
-
-    /// TODO
-    fn typed_continuations_cont_obj_get_cont_ref(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        contobj: ir::Value,
-    ) -> ir::Value {
-        wasmfx_impl::typed_continuations_cont_obj_get_cont_ref(self, builder, contobj)
-    }
-
-    /// TODO
-    fn typed_continuations_store_resume_args(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        values: &[ir::Value],
-        remaining_arg_count: ir::Value,
-        contref: ir::Value,
-    ) {
-        wasmfx_impl::typed_continuations_store_resume_args(
-            self,
-            builder,
-            values,
-            remaining_arg_count,
-            contref,
-        )
-    }
-
-    //TODO(frank-emrich) Consider removing `valtypes` argument, as values are inherently typed
-    fn typed_continuations_store_payloads(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        valtypes: &[WasmValType],
-        values: &[ir::Value],
-    ) {
-        wasmfx_impl::typed_continuations_store_payloads(self, builder, valtypes, values)
-    }
-
-    fn typed_continuations_load_continuation_reference(
-        &mut self,
-        builder: &mut FunctionBuilder,
-    ) -> ir::Value {
-        wasmfx_impl::typed_continuations_load_continuation_reference(self, builder)
-    }
-
-    fn typed_continuations_new_cont_obj(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        contref_addr: ir::Value,
-    ) -> ir::Value {
-        wasmfx_impl::typed_continuations_new_cont_obj(self, builder, contref_addr)
-    }
-
     fn use_x86_blendv_for_relaxed_laneselect(&self, ty: Type) -> bool {
         self.isa.has_x86_blendv_lowering(ty)
     }

crates/cranelift/src/wasmfx/baseline.rs

@@ -8,13 +8,10 @@ use cranelift_codegen::ir::InstBuilder;
 use cranelift_frontend::{FunctionBuilder, Switch};
 use cranelift_wasm::FuncEnvironment;
 use cranelift_wasm::{FuncTranslationState, WasmResult, WasmValType};
+use shared::typed_continuations_cont_obj_get_cont_ref;
+use shared::typed_continuations_new_cont_obj;
 use wasmtime_environ::PtrSize;
 
-#[allow(unused_imports)]
-pub(crate) use shared::typed_continuations_cont_obj_get_cont_ref;
-#[allow(unused_imports)]
-pub(crate) use shared::typed_continuations_new_cont_obj;
-
 fn typed_continuations_load_payloads<'a>(
     env: &mut crate::func_environ::FuncEnvironment<'a>,
     builder: &mut FunctionBuilder,
@@ -122,11 +119,9 @@ pub(crate) fn typed_continuations_store_resume_args<'a>(
 pub(crate) fn typed_continuations_store_payloads<'a>(
     env: &mut crate::func_environ::FuncEnvironment<'a>,
     builder: &mut FunctionBuilder,
-    valtypes: &[WasmValType],
     values: &[ir::Value],
 ) {
-    assert_eq!(values.len(), valtypes.len());
-    if valtypes.len() > 0 {
+    if values.len() > 0 {
         // Retrieve the pointer to the payloads buffer.
         let nargs = builder.ins().iconst(I64, values.len() as i64);
         call_builtin!(builder, env, let payloads_ptr = tc_baseline_get_payloads_ptr(nargs));
@@ -376,6 +371,20 @@ pub(crate) fn translate_resume<'a>(
     }
 }
 
+pub(crate) fn translate_cont_bind<'a>(
+    env: &mut crate::func_environ::FuncEnvironment<'a>,
+    builder: &mut FunctionBuilder,
+    contobj: ir::Value,
+    args: &[ir::Value],
+    remaining_arg_count: usize,
+) -> ir::Value {
+    let contref = typed_continuations_cont_obj_get_cont_ref(env, builder, contobj);
+    let remaining_arg_count = builder.ins().iconst(I32, remaining_arg_count as i64);
+    typed_continuations_store_resume_args(env, builder, args, remaining_arg_count, contref);
+
+    typed_continuations_new_cont_obj(env, builder, contref)
+}
+
 pub(crate) fn translate_cont_new<'a>(
     env: &mut crate::func_environ::FuncEnvironment<'a>,
     builder: &mut FunctionBuilder,
@@ -388,15 +397,23 @@ pub(crate) fn translate_cont_new<'a>(
     let nargs = builder.ins().iconst(I64, arg_types.len() as i64);
     let nreturns = builder.ins().iconst(I64, return_types.len() as i64);
     call_builtin!(builder, env, let contref = tc_baseline_cont_new(func, nargs, nreturns));
-
-    Ok(contref)
+    let contobj = typed_continuations_new_cont_obj(env, builder, contref);
+    Ok(contobj)
 }
 
 pub(crate) fn translate_suspend<'a>(
     env: &mut crate::func_environ::FuncEnvironment<'a>,
     builder: &mut FunctionBuilder,
     tag_index: ir::Value,
-) -> ir::Value {
+    suspend_args: &[ir::Value],
+    tag_return_types: &[WasmValType],
+) -> Vec<ir::Value> {
+    typed_continuations_store_payloads(env, builder, suspend_args);
     call_builtin!(builder, env, tc_baseline_suspend(tag_index));
-    return env.vmctx_val(&mut builder.cursor());
+    let contref = typed_continuations_load_continuation_reference(env, builder);
+
+    let return_values =
+        typed_continuations_load_tag_return_values(env, builder, contref, tag_return_types);
+
+    return_values
 }