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

build.rs

@@ -207,15 +207,20 @@ fn ignore(testsuite: &str, testname: &str, strategy: &str) -> bool {
     if strategy == "Winch" {
         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 == "spec_testsuite" {
+            // The official table init and table copy tests are now supported.
+            return !["table_init", "table_copy"].contains(&testname);
+        }
         if testsuite != "winch" {
             return true;
         }
 
-        let assert_trap = ["i32", "i64", "call_indirect"].contains(&testname);
+        let assert_trap =
+            ["i32", "i64", "call_indirect", "table_init", "table_copy"].contains(&testname);
 
         if assert_trap && env::var("CARGO_CFG_TARGET_OS").unwrap().as_str() == "windows" {
             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 WasmHeatType: {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 the 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),
         }
     }
 

winch/codegen/src/codegen/env.rs

@@ -1,12 +1,20 @@
-use crate::{codegen::BuiltinFunctions, CallingConvention};
+use crate::{
+    codegen::{BuiltinFunctions, OperandSize},
+    CallingConvention,
+};
 use smallvec::{smallvec, SmallVec};
+use std::collections::{
+    hash_map::Entry::{Occupied, Vacant},
+    HashMap,
+};
 use wasmparser::BlockType;
 use wasmtime_environ::{
-    FuncIndex, GlobalIndex, ModuleTranslation, ModuleTypes, PtrSize, TableIndex, TypeConvert,
-    TypeIndex, VMOffsets, WasmFuncType, WasmType,
+    FuncIndex, GlobalIndex, ModuleTranslation, ModuleTypes, PtrSize, TableIndex, TablePlan,
+    TypeConvert, TypeIndex, VMOffsets, WasmFuncType, WasmType,
 };
 
 /// Table metadata.
+#[derive(Debug, Copy, Clone)]
 pub struct TableData {
     /// The offset to the base of the table.
     pub offset: u32,
@@ -15,8 +23,10 @@ pub struct TableData {
     /// If the table is imported, return the base
     /// offset of the `from` field in `VMTableImport`.
     pub base: Option<u32>,
-    /// The size of the table elements, in bytes.
-    pub element_size: u8,
+    /// The size of the table elements.
+    pub(crate) element_size: OperandSize,
+    /// The size of the current elements field.
+    pub(crate) current_elements_size: OperandSize,
 }
 
 /// A function callee.
@@ -53,6 +63,8 @@ pub struct FuncEnv<'a, P: PtrSize> {
     pub builtins: BuiltinFunctions,
     /// The module's function types.
     pub types: &'a ModuleTypes,
+    /// Track resolved table information.
+    resolved_tables: HashMap<TableIndex, TableData>,
 }
 
 pub fn ptr_type_from_ptr_size(size: u8) -> WasmType {
@@ -77,6 +89,7 @@ impl<'a, P: PtrSize> FuncEnv<'a, P> {
             translation,
             builtins: BuiltinFunctions::new(size, call_conv, builtins_base),
             types,
+            resolved_tables: HashMap::new(),
         }
     }
 
@@ -141,27 +154,40 @@ impl<'a, P: PtrSize> FuncEnv<'a, P> {
     }
 
     /// Returns the table information for the given table index.
-    pub fn resolve_table_data(&self, index: TableIndex) -> TableData {
-        let (from_offset, base_offset, current_elems_offset) =
-            match self.translation.module.defined_table_index(index) {
-                Some(defined) => (
-                    None,
-                    self.vmoffsets.vmctx_vmtable_definition_base(defined),
-                    self.vmoffsets
-                        .vmctx_vmtable_definition_current_elements(defined),
-                ),
-                None => (
-                    Some(self.vmoffsets.vmctx_vmtable_import_from(index)),
-                    self.vmoffsets.vmtable_definition_base().into(),
-                    self.vmoffsets.vmtable_definition_current_elements().into(),
-                ),
-            };
-
-        TableData {
-            base: from_offset,
-            offset: base_offset,
-            current_elems_offset,
-            element_size: self.vmoffsets.ptr.size(),
+    pub fn resolve_table_data(&mut self, index: TableIndex) -> TableData {
+        match self.resolved_tables.entry(index) {
+            Occupied(entry) => *entry.get(),
+            Vacant(entry) => {
+                let (from_offset, base_offset, current_elems_offset) =
+                    match self.translation.module.defined_table_index(index) {
+                        Some(defined) => (
+                            None,
+                            self.vmoffsets.vmctx_vmtable_definition_base(defined),
+                            self.vmoffsets
+                                .vmctx_vmtable_definition_current_elements(defined),
+                        ),
+                        None => (
+                            Some(self.vmoffsets.vmctx_vmtable_import_from(index)),
+                            self.vmoffsets.vmtable_definition_base().into(),
+                            self.vmoffsets.vmtable_definition_current_elements().into(),
+                        ),
+                    };
+
+                *entry.insert(TableData {
+                    base: from_offset,
+                    offset: base_offset,
+                    current_elems_offset,
+                    element_size: OperandSize::from_bytes(self.vmoffsets.ptr.size()),
+                    current_elements_size: OperandSize::from_bytes(
+                        self.vmoffsets.size_of_vmtable_definition_current_elements(),
+                    ),
+                })
+            }
         }
     }
+
+    /// Get a [`TablePlan`] from a [`TableIndex`].
+    pub fn table_plan(&mut self, index: TableIndex) -> &TablePlan {
+        &self.translation.module.table_plans[index]
+    }
 }