Codegen for switch instructions (#264)

This PR provides the missing piece to support `switch` instructions, by
adding the necessary codegen (but only for the optimized implementation,
support in the baseline implementation is not included).

Thus, the main change is the addition of
`wasmtime_cranelift::wasmfx::optimized::translate_switch`.

In general, the implementation piggybacks on #256: On `resume`, we now
fill the `HandlerList` with _m_ entries for all tags with suspend
handlers, followed by _n_ entries for all tags with switch handlers. The
`search_handler` code is changed so that on `suspend` and `switch`, we
look in the correct part of the `HandlerList`. We could use two separate
`HandlerLists` instead, but then we would have yet another allocation to
manage, so putting both kinds of tags into the same list, and then only
searching part of it seems preferable.


A few more notes:
- `translate_resume` and `translate_switch` now return a `WasmResult`.
This is used so that the baseline implementation can bail out without
needing to `panic`.
- The test runner in `stack_switchting.rs` now takes an extra parameter
that allows us to enable the gc proposal, which is required for the
tests using recursive types.

---------

Co-authored-by: Daniel Hillerström <[email protected]>

crates/continuations/src/lib.rs

@@ -77,13 +77,30 @@ pub struct CommonStackInformation {
     /// Otherwise, the list may remain allocated, but its `length` must be 0.
     ///
     /// Represents the handlers that this stack installed when resume-ing a
-    /// continuation with
-    /// (resume $ct handler_clause_1 ... handler_clause_n)
-    /// The actual entries are tag identifiers (i.e., *mut VMTagDefinition). The
-    /// order of entries is relevant: The tag in the i-th entry in the list is
-    /// handled by the block mentioned in the i-th handler clause in the resume
-    /// instruction above.
+    /// continuation.
+    ///
+    /// Note that for any resume instruction, we can re-order the handler
+    /// clauses without changing behavior such that all the suspend handlers
+    /// come first, followed by all the switch handler (while maintaining the
+    /// original ordering within the two groups).
+    /// Thus, we assume that the given resume instruction has the following
+    /// shape:
+    ///
+    /// (resume $ct
+    ///   (on $tag_0 $block_0) ... (on $tag_{n-1} $block_{n-1})
+    ///   (on $tag_n switch) ... (on $tag_m switch)
+    /// )
+    ///
+    /// On resume, the handler list is then filled with m + 1 (i.e., one per
+    /// handler clause) entries such that the i-th entry, using 0-based
+    /// indexing, is the identifier of $tag_i (represented as *mut
+    /// VMTagDefinition).
+    /// Further, `first_switch_handler_index` (see below) is set to n (i.e., the
+    /// 0-based index of the first switch handler).
     pub handlers: HandlerList,
+
+    /// Only used when state is `Parent`. See documentation of `handlers` above.
+    pub first_switch_handler_index: u32,
 }
 
 impl CommonStackInformation {
@@ -92,6 +109,7 @@ impl CommonStackInformation {
             limits: StackLimits::default(),
             state: State::Running,
             handlers: HandlerList::new(INITIAL_HANDLER_LIST_CAPACITY as u32),
+            first_switch_handler_index: 0,
         }
     }
 }
@@ -208,7 +226,7 @@ pub const STACK_CHAIN_CONTINUATION_DISCRIMINANT: usize = 2;
 #[derive(Debug, Clone, Copy, PartialEq)]
 #[repr(i32)]
 pub enum State {
-    /// The `VMContRef` has been created, but `resume` has never been
+    /// The `VMContRef` has been created, but neither `resume` or `switch` has ever been
     /// called on it. During this stage, we may add arguments using `cont.bind`.
     Fresh,
     /// The continuation is running, meaning that it is the one currently
@@ -219,7 +237,7 @@ pub enum State {
     /// another continuation (which may itself be `Running`, a `Parent`, or
     /// `Suspended`).
     Parent,
-    /// The continuation was suspended by a `suspend` instruction.
+    /// The continuation was suspended by a `suspend` or `switch` instruction.
     Suspended,
     /// The function originally passed to `cont.new` has returned normally.
     /// Note that there is no guarantee that a VMContRef will ever
@@ -295,6 +313,8 @@ pub mod offsets {
         pub const LIMITS: usize = offset_of!(CommonStackInformation, limits);
         pub const STATE: usize = offset_of!(CommonStackInformation, state);
         pub const HANDLERS: usize = offset_of!(CommonStackInformation, handlers);
+        pub const FIRST_SWITCH_HANDLER_INDEX: usize =
+            offset_of!(CommonStackInformation, first_switch_handler_index);
     }
 
     /// Size of wasmtime_runtime::continuation::FiberStack.
@@ -354,6 +374,9 @@ pub const CONTROL_EFFECT_RESUME_DISCRIMINANT: u32 = 1;
 /// Discriminant of variant `Suspend` in
 /// `ControlEffect`.
 pub const CONTROL_EFFECT_SUSPEND_DISCRIMINANT: u32 = 2;
+/// Discriminant of variant `Switch` in
+/// `ControlEffect`.
+pub const CONTROL_EFFECT_SWITCH_DISCRIMINANT: u32 = 3;
 
 /// Universal control effect. This structure encodes return signal,
 /// resume signal, suspension signal, and the handler to suspend to in a single variant type.
@@ -365,6 +388,7 @@ pub enum ControlEffect {
     Return = CONTROL_EFFECT_RETURN_DISCRIMINANT,
     Resume = CONTROL_EFFECT_RESUME_DISCRIMINANT,
     Suspend { handler_index: u32 } = CONTROL_EFFECT_SUSPEND_DISCRIMINANT,
+    Switch = CONTROL_EFFECT_SWITCH_DISCRIMINANT,
 }
 
 // TODO(frank-emrich) This conversion assumes little-endian data layout.

crates/cranelift/src/func_environ.rs

@@ -3261,8 +3261,8 @@ impl<'module_environment> crate::translate::FuncEnvironment
         type_index: u32,
         contobj: ir::Value,
         resume_args: &[ir::Value],
-        resumetable: &[(u32, ir::Block)],
-    ) -> Vec<ir::Value> {
+        resumetable: &[(u32, Option<ir::Block>)],
+    ) -> WasmResult<Vec<ir::Value>> {
         wasmfx_impl::translate_resume(self, builder, type_index, contobj, resume_args, resumetable)
     }
 
@@ -3286,6 +3286,18 @@ impl<'module_environment> crate::translate::FuncEnvironment
         wasmfx_impl::translate_suspend(self, builder, tag_index, suspend_args, tag_return_types)
     }
 
+    /// Translates switch instructions.
+    fn translate_switch(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        tag_index: u32,
+        contobj: ir::Value,
+        switch_args: &[ir::Value],
+        return_types: &[WasmValType],
+    ) -> WasmResult<Vec<ir::Value>> {
+        wasmfx_impl::translate_switch(self, builder, tag_index, contobj, switch_args, return_types)
+    }
+
     fn continuation_arguments(&self, index: u32) -> &[WasmValType] {
         let idx = self.module.types[TypeIndex::from_u32(index)];
         self.types[self.types[idx].unwrap_cont().clone().interned_type_index()]

crates/cranelift/src/translate/code_translator.rs

@@ -93,7 +93,7 @@ use std::vec::Vec;
 use wasmparser::{FuncValidator, MemArg, Operator, WasmModuleResources};
 use wasmtime_environ::{
     wasm_unsupported, DataIndex, ElemIndex, FuncIndex, GlobalIndex, MemoryIndex, Signed,
-    TableIndex, TypeIndex, Unsigned, WasmRefType, WasmResult,
+    TableIndex, TypeIndex, Unsigned, WasmHeapType, WasmRefType, WasmResult,
 };
 
 /// Given a `Reachability<T>`, unwrap the inner `T` or, when unreachable, set
@@ -2905,9 +2905,11 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                         let frame = &mut state.control_stack[i];
                         // This is side-effecting!
                         frame.set_branched_to_exit();
-                        resumetable.push((*tag, frame.br_destination()));
+                        resumetable.push((*tag, Some(frame.br_destination())));
+                    }
+                    wasmparser::Handle::OnSwitch { tag } => {
+                        resumetable.push((*tag, None));
                     }
-                    wasmparser::Handle::OnSwitch { tag: _ } => unimplemented!(),
                 }
             }
 
@@ -2920,7 +2922,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                 *contobj,
                 call_args,
                 resumetable.as_slice(),
-            );
+            )?;
 
             state.popn(arity + 1); // arguments + continuation
             state.pushn(&cont_return_vals);
@@ -2929,11 +2931,48 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             cont_type_index: _,
             tag_index: _,
             resume_table: _,
+        } => todo!("unimplemented stack switching instruction"),
+        Operator::Switch {
+            cont_type_index,
+            tag_index,
+        } => {
+            // Arguments of the continuation we are going to switch to
+            let continuation_argument_types =
+                environ.continuation_arguments(*cont_type_index).to_vec();
+            // Arity includes the continuation argument
+            let arity = continuation_argument_types.len();
+            let (contobj, switch_args) = state.peekn(arity).split_last().unwrap();
+
+            // Type of the continuation we are going to create by suspending the
+            // currently running stack
+            let current_continuation_type = continuation_argument_types.last().unwrap();
+            let current_continuation_type = current_continuation_type.unwrap_ref_type();
+
+            // Argument types of current_continuation_type. These will in turn
+            // be the types of the arguments we receive when someone switches
+            // back to this switch instruction
+            let current_continuation_arg_types = match current_continuation_type.heap_type {
+                WasmHeapType::ConcreteCont(index) => {
+                    let mti = index
+                        .as_module_type_index()
+                        .expect("Only supporting module type indices on switch for now");
+
+                    environ.continuation_arguments(mti.as_u32()).to_vec()
+                }
+                _ => panic!("Invalid type on switch"),
+            };
+
+            let switch_return_values = environ.translate_switch(
+                builder,
+                *tag_index,
+                *contobj,
+                switch_args,
+                &current_continuation_arg_types,
+            )?;
+
+            state.popn(arity);
+            state.pushn(&switch_return_values)
         }
-        | Operator::Switch {
-            cont_type_index: _,
-            tag_index: _,
-        } => todo!("unimplemented stack switching instructions"),
 
         Operator::GlobalAtomicGet { .. }
         | Operator::GlobalAtomicSet { .. }

crates/cranelift/src/translate/environ/spec.rs

@@ -845,8 +845,18 @@ pub trait FuncEnvironment: TargetEnvironment {
         type_index: u32,
         contobj: ir::Value,
         resume_args: &[ir::Value],
-        resumetable: &[(u32, ir::Block)],
-    ) -> Vec<ir::Value>;
+        resumetable: &[(u32, Option<ir::Block>)],
+    ) -> WasmResult<Vec<ir::Value>>;
+
+    /// Translates switch instructions.
+    fn translate_switch(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        tag_index: u32,
+        contobj: ir::Value,
+        switch_args: &[ir::Value],
+        return_types: &[WasmValType],
+    ) -> WasmResult<Vec<ir::Value>>;
 
     /// TODO(dhil): write documentation.
     #[allow(unused, reason = "TODO")]

crates/cranelift/src/wasmfx/baseline.rs

@@ -8,7 +8,7 @@ use cranelift_codegen::ir::types::*;
 use cranelift_codegen::ir::InstBuilder;
 use cranelift_frontend::{FunctionBuilder, Switch};
 use wasmtime_environ::PtrSize;
-use wasmtime_environ::{WasmResult, WasmValType};
+use wasmtime_environ::{WasmError, WasmResult, WasmValType};
 
 #[cfg_attr(
     not(feature = "wasmfx_baseline"),
@@ -182,8 +182,8 @@ pub(crate) fn translate_resume<'a>(
     type_index: u32,
     resumee_obj: ir::Value,
     resume_args: &[ir::Value],
-    resumetable: &[(u32, ir::Block)],
-) -> Vec<ir::Value> {
+    resumetable: &[(u32, Option<ir::Block>)],
+) -> WasmResult<Vec<ir::Value>> {
     // The resume instruction is by far the most involved
     // instruction to compile as it is responsible for both
     // continuation application and effect dispatch.
@@ -311,11 +311,16 @@ pub(crate) fn translate_resume<'a>(
     // Second, we consume the resume table entry-wise.
     let mut case_blocks = vec![];
     let mut tag_seen = std::collections::HashSet::new(); // Used to keep track of tags
-    for &(tag, label) in resumetable {
+    for &(tag, label_opt) in resumetable {
         // Skip if this `tag` has been seen previously.
         if !tag_seen.insert(tag) {
             continue;
         }
+        let label = label_opt.ok_or_else(|| {
+            WasmError::Unsupported(String::from(
+                "switch handlers not supported in the baseline implementation",
+            ))
+        })?;
         let case = builder.create_block();
         switch.set_entry(tag as u128, case);
         builder.switch_to_block(case);
@@ -401,7 +406,7 @@ pub(crate) fn translate_resume<'a>(
             tc_baseline_drop_continuation_reference(resumee_fiber)
         );
 
-        return values;
+        Ok(values)
     }
 }
 
@@ -453,3 +458,16 @@ pub(crate) fn translate_suspend<'a>(
 
     return_values
 }
+
+pub(crate) fn translate_switch<'a>(
+    _env: &mut crate::func_environ::FuncEnvironment<'a>,
+    _builder: &mut FunctionBuilder,
+    _tag_index: u32,
+    _switchee_contobj: ir::Value,
+    _switch_args: &[ir::Value],
+    _return_types: &[WasmValType],
+) -> WasmResult<Vec<ir::Value>> {
+    WasmResult::Err(WasmError::Unsupported(String::from(
+        "switch instructions are unsuported",
+    )))
+}