winch(x64): Add support for table instructions (bytecodealliance#7155)

* winch(x64): Add support for table instructions

This change adds support for the following table insructions:
`elem.drop`, `table.copy`, `table.set`, `table.get`, `table.fill`,
`table.grow`, `table.size`, `table.init`.

This change also introduces partial support for the `Ref` WebAssembly
type, more conretely the `Func` heap type, which means that all the
table instructions above, only work this WebAssembly type as of this
change.

Finally, this change is also a small follow up to the primitives
introduced in bytecodealliance#7100,
more concretely:

* `FnCall::with_lib`: tracks the presence of a libcall and ensures that
  any result registers are freed right when the call is emitted.
* `MacroAssembler::table_elem_addr` returns an address rather than the
  value of the address, making it convenient for other use cases like
  `table.set`.

--

prtest:full

* chore: Make stack functions take impl IntoIterator<..>

* Update winch/codegen/src/codegen/call.rs

Co-authored-by: Trevor Elliott <awesomelyawesome@gmail.com>

* Remove a dangling `dbg!`

* Add comment on branching

---------

Co-authored-by: Trevor Elliott <awesomelyawesome@gmail.com>

build.rs

@@ -205,19 +205,34 @@ fn ignore(testsuite: &str, testname: &str, strategy: &str) -> bool {
     // We ignore tests that assert for traps on windows, given
     // that Winch doesn't encode unwind information for Windows, yet.
     if strategy == "Winch" {
+        let assert_trap = [
+            "i32",
+            "i64",
+            "call_indirect",
+            "table_fill",
+            "table_init",
+            "table_copy",
+            "table_set",
+            "table_get",
+        ]
+        .contains(&testname);
+
+        if assert_trap && env::var("CARGO_CFG_TARGET_OS").unwrap().as_str() == "windows" {
+            return true;
+        }
+
         if testsuite == "misc_testsuite" {
             // The misc/call_indirect is fully supported by Winch.
-            if testname == "call_indirect" {
-                return false;
+            if testname != "call_indirect" {
+                return true;
             }
         }
-        if testsuite != "winch" {
-            return true;
+        if testsuite == "spec_testsuite" {
+            // The official table init and table copy tests are now supported.
+            return !["table_init", "table_copy"].contains(&testname);
         }
 
-        let assert_trap = ["i32", "i64", "call_indirect"].contains(&testname);
-
-        if assert_trap && env::var("CARGO_CFG_TARGET_OS").unwrap().as_str() == "windows" {
+        if testsuite != "winch" {
             return true;
         }
     }

fuzz/fuzz_targets/differential.rs

@@ -375,7 +375,15 @@ fn winch_supports_module(module: &[u8]) -> bool {
                         | F64Abs { .. }
                         | F32Neg { .. }
                         | F64Neg { .. }
-                        | CallIndirect { .. } => {}
+                        | CallIndirect { .. }
+                        | ElemDrop { .. }
+                        | TableCopy { .. }
+                        | TableSet { .. }
+                        | TableGet { .. }
+                        | TableFill { .. }
+                        | TableGrow { .. }
+                        | TableSize { .. }
+                        | TableInit { .. } => {}
                         _ => {
                             supported = false;
                             break 'main;

tests/misc_testsuite/winch/table_fill.wast

@@ -0,0 +1,65 @@
+(module
+  (type $t0 (func))
+  (func $f1 (type $t0))
+  (func $f2 (type $t0))
+  (func $f3 (type $t0))
+
+  ;; Define two tables of funcref
+  (table $t1 3 funcref)
+  (table $t2 10 funcref)
+
+  ;; Initialize table $t1 with functions $f1, $f2, $f3
+  (elem (i32.const 0) $f1 $f2 $f3)
+
+  ;; Function to fill table $t1 using a function reference from table $t2
+  (func (export "fill") (param $i i32) (param $r i32) (param $n i32)
+    (local $ref funcref)
+    (local.set $ref (table.get $t1 (local.get $r)))
+    (table.fill $t2 (local.get $i) (local.get $ref) (local.get $n))
+  )
+
+  (func (export "get") (param $i i32) (result funcref)
+    (table.get $t2 (local.get $i))
+  )
+)
+
+(assert_return (invoke "get" (i32.const 1)) (ref.null func))
+(assert_return (invoke "get" (i32.const 2)) (ref.null func))
+(assert_return (invoke "get" (i32.const 3)) (ref.null func))
+(assert_return (invoke "get" (i32.const 4)) (ref.null func))
+(assert_return (invoke "get" (i32.const 5)) (ref.null func))
+
+(assert_return (invoke "fill" (i32.const 2) (i32.const 0) (i32.const 3)))
+(assert_return (invoke "get" (i32.const 1)) (ref.null func))
+(assert_return (invoke "get" (i32.const 2)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 3)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 4)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 5)) (ref.null func))
+
+(assert_return (invoke "fill" (i32.const 4) (i32.const 1) (i32.const 2)))
+(assert_return (invoke "get" (i32.const 3)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 4)) (ref.func 1))
+(assert_return (invoke "get" (i32.const 5)) (ref.func 1))
+(assert_return (invoke "get" (i32.const 6)) (ref.null func))
+
+(assert_return (invoke "fill" (i32.const 4) (i32.const 2) (i32.const 0)))
+(assert_return (invoke "get" (i32.const 3)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 4)) (ref.func 1))
+(assert_return (invoke "get" (i32.const 5)) (ref.func 1))
+
+(assert_return (invoke "fill" (i32.const 8) (i32.const 0) (i32.const 2)))
+(assert_return (invoke "get" (i32.const 7)) (ref.null func))
+(assert_return (invoke "get" (i32.const 8)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 9)) (ref.func 0))
+
+(assert_return (invoke "fill" (i32.const 9) (i32.const 2) (i32.const 1)))
+(assert_return (invoke "get" (i32.const 8)) (ref.func 0))
+(assert_return (invoke "get" (i32.const 9)) (ref.func 2))
+
+(assert_return (invoke "fill" (i32.const 10) (i32.const 1) (i32.const 0)))
+(assert_return (invoke "get" (i32.const 9)) (ref.func 2))
+
+(assert_trap
+  (invoke "fill" (i32.const 8) (i32.const 0) (i32.const 3))
+  "out of bounds table access"
+)

tests/misc_testsuite/winch/table_get.wast

@@ -0,0 +1,13 @@
+(module
+  (table $t3 3 funcref)
+  (elem (table $t3) (i32.const 1) func $dummy)
+  (func $dummy)
+  (func $f3 (export "get-funcref") (param $i i32) (result funcref)
+    (table.get $t3 (local.get $i))
+  )
+)
+
+(assert_return (invoke "get-funcref" (i32.const 0)) (ref.null func))
+(assert_trap (invoke "get-funcref" (i32.const 3)) "out of bounds table access")
+(assert_trap (invoke "get-funcref" (i32.const -1)) "out of bounds table access")
+

tests/misc_testsuite/winch/table_grow.wast

@@ -0,0 +1,26 @@
+(module
+  (table $t1 0 funcref)
+
+  (func (export "grow-by-10") (param $r funcref) (result i32)
+    (table.grow $t1 (local.get $r) (i32.const 10))
+  )
+  (func (export "grow-over") (param $r funcref) (result i32)
+    (table.grow $t1 (local.get $r) (i32.const 0xffff_fff0))
+  )
+
+  (func (export "size") (result i32)
+    (table.size $t1))
+)
+
+(assert_return (invoke "size") (i32.const 0))
+(assert_return (invoke "grow-by-10" (ref.null func)) (i32.const 0))
+(assert_return (invoke "size") (i32.const 10))
+
+(module
+  (table $t 0x10 funcref)
+  (func $f (export "grow") (param $r funcref) (result i32)
+    (table.grow $t (local.get $r) (i32.const 0xffff_fff0))
+  )
+)
+
+(assert_return (invoke "grow" (ref.null func)) (i32.const -1))

tests/misc_testsuite/winch/table_set.wast

@@ -0,0 +1,27 @@
+(module
+  (table $t3 2 funcref)
+  (elem (table $t3) (i32.const 1) func $dummy)
+  (func $dummy)
+
+  (func $f3 (export "get-funcref") (param $i i32) (result funcref)
+    (table.get $t3 (local.get $i))
+  )
+
+  (func (export "set-funcref") (param $i i32) (param $r funcref)
+    (table.set $t3 (local.get $i) (local.get $r))
+  )
+  (func (export "set-funcref-from") (param $i i32) (param $j i32)
+    (table.set $t3 (local.get $i) (table.get $t3 (local.get $j)))
+  )
+)
+
+(assert_return (invoke "get-funcref" (i32.const 0)) (ref.null func))
+(assert_return (invoke "set-funcref-from" (i32.const 0) (i32.const 1)))
+(assert_return (invoke "set-funcref" (i32.const 0) (ref.null func)))
+(assert_return (invoke "get-funcref" (i32.const 0)) (ref.null func))
+
+(assert_trap (invoke "set-funcref" (i32.const 3) (ref.null func)) "out of bounds table access")
+(assert_trap (invoke "set-funcref" (i32.const -1) (ref.null func)) "out of bounds table access")
+
+(assert_trap (invoke "set-funcref-from" (i32.const 3) (i32.const 1)) "out of bounds table access")
+(assert_trap (invoke "set-funcref-from" (i32.const -1) (i32.const 1)) "out of bounds table access")

winch/codegen/src/abi/mod.rs

@@ -47,7 +47,7 @@ use crate::isa::{reg::Reg, CallingConvention};
 use crate::masm::OperandSize;
 use smallvec::SmallVec;
 use std::ops::{Add, BitAnd, Not, Sub};
-use wasmtime_environ::{WasmFuncType, WasmType};
+use wasmtime_environ::{WasmFuncType, WasmHeapType, WasmType};
 
 pub(crate) mod local;
 pub(crate) use local::*;
@@ -267,7 +267,15 @@ pub(crate) fn ty_size(ty: &WasmType) -> u32 {
     match *ty {
         WasmType::I32 | WasmType::F32 => 4,
         WasmType::I64 | WasmType::F64 => 8,
-        _ => panic!(),
+        WasmType::Ref(rt) => match rt.heap_type {
+            // TODO: Similar to the comment in visitor.rs at impl From<WasmType> for
+            // OperandSize, Once Wasmtime supports 32-bit architectures, this will
+            // need to be updated to derive operand size from the target's pointer
+            // size.
+            WasmHeapType::Func => 8,
+            ht => unimplemented!("Support for WasmHeapType: {ht}"),
+        },
+        t => unimplemented!("Support for WasmType: {t}"),
     }
 }
 

winch/codegen/src/codegen/call.rs

@@ -1,7 +1,7 @@
 //! Function call emission.  For more details around the ABI and
 //! calling convention, see [ABI].
 use crate::{
-    abi::{ABIArg, ABISig, ABI},
+    abi::{ABIArg, ABIResult, ABISig, ABI},
     codegen::{BuiltinFunction, CodeGenContext},
     masm::{CalleeKind, MacroAssembler, OperandSize},
     reg::Reg,
@@ -65,6 +65,8 @@ pub(crate) struct FnCall<'a> {
     arg_stack_space: u32,
     /// The ABI-specific signature of the callee.
     pub abi_sig: &'a ABISig,
+    /// Whether this a built-in function call.
+    lib: bool,
 }
 
 impl<'a> FnCall<'a> {
@@ -74,6 +76,7 @@ impl<'a> FnCall<'a> {
             abi_sig: &callee_sig,
             arg_stack_space: callee_sig.stack_bytes,
             call_stack_space: None,
+            lib: false,
         }
     }
 
@@ -238,6 +241,7 @@ impl<'a> FnCall<'a> {
     ) where
         F: FnMut(&mut CodeGenContext, &mut M, &mut Self, Reg),
     {
+        self.lib = true;
         // When dealing with libcalls, we don't have all the information
         // upfront (all necessary arguments in the stack) in order to optimize
         // saving the live registers, so we save all the values available in
@@ -288,6 +292,26 @@ impl<'a> FnCall<'a> {
                 regalloc.free(v.get_reg().into());
             }
         });
+
+        // When emitting built-calls we ensure that none of the registers
+        // (params and results) used as part of the ABI signature are
+        // allocatable throughout the lifetime of the `with_lib` callback, since
+        // such registers will be used to assign arguments and hold results.
+        // After executing the callback, it's only safe to free the param
+        // registers, since depending on the signature, the caller
+        // will push any result registers to the stack, keeping those registers allocated.
+        // Here we ensure that any allocated result registers are correctly
+        // freed before finalizing the function call and pushing any results to
+        // the value stack.
+        if self.lib {
+            match self.abi_sig.result {
+                ABIResult::Reg { reg, .. } => {
+                    assert!(!context.regalloc.reg_available(reg));
+                    context.free_reg(reg);
+                }
+                _ => {}
+            }
+        }
         context.push_abi_results(&self.abi_sig.result, masm);
     }
 

winch/codegen/src/codegen/context.rs

@@ -1,4 +1,4 @@
-use wasmtime_environ::WasmType;
+use wasmtime_environ::{WasmHeapType, WasmType};
 
 use super::ControlStackFrame;
 use crate::{
@@ -63,7 +63,11 @@ impl<'a> CodeGenContext<'a> {
         match ty {
             I32 | I64 => self.reg_for_class(RegClass::Int, masm),
             F32 | F64 => self.reg_for_class(RegClass::Float, masm),
-            t => panic!("unsupported type {:?}", t),
+            Ref(rt) => match rt.heap_type {
+                WasmHeapType::Func => self.reg_for_class(RegClass::Int, masm),
+                ht => unimplemented!("Support for WasmHeapType: {ht}"),
+            },
+            t => unimplemented!("Support for WasmType: {t}"),
         }
     }