winch: Properly derive a scratch register for arg assignment (bytecod…

…ealliance#7501)

This commit properly derives a scratch register for a particular
WebAssembly type. The included spec test uncovered that the previous
implementation used a int scratch register to assign float stack
arguments, which resulted in a panic.

tests/misc_testsuite/winch/call.wast

@@ -144,6 +144,13 @@
     (block (result i32) (i32.const 2) (call $const-i32) (br_table 0 0))
   )
 
+  (func $return-from-long-argument-list-helper (param f32 i32 i32 f64 f32 f32 f32 f64 f32 i32 i32 f32 f64 i64 i64 i32 i64 i64 f32 i64 i64 i64 i32 f32 f32 f32 f64 f32 i32 i64 f32 f64 f64 f32 i32 f32 f32 f64 i64 f64 i32 i64 f32 f64 i32 i32 i32 i64 f64 i32 i64 i64 f64 f64 f64 f64 f64 f64 i32 f32 f64 f64 i32 i64 f32 f32 f32 i32 f64 f64 f64 f64 f64 f32 i64 i64 i32 i32 i32 f32 f64 i32 i64 f32 f32 f32 i32 i32 f32 f64 i64 f32 f64 f32 f32 f32 i32 f32 i64 i32) (result i32)
+    (local.get 99)
+  )
+
+  (func (export "return-from-long-argument-list") (param i32) (result i32)
+    (call $return-from-long-argument-list-helper (f32.const 0) (i32.const 0) (i32.const 0) (f64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (f64.const 0) (f32.const 0) (i32.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (i64.const 0) (i64.const 0) (i32.const 0) (i64.const 0) (i64.const 0) (f32.const 0) (i64.const 0) (i64.const 0) (i64.const 0) (i32.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (f64.const 0) (f32.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f64.const 0) (f64.const 0) (f32.const 0) (i32.const 0) (f32.const 0) (f32.const 0) (f64.const 0) (i64.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f64.const 0) (i32.const 0) (i32.const 0) (i32.const 0) (i64.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (i64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (i32.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f32.const 0) (i64.const 0) (i64.const 0) (i32.const 0) (i32.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (i32.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (i64.const 0) (f32.const 0) (f64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (i32.const 0) (f32.const 0) (i64.const 0) (local.get 0))
+  )
 )
 
 (assert_return (invoke "type-i32") (i32.const 0x132))
@@ -193,3 +200,4 @@
 
 (assert_return (invoke "as-br_table-first") (i32.const 0x132))
 (assert_return (invoke "as-br_table-last") (i32.const 2))
+(assert_return (invoke "return-from-long-argument-list" (i32.const 42)) (i32.const 42))

winch/codegen/src/abi/mod.rs

@@ -95,6 +95,15 @@ pub(crate) trait ABI {
     /// Returns the designated floating point scratch register.
     fn float_scratch_reg() -> Reg;
 
+    /// Returns the designated scratch register for the given [WasmType].
+    fn scratch_for(ty: &WasmType) -> Reg {
+        match ty {
+            WasmType::I32 | WasmType::I64 => Self::scratch_reg(),
+            WasmType::F32 | WasmType::F64 => Self::float_scratch_reg(),
+            _ => unimplemented!(),
+        }
+    }
+
     /// Returns the frame pointer register.
     fn fp_reg() -> Reg;
 

winch/codegen/src/codegen/call.rs

@@ -99,8 +99,7 @@ impl FnCall {
         let call_stack_space = Self::save(context, masm, &sig);
 
         let reserved_stack = masm.call(arg_stack_space, |masm| {
-            let scratch = <M::ABI as ABI>::scratch_reg();
-            Self::assign(sig, context, masm, scratch);
+            Self::assign(sig, context, masm);
             kind
         });
 
@@ -247,15 +246,11 @@ impl FnCall {
     }
 
     /// Assign arguments for the function call.
-    fn assign<M: MacroAssembler>(
-        sig: &ABISig,
-        context: &mut CodeGenContext,
-        masm: &mut M,
-        scratch: Reg,
-    ) {
+    fn assign<M: MacroAssembler>(sig: &ABISig, context: &mut CodeGenContext, masm: &mut M) {
         let arg_count = sig.params.len();
         let stack = &context.stack;
         let mut stack_values = stack.peekn(arg_count);
+
         for arg in &sig.params {
             let val = stack_values
                 .next()
@@ -267,6 +262,7 @@ impl FnCall {
                 &ABIArg::Stack { ty, offset } => {
                     let addr = masm.address_at_sp(*offset);
                     let size: OperandSize = (*ty).into();
+                    let scratch = <M::ABI as ABI>::scratch_for(&ty);
                     context.move_val_to_reg(val, scratch, masm);
                     masm.store(scratch.into(), addr, size);
                 }

winch/codegen/src/codegen/context.rs

@@ -454,11 +454,7 @@ impl<'a, 'builtins> CodeGenContext<'a, 'builtins> {
             Val::Local(local) => {
                 let slot = frame.get_local(local.index).expect("valid local at slot");
                 let addr = masm.local_address(&slot);
-                let scratch = match slot.ty {
-                    WasmType::I32 | WasmType::I64 => <M::ABI as ABI>::scratch_reg(),
-                    WasmType::F32 | WasmType::F64 => <M::ABI as ABI>::float_scratch_reg(),
-                    WasmType::V128 | WasmType::Ref(_) => unimplemented!(),
-                };
+                let scratch = <M::ABI as ABI>::scratch_for(&slot.ty);
                 masm.load(addr, scratch, slot.ty.into());
                 let stack_slot = masm.push(scratch, slot.ty.into());
                 *v = Val::mem(slot.ty, stack_slot);