stack-switching: define VMContObj as two words

This change migrates VMContObj and its usages in cranelift and runtime
to work with the VMContObj fat pointer as two words in order to better
target different architectures (still gated to x86_64 for now).

To support this, a size type was plumbed into the builtins function
signature types (as is done for component types) that maps to
usize.

Author: posborne

crates/cranelift/src/func_environ.rs

@@ -2397,7 +2397,7 @@ impl FuncEnvironment<'_> {
             WasmHeapTopType::Cont => {
                 let (elem_addr, flags) = table_data.prepare_table_addr(self, builder, index);
                 Ok(builder.ins().load(
-                    stack_switching::fatpointer::POINTER_TYPE,
+                    stack_switching::fatpointer::fatpointer_type(self),
                     flags,
                     elem_addr,
                     0,

crates/cranelift/src/func_environ/stack_switching/fatpointer.rs

@@ -1,15 +1,16 @@
 use cranelift_codegen::ir;
 use cranelift_codegen::ir::InstBuilder;
 
-/// The Cranelift type used to represent all of the following:
+/// Returns the Cranelift type used to represent all of the following:
 /// - wasm values of type `(ref null $ct)` and `(ref $ct)`
 /// - equivalently: runtime values of type `Option<VMContObj>` and `VMContObj`
-/// Note that a `VMContObj` is a fat pointer
-/// consisting of a pointer to `VMContRef` and a 64 bit sequence
-/// counter.
-/// We represent this here as a 128bit value, with the same representation as
-/// `core::mem::transmute::<i128, VMContObj>`.
-pub const POINTER_TYPE: ir::Type = ir::types::I128;
+/// Note that a `VMContObj` is a fat pointer consisting of a pointer to
+/// `VMContRef` and a pointer-sized revision counter. We represent this as 2 words
+/// (pointer and usize).
+pub fn fatpointer_type(env: &crate::func_environ::FuncEnvironment) -> ir::Type {
+    let ptr_bits = env.pointer_type().bits();
+    ir::Type::int((2 * ptr_bits).try_into().unwrap()).unwrap()
+}
 
 /// Turns a (possibly null) reference to a continuation object into a tuple
 /// (revision, contref_ptr). If `contobj` denotes a wasm null reference, the
@@ -19,13 +20,13 @@ pub(crate) fn deconstruct<'a>(
     pos: &mut cranelift_codegen::cursor::FuncCursor,
     contobj: ir::Value,
 ) -> (ir::Value, ir::Value) {
-    debug_assert_eq!(pos.func.dfg.value_type(contobj), POINTER_TYPE);
+    debug_assert_eq!(pos.func.dfg.value_type(contobj), fatpointer_type(env));
     let ptr_ty = env.pointer_type();
-    assert!(ptr_ty.bits() <= 64);
+    let ptr_bits = ptr_ty.bits();
 
     let contref = pos.ins().ireduce(ptr_ty, contobj);
-    let shifted = pos.ins().ushr_imm(contobj, 64);
-    let revision_counter = pos.ins().ireduce(ir::types::I64, shifted);
+    let shifted = pos.ins().ushr_imm(contobj, i64::from(ptr_bits));
+    let revision_counter = pos.ins().ireduce(ptr_ty, shifted);
 
     (revision_counter, contref)
 }
@@ -38,13 +39,16 @@ pub(crate) fn construct<'a>(
     revision_counter: ir::Value,
     contref_addr: ir::Value,
 ) -> ir::Value {
-    debug_assert_eq!(pos.func.dfg.value_type(contref_addr), env.pointer_type());
-    debug_assert_eq!(pos.func.dfg.value_type(revision_counter), ir::types::I64);
-    assert!(env.pointer_type().bits() <= 64);
+    let ptr_ty = env.pointer_type();
+    let ptr_bits = ptr_ty.bits();
+    let fat_ptr_ty = fatpointer_type(env);
+
+    debug_assert_eq!(pos.func.dfg.value_type(contref_addr), ptr_ty);
+    debug_assert_eq!(pos.func.dfg.value_type(revision_counter), ptr_ty);
 
-    let contref_addr = pos.ins().uextend(ir::types::I128, contref_addr);
-    let revision_counter = pos.ins().uextend(ir::types::I128, revision_counter);
-    let shifted_counter = pos.ins().ishl_imm(revision_counter, 64);
+    let contref_addr = pos.ins().uextend(fat_ptr_ty, contref_addr);
+    let revision_counter = pos.ins().uextend(fat_ptr_ty, revision_counter);
+    let shifted_counter = pos.ins().ishl_imm(revision_counter, i64::from(ptr_bits));
     let contobj = pos.ins().bor(shifted_counter, contref_addr);
 
     contobj

crates/cranelift/src/lib.rs

@@ -213,7 +213,10 @@ fn reference_type(wasm_ht: WasmHeapType, pointer_type: ir::Type) -> ir::Type {
     match wasm_ht.top() {
         WasmHeapTopType::Func => pointer_type,
         WasmHeapTopType::Any | WasmHeapTopType::Extern | WasmHeapTopType::Exn => ir::types::I32,
-        WasmHeapTopType::Cont => func_environ::stack_switching::fatpointer::POINTER_TYPE,
+        WasmHeapTopType::Cont => {
+            // VMContObj is 2 * pointer_size (pointer + usize revision)
+            ir::Type::int((2 * pointer_type.bits()).try_into().unwrap()).unwrap()
+        }
     }
 }
 
@@ -372,6 +375,10 @@ impl BuiltinFunctionSignatures {
         AbiParam::new(ir::types::I8)
     }
 
+    fn size(&self) -> AbiParam {
+        AbiParam::new(self.pointer_type)
+    }
+
     fn wasm_signature(&self, builtin: BuiltinFunctionIndex) -> Signature {
         let mut _cur = 0;
         macro_rules! iter {

crates/environ/src/builtin.rs

@@ -233,12 +233,12 @@ macro_rules! foreach_builtin_function {
             // denote the continuation being `None`, `init_contref`
             // may be 0.
             #[cfg(feature = "stack-switching")]
-            table_grow_cont_obj(vmctx: vmctx, table: u32, delta: u64, init_contref: pointer, init_revision: u64) -> pointer;
+            table_grow_cont_obj(vmctx: vmctx, table: u32, delta: u64, init_contref: pointer, init_revision: size) -> pointer;
 
             // `value_contref` and `value_revision` together encode
             // the Option<VMContObj>, as in previous libcall.
             #[cfg(feature = "stack-switching")]
-            table_fill_cont_obj(vmctx: vmctx, table: u32, dst: u64, value_contref: pointer, value_revision: u64, len: u64) -> bool;
+            table_fill_cont_obj(vmctx: vmctx, table: u32, dst: u64, value_contref: pointer, value_revision: size, len: u64) -> bool;
 
             // Return the instance ID for a given vmctx.
             #[cfg(feature = "gc")]

crates/environ/src/vmoffsets.rs

@@ -348,12 +348,17 @@ pub trait PtrSize {
 
     /// Return the offset of `VMContObj::revision`
     fn vmcontobj_revision(&self) -> u8 {
-        8
+        self.size()
     }
 
-    /// Return the size of `VMHostArray`.
+    /// Return the size of `VMContObj`.
     fn size_of_vmcontobj(&self) -> u8 {
-        16
+        u8::try_from(align(
+            u32::from(self.vmcontobj_revision())
+                + u32::try_from(core::mem::size_of::<usize>()).unwrap(),
+            u32::from(self.size()),
+        ))
+        .unwrap()
     }
 
     // Offsets within `VMContRef`

crates/wasmtime/src/runtime/vm/libcalls.rs

@@ -157,6 +157,7 @@ pub mod raw {
         (@ty f64x2) => (f64x2);
         (@ty bool) => (bool);
         (@ty pointer) => (*mut u8);
+        (@ty size) => (usize);
     }
 
     wasmtime_environ::foreach_builtin_function!(libcall);
@@ -339,7 +340,7 @@ unsafe fn table_grow_cont_obj(
     // The following two values together form the initial Option<VMContObj>.
     // A None value is indicated by the pointer being null.
     init_value_contref: *mut u8,
-    init_value_revision: u64,
+    init_value_revision: usize,
 ) -> Result<Option<AllocationSize>> {
     let defined_table_index = DefinedTableIndex::from_u32(defined_table_index);
     let element = unsafe { VMContObj::from_raw_parts(init_value_contref, init_value_revision) };
@@ -416,7 +417,7 @@ unsafe fn table_fill_cont_obj(
     table_index: u32,
     dst: u64,
     value_contref: *mut u8,
-    value_revision: u64,
+    value_revision: usize,
     len: u64,
 ) -> Result<()> {
     let instance = store.instance_mut(instance);

crates/wasmtime/src/runtime/vm/stack_switching.rs

@@ -39,15 +39,15 @@ pub use stack::*;
 /// For performance reasons, the VMContRef at the bottom of this chain
 /// (i.e., the one pointed to by the VMContObj) has a pointer to the
 /// other end of the chain (i.e., its last ancestor).
-#[repr(C, align(16))]
+#[repr(C)]
 #[derive(Debug, Clone, Copy)]
 pub struct VMContObj {
     pub contref: NonNull<VMContRef>,
-    pub revision: u64,
+    pub revision: usize,
 }
 
 impl VMContObj {
-    pub fn new(contref: NonNull<VMContRef>, revision: u64) -> Self {
+    pub fn new(contref: NonNull<VMContRef>, revision: usize) -> Self {
         Self { contref, revision }
     }
 
@@ -59,7 +59,7 @@ impl VMContObj {
     ///
     /// Behavior will be undefined if a pointer to data that is not a
     /// VMContRef is provided.
-    pub unsafe fn from_raw_parts(contref: *mut u8, revision: u64) -> Option<Self> {
+    pub unsafe fn from_raw_parts(contref: *mut u8, revision: usize) -> Option<Self> {
         NonNull::new(contref.cast::<VMContRef>()).map(|contref| Self::new(contref, revision))
     }
 }
@@ -204,7 +204,7 @@ pub struct VMContRef {
     pub last_ancestor: *mut VMContRef,
 
     /// Revision counter.
-    pub revision: u64,
+    pub revision: usize,
 
     /// The underlying stack.
     pub stack: VMContinuationStack,

crates/wasmtime/src/runtime/vm/vmcontext.rs

@@ -1056,6 +1056,7 @@ macro_rules! define_builtin_array {
     (@ty f64x2) => (f64x2);
     (@ty bool) => (bool);
     (@ty pointer) => (*mut u8);
+    (@ty size) => (usize);
     (@ty vmctx) => (NonNull<VMContext>);
 }