[stack-switching] GC interaction

cranelift/codegen/src/isa/x64/inst/emit.rs

@@ -550,6 +550,8 @@ pub(crate) fn emit(
             in_payload0,
             out_payload0,
         } => {
+            let stack_map = state.take_stack_map();
+
             // Note that we do not emit anything for preserving and restoring
             // ordinary registers here: That's taken care of by regalloc for us,
             // since we marked this instruction as clobbering all registers.
@@ -643,6 +645,9 @@ pub(crate) fn emit(
             asm::inst::jmpq_m::new(tmp1.to_reg()).emit(sink, info, state);
 
             sink.bind_label(resume, state.ctrl_plane_mut());
+            if let Some(s) = stack_map {
+                sink.push_user_stack_map(state, sink.cur_offset(), s);
+            }
         }
 
         Inst::JmpKnown { dst } => uncond_jmp(sink, *dst),

cranelift/codegen/src/isa/x64/inst/mod.rs

@@ -1467,7 +1467,9 @@ impl MachInst for Inst {
 
     fn is_safepoint(&self) -> bool {
         match self {
-            Inst::CallKnown { .. } | Inst::CallUnknown { .. } => true,
+            Inst::CallKnown { .. } | Inst::CallUnknown { .. } | Inst::StackSwitchBasic { .. } => {
+                true
+            }
             _ => false,
         }
     }

cranelift/filetests/filetests/isa/x64/stack_switch_stack_map.clif

@@ -0,0 +1,14 @@
+test compile
+set opt_level=speed
+set stack_switch_model=basic
+target x86_64
+
+;; A stack switch is a GC safepoint. Its stack map is associated with the
+;; resume PC saved in the outgoing control context.
+function %switch_with_stack_map(i64, i64, i64) -> i64 {
+    ss0 = explicit_slot 4
+
+block0(v0: i64, v1: i64, v2: i64):
+    v3 = stack_switch v0, v1, v2, stack_map=[i32 @ ss0+0]
+    return v3
+}

crates/cranelift/src/func_environ.rs

@@ -46,6 +46,20 @@ use wasmtime_environ::{
 };
 use wasmtime_environ::{FUNCREF_INIT_BIT, FUNCREF_MASK};
 
+/// Function-local stack slots backing a continuation's
+/// `VMPayloads::values`.
+///
+/// The values and their GC-reference markers are logically one
+/// payload descriptor but use distinct stack slots so that Cranelift
+/// can assign them distinct alias regions. The marker slot is only
+/// allocated when this function contains a stack switching site whose
+/// payloads may contain GC references.
+#[derive(Clone, Copy)]
+pub(crate) struct VMPayloadStackSlots {
+    pub(crate) values: ir::StackSlot,
+    pub(crate) gc_ref_markers: Option<ir::StackSlot>,
+}
+
 #[derive(Copy, Clone, Debug)]
 pub(crate) enum Extension {
     Sign,
@@ -216,10 +230,9 @@ pub struct FuncEnvironment<'module_environment> {
     /// current stack's `handler_list` field.
     stack_switching_handler_list_buffer: Option<ir::StackSlot>,
 
-    /// Used by the stack switching feature. If set, we have a allocated a
-    /// slot on this function's stack to be used for the
-    /// current continuation's `values` field.
-    stack_switching_values_buffer: Option<ir::StackSlot>,
+    /// Used by the stack switching feature. If set, these are the stack slots
+    /// backing the current continuation's `values` field.
+    stack_switching_values_storage: Option<VMPayloadStackSlots>,
 
     /// The stack-slot used for exposing Wasm state via debug
     /// instrumentation, if any, and the builder containing its metadata.
@@ -300,7 +313,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             stack_limit_at_function_entry: None,
 
             stack_switching_handler_list_buffer: None,
-            stack_switching_values_buffer: None,
+            stack_switching_values_storage: None,
 
             state_slot: None,
             next_srcloc: ir::SourceLoc::default(),
@@ -5259,8 +5272,9 @@ impl FuncEnvironment<'_> {
         builder: &mut FunctionBuilder<'_>,
         contobj: ir::Value,
         args: &[ir::Value],
+        arg_types: &[WasmValType],
     ) -> ir::Value {
-        stack_switching::instructions::translate_cont_bind(self, builder, contobj, args)
+        stack_switching::instructions::translate_cont_bind(self, builder, contobj, args, arg_types)
     }
 
     pub fn translate_cont_new(
@@ -5340,13 +5354,15 @@ impl FuncEnvironment<'_> {
         builder: &mut FunctionBuilder<'_>,
         tag_index: u32,
         suspend_args: &[ir::Value],
-        tag_return_types: &[ir::Type],
+        suspend_arg_types: &[WasmValType],
+        tag_return_types: &[WasmValType],
     ) -> WasmResult<Vec<ir::Value>> {
         stack_switching::instructions::translate_suspend(
             self,
             builder,
             tag_index,
             suspend_args,
+            suspend_arg_types,
             tag_return_types,
         )
     }
@@ -5358,14 +5374,16 @@ impl FuncEnvironment<'_> {
         tag_index: u32,
         contobj: ir::Value,
         switch_args: &[ir::Value],
-        return_types: &[ir::Type],
+        switch_arg_types: &[WasmValType],
+        return_types: &[WasmValType],
     ) -> WasmResult<Vec<ir::Value>> {
         stack_switching::instructions::translate_switch(
             self,
             builder,
             tag_index,
             contobj,
             switch_args,
+            switch_arg_types,
             return_types,
         )
     }

crates/cranelift/src/func_environ/gc.rs

@@ -2,7 +2,7 @@
 
 use crate::TRAP_ARRAY_OUT_OF_BOUNDS;
 use crate::bounds_checks::BoundsCheck;
-use crate::func_environ::{CheckedEntity, Extension, FuncEnvironment};
+use crate::func_environ::{CheckedEntity, Extension, FuncEnvironment, stack_switching::fatpointer};
 use crate::translate::{
     Heap, HeapData, MemoryKind, StructFieldsVec, TargetEnvironment, VmctxLoadChain,
 };
@@ -377,10 +377,7 @@ pub fn read_field_at_addr(
                         .call(get_interned_func_ref, &[vmctx, func_ref_id, expected_ty]);
                     builder.func.dfg.first_result(call_inst)
                 }
-                WasmHeapTopType::Cont => {
-                    // TODO(#10248) GC integration for stack switching
-                    return stack_switching_unsupported();
-                }
+                WasmHeapTopType::Cont => read_cont_ref_at_addr(func_env, builder, addr, flags)?,
             },
         },
     };
@@ -431,6 +428,89 @@ pub fn intern_func_ref(
     Ok(builder.ins().ireduce(ir::types::I32, func_ref_id))
 }
 
+fn intern_cont_ref(
+    func_env: &mut FuncEnvironment<'_>,
+    builder: &mut FunctionBuilder<'_>,
+    ref_type: WasmRefType,
+    contobj: ir::Value,
+) -> ir::Value {
+    assert_eq!(ref_type.heap_type.top(), WasmHeapTopType::Cont);
+
+    if ref_type.heap_type == WasmHeapType::NoCont {
+        let zero = builder.ins().iconst(ir::types::I32, 0);
+
+        if !ref_type.nullable {
+            builder.ins().trapz(zero, TRAP_INTERNAL_ASSERT);
+        }
+
+        return zero;
+    }
+
+    let (revision, contref) = fatpointer::deconstruct(func_env, &mut builder.cursor(), contobj);
+    let vmctx = func_env.vmctx_val(&mut builder.cursor());
+    let intern = func_env
+        .builtin_functions
+        .intern_contref_for_gc_heap(builder.func);
+    let call = builder.ins().call(intern, &[vmctx, contref, revision]);
+    let id = builder.func.dfg.first_result(call);
+    builder.ins().ireduce(ir::types::I32, id)
+}
+
+fn read_cont_ref_at_addr(
+    func_env: &mut FuncEnvironment<'_>,
+    builder: &mut FunctionBuilder<'_>,
+    addr: ir::Value,
+    flags: ir::MemFlagsData,
+) -> WasmResult<ir::Value> {
+    let id = builder.ins().load(ir::types::I32, flags, addr, 0);
+
+    // We create a stack slot to hold the continuation values (16
+    // bytes), and pass the address of this slot as the out parameter
+    // to the builtin.
+    let pointer_type = func_env.pointer_type();
+    let pointer_bytes = pointer_type.bytes();
+    let slot = builder.create_sized_stack_slot(ir::StackSlotData::new(
+        ir::StackSlotKind::ExplicitSlot,
+        2 * pointer_bytes,
+        u8::try_from(pointer_bytes.trailing_zeros()).unwrap(),
+    ));
+    let out_result = builder.ins().stack_addr(pointer_type, slot, 0);
+
+    let vmctx = func_env.vmctx_val(&mut builder.cursor());
+    let get = func_env
+        .builtin_functions
+        .get_interned_contref(builder.func);
+    builder.ins().call(get, &[vmctx, id, out_result]);
+
+    let region = func_env.alias_regions.stack_slot_region(builder.func, slot);
+    let flags = ir::MemFlagsData::trusted().with_alias_region(Some(region));
+    let contref = builder.ins().load(pointer_type, flags, out_result, 0);
+    let revision = builder.ins().load(
+        pointer_type,
+        flags,
+        out_result,
+        i32::try_from(pointer_bytes).unwrap(),
+    );
+    Ok(fatpointer::construct(
+        func_env,
+        &mut builder.cursor(),
+        revision,
+        contref,
+    ))
+}
+
+fn write_cont_ref_at_addr(
+    func_env: &mut FuncEnvironment<'_>,
+    builder: &mut FunctionBuilder<'_>,
+    ref_type: WasmRefType,
+    flags: ir::MemFlagsData,
+    field_addr: ir::Value,
+    contobj: ir::Value,
+) {
+    let id = intern_cont_ref(func_env, builder, ref_type, contobj);
+    builder.ins().store(flags, id, field_addr, 0);
+}
+
 fn write_func_ref_at_addr(
     func_env: &mut FuncEnvironment<'_>,
     builder: &mut FunctionBuilder<'_>,
@@ -487,7 +567,9 @@ pub fn write_field_at_addr(
                     func_env, builder, r, field_addr, new_val, flags,
                 )?;
             }
-            WasmHeapTopType::Cont => return stack_switching_unsupported(),
+            WasmHeapTopType::Cont => {
+                write_cont_ref_at_addr(func_env, builder, r, flags, field_addr, new_val)
+            }
         },
         WasmStorageType::Val(_) => {
             assert_eq!(
@@ -527,6 +609,10 @@ fn default_value(
             }
             WasmValType::Ref(r) => {
                 assert!(r.nullable);
+                if r.heap_type.top() == WasmHeapTopType::Cont {
+                    let zero = cursor.ins().iconst(func_env.pointer_type(), 0);
+                    return fatpointer::construct(func_env, cursor, zero, zero);
+                }
                 let (ty, needs_stack_map) = func_env.reference_type(r.heap_type);
 
                 // NB: The collector doesn't need to know about null references.

crates/cranelift/src/func_environ/gc/copying.rs

@@ -440,7 +440,9 @@ impl GcCompiler for CopyingCompiler {
                     // store the reference directly.
                     unbarriered_store_gc_ref(builder, r.heap_type, field_addr, val, flags)?;
                 }
-                WasmHeapTopType::Cont => return super::stack_switching_unsupported(),
+                WasmHeapTopType::Cont => {
+                    write_field_at_addr(func_env, builder, ty, field_addr, val)?
+                }
             },
             WasmStorageType::I8 => {
                 assert_eq!(builder.func.dfg.value_type(val), ir::types::I32);

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

@@ -34,7 +34,7 @@ pub(crate) fn deconstruct<'a>(
 /// Constructs a continuation object from a given contref and revision pointer.
 /// The contref_addr may be 0, to indicate that we want to build a wasm null reference.
 pub(crate) fn construct<'a>(
-    env: &mut crate::func_environ::FuncEnvironment<'a>,
+    env: &crate::func_environ::FuncEnvironment<'a>,
     pos: &mut cranelift_codegen::cursor::FuncCursor,
     revision_counter: ir::Value,
     contref_addr: ir::Value,