Use is_wasm_parameter in translating wasm calls (#1352)

* Use `is_wasm_parameter` in translating wasm calls

Added in #1329 it's now possible for multiple parameters to be non-wasm
parameters, so the previous `param_types` method is no longer suitable
for acquiring all wasm-related parameters, rather then `FuncEnvironment`
must be consulted. This removes usage of `param_types()` as a method
from the wasm translation and instead adds a custom method inline for
filtering the parameters based on `is_wasm_parameter`.

* Apply feedback

* Run rustfmt

* Don't require `mut`

* Run rustfmt

cranelift/codegen/src/ir/extfunc.rs

@@ -102,24 +102,6 @@ impl Signature {
             .count()
             > 1
     }
-
-    /// Collect the normal parameter types of the signature; see `[ArgumentPurpose::Normal]`.
-    pub fn param_types(&self) -> Vec<Type> {
-        self.params
-            .iter()
-            .filter(|ap| ap.purpose == ArgumentPurpose::Normal)
-            .map(|ap| ap.value_type)
-            .collect()
-    }
-
-    /// Collect the normal return types of the signature; see `[ArgumentPurpose::Normal]`.
-    pub fn return_types(&self) -> Vec<Type> {
-        self.returns
-            .iter()
-            .filter(|ap| ap.purpose == ArgumentPurpose::Normal)
-            .map(|ap| ap.value_type)
-            .collect()
-    }
 }
 
 /// Wrapper type capable of displaying a `Signature` with correct register names.

cranelift/wasm/src/code_translator.rs

@@ -38,6 +38,7 @@ use cranelift_codegen::ir::{
 };
 use cranelift_codegen::packed_option::ReservedValue;
 use cranelift_frontend::{FunctionBuilder, Variable};
+use std::vec::Vec;
 use wasmparser::{MemoryImmediate, Operator};
 
 // Clippy warns about "flags: _" but its important to document that the flags field is ignored
@@ -439,7 +440,9 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             };
             {
                 let return_args = state.peekn_mut(return_count);
-                let return_types = &builder.func.signature.return_types();
+                let return_types = wasm_param_types(&builder.func.signature.returns, |i| {
+                    environ.is_wasm_return(&builder.func.signature, i)
+                });
                 bitcast_arguments(return_args, &return_types, builder);
                 match environ.return_mode() {
                     ReturnMode::NormalReturns => builder.ins().return_(return_args),
@@ -463,7 +466,10 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let callee_signature =
                 &builder.func.dfg.signatures[builder.func.dfg.ext_funcs[fref].signature];
             let args = state.peekn_mut(num_args);
-            bitcast_arguments(args, &callee_signature.param_types(), builder);
+            let types = wasm_param_types(&callee_signature.params, |i| {
+                environ.is_wasm_parameter(&callee_signature, i)
+            });
+            bitcast_arguments(args, &types, builder);
 
             let call = environ.translate_call(
                 builder.cursor(),
@@ -492,7 +498,10 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             // Bitcast any vector arguments to their default type, I8X16, before calling.
             let callee_signature = &builder.func.dfg.signatures[sigref];
             let args = state.peekn_mut(num_args);
-            bitcast_arguments(args, &callee_signature.param_types(), builder);
+            let types = wasm_param_types(&callee_signature.params, |i| {
+                environ.is_wasm_parameter(&callee_signature, i)
+            });
+            bitcast_arguments(args, &types, builder);
 
             let call = environ.translate_call_indirect(
                 builder.cursor(),
@@ -1932,3 +1941,16 @@ pub fn bitcast_arguments(
         }
     }
 }
+
+/// A helper to extract all the `Type` listings of each variable in `params`
+/// for only parameters the return true for `is_wasm`, typically paired with
+/// `is_wasm_return` or `is_wasm_parameter`.
+pub fn wasm_param_types(params: &[ir::AbiParam], is_wasm: impl Fn(usize) -> bool) -> Vec<Type> {
+    let mut ret = Vec::with_capacity(params.len());
+    for (i, param) in params.iter().enumerate() {
+        if is_wasm(i) {
+            ret.push(param.value_type);
+        }
+    }
+    return ret;
+}

cranelift/wasm/src/environ/spec.rs

@@ -144,6 +144,12 @@ pub trait FuncEnvironment: TargetEnvironment {
         signature.params[index].purpose == ir::ArgumentPurpose::Normal
     }
 
+    /// Is the given return of the given function a wasm-level parameter, as
+    /// opposed to a hidden parameter added for use by the implementation?
+    fn is_wasm_return(&self, signature: &ir::Signature, index: usize) -> bool {
+        signature.returns[index].purpose == ir::ArgumentPurpose::Normal
+    }
+
     /// Should the code be structured to use a single `fallthrough_return` instruction at the end
     /// of the function body, rather than `return` instructions as needed? This is used by VMs
     /// to append custom epilogues.

cranelift/wasm/src/func_translator.rs

@@ -4,7 +4,7 @@
 //! function to Cranelift IR guided by a `FuncEnvironment` which provides information about the
 //! WebAssembly module and the runtime environment.
 
-use crate::code_translator::{bitcast_arguments, translate_operator};
+use crate::code_translator::{bitcast_arguments, translate_operator, wasm_param_types};
 use crate::environ::{FuncEnvironment, ReturnMode, WasmResult};
 use crate::state::{FuncTranslationState, ModuleTranslationState};
 use crate::translation_utils::get_vmctx_value_label;
@@ -245,7 +245,9 @@ fn parse_function_body<FE: FuncEnvironment + ?Sized>(
         if !builder.is_unreachable() {
             match environ.return_mode() {
                 ReturnMode::NormalReturns => {
-                    let return_types = &builder.func.signature.return_types();
+                    let return_types = wasm_param_types(&builder.func.signature.returns, |i| {
+                        environ.is_wasm_return(&builder.func.signature, i)
+                    });
                     bitcast_arguments(&mut state.stack, &return_types, builder);
                     builder.ins().return_(&state.stack)
                 }