Add support for the stack switching proposal (#1802)

* Stack switching feature flag

* Stubs for stack switching operators.

* Type structure

* Stack switching instructions validation.

* Simplify validation API

* Wast support

* Skip AssertSuspension during roundtrip testing

* Fix bugs; particularly label resolution bug in resume tables. Add test file.

* Add snapshots

* Fix validation bugs

* Fix validation bugs. Test interaction between continuation types and gc types.

* Add more tests

* Add sub_type_at_id to EmptyResources. Formatting.

* Format

* Comma

* Fix validation bugs in cont.bind and switch. Extend validation testsuite.

* Address Alex's feedback

* Update signature of sub_type_at_id in EmptyResources

* Format

* [wast] Translate Cont and NoCont heap types.

* Address Alex's second round of feedback.

* Revert encoding of continuation types back to s33 to stay in sync with the reference interpreter.

Author: dhil

crates/wasm-encoder/src/core/code.rs

@@ -1203,6 +1203,26 @@ pub enum Instruction<'a> {
         array_type_index: u32,
     },
     RefI31Shared,
+    // Stack switching
+    ContNew(u32),
+    ContBind {
+        argument_index: u32,
+        result_index: u32,
+    },
+    Suspend(u32),
+    Resume {
+        cont_type_index: u32,
+        resume_table: Cow<'a, [Handle]>,
+    },
+    ResumeThrow {
+        cont_type_index: u32,
+        tag_index: u32,
+        resume_table: Cow<'a, [Handle]>,
+    },
+    Switch {
+        cont_type_index: u32,
+        tag_index: u32,
+    },
 }
 
 impl Encode for Instruction<'_> {
@@ -3695,6 +3715,48 @@ impl Encode for Instruction<'_> {
                 sink.push(0xFE);
                 sink.push(0x72);
             }
+            Instruction::ContNew(type_index) => {
+                sink.push(0xE0);
+                type_index.encode(sink);
+            }
+            Instruction::ContBind {
+                argument_index,
+                result_index,
+            } => {
+                sink.push(0xE1);
+                argument_index.encode(sink);
+                result_index.encode(sink);
+            }
+            Instruction::Suspend(tag_index) => {
+                sink.push(0xE2);
+                tag_index.encode(sink);
+            }
+            Instruction::Resume {
+                cont_type_index,
+                ref resume_table,
+            } => {
+                sink.push(0xE3);
+                cont_type_index.encode(sink);
+                resume_table.encode(sink);
+            }
+            Instruction::ResumeThrow {
+                cont_type_index,
+                tag_index,
+                ref resume_table,
+            } => {
+                sink.push(0xE4);
+                cont_type_index.encode(sink);
+                tag_index.encode(sink);
+                resume_table.encode(sink);
+            }
+            Instruction::Switch {
+                cont_type_index,
+                tag_index,
+            } => {
+                sink.push(0xE5);
+                cont_type_index.encode(sink);
+                tag_index.encode(sink);
+            }
         }
     }
 }
@@ -3733,6 +3795,29 @@ impl Encode for Catch {
     }
 }
 
+#[derive(Clone, Debug)]
+#[allow(missing_docs)]
+pub enum Handle {
+    OnLabel { tag: u32, label: u32 },
+    OnSwitch { tag: u32 },
+}
+
+impl Encode for Handle {
+    fn encode(&self, sink: &mut Vec<u8>) {
+        match self {
+            Handle::OnLabel { tag, label } => {
+                sink.push(0x00);
+                tag.encode(sink);
+                label.encode(sink);
+            }
+            Handle::OnSwitch { tag } => {
+                sink.push(0x01);
+                tag.encode(sink);
+            }
+        }
+    }
+}
+
 /// A constant expression.
 ///
 /// Usable in contexts such as offsets or initializers.

crates/wasm-encoder/src/core/types.rs

@@ -31,6 +31,8 @@ pub enum CompositeInnerType {
     Array(ArrayType),
     /// The type is for a struct.
     Struct(StructType),
+    /// The type is for a continuation.
+    Cont(ContType),
 }
 
 /// Represents a type of a function in a WebAssembly module.
@@ -88,6 +90,10 @@ impl StorageType {
     }
 }
 
+/// Represents a type of a continuation in a WebAssembly module.
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+pub struct ContType(pub u32);
+
 /// The type of a core WebAssembly value.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
 pub enum ValType {
@@ -443,6 +449,12 @@ pub enum AbstractHeapType {
 
     /// The abstract `noexn` heap type.
     NoExn,
+
+    /// The abstract `cont` heap type.
+    Cont,
+
+    /// The abstract `nocont` heap type.
+    NoCont,
 }
 
 impl Encode for AbstractHeapType {
@@ -461,6 +473,8 @@ impl Encode for AbstractHeapType {
             I31 => sink.push(0x6C),
             Exn => sink.push(0x69),
             NoExn => sink.push(0x74),
+            Cont => sink.push(0x68),
+            NoCont => sink.push(0x75),
         }
     }
 }
@@ -611,6 +625,11 @@ impl<'a> CoreTypeEncoder<'a> {
         }
     }
 
+    fn encode_cont(&mut self, ty: &ContType) {
+        self.bytes.push(0x5d);
+        i64::from(ty.0).encode(self.bytes);
+    }
+
     /// Define an explicit subtype in this type section.
     pub fn subtype(mut self, ty: &SubType) {
         self.encode_subtype(ty)
@@ -643,6 +662,7 @@ impl<'a> CoreTypeEncoder<'a> {
                 self.encode_array(&ty.element_type, ty.mutable)
             }
             CompositeInnerType::Struct(ty) => self.encode_struct(ty.fields.iter().cloned()),
+            CompositeInnerType::Cont(ty) => self.encode_cont(ty),
         }
     }
 

crates/wasm-encoder/src/reencode.rs

@@ -127,13 +127,24 @@ pub trait Reencode {
         utils::func_type(self, func_ty)
     }
 
+    fn cont_type(
+        &mut self,
+        cont_ty: wasmparser::ContType,
+    ) -> Result<crate::ContType, Error<Self::Error>> {
+        utils::cont_type(self, cont_ty)
+    }
+
     fn global_type(
         &mut self,
         global_ty: wasmparser::GlobalType,
     ) -> Result<crate::GlobalType, Error<Self::Error>> {
         utils::global_type(self, global_ty)
     }
 
+    fn handle(&mut self, on: wasmparser::Handle) -> crate::Handle {
+        utils::handle(self, on)
+    }
+
     fn heap_type(
         &mut self,
         heap_type: wasmparser::HeapType,
@@ -883,6 +894,21 @@ pub mod utils {
         }
     }
 
+    pub fn handle<T: ?Sized + Reencode>(
+        reencoder: &mut T,
+        arg: wasmparser::Handle,
+    ) -> crate::Handle {
+        match arg {
+            wasmparser::Handle::OnLabel { tag, label } => crate::Handle::OnLabel {
+                tag: reencoder.tag_index(tag),
+                label,
+            },
+            wasmparser::Handle::OnSwitch { tag } => crate::Handle::OnSwitch {
+                tag: reencoder.tag_index(tag),
+            },
+        }
+    }
+
     /// Parses the input `section` given from the `wasmparser` crate and
     /// adds the custom section to the `module`.
     pub fn parse_custom_section<T: ?Sized + Reencode>(
@@ -989,6 +1015,8 @@ pub mod utils {
             I31 => crate::AbstractHeapType::I31,
             Exn => crate::AbstractHeapType::Exn,
             NoExn => crate::AbstractHeapType::NoExn,
+            Cont => crate::AbstractHeapType::Cont,
+            NoCont => crate::AbstractHeapType::NoCont,
         }
     }
 
@@ -1052,6 +1080,9 @@ pub mod utils {
             wasmparser::CompositeInnerType::Struct(s) => {
                 crate::CompositeInnerType::Struct(reencoder.struct_type(s)?)
             }
+            wasmparser::CompositeInnerType::Cont(c) => {
+                crate::CompositeInnerType::Cont(reencoder.cont_type(c)?)
+            }
         };
         Ok(crate::CompositeType {
             inner,
@@ -1114,6 +1145,15 @@ pub mod utils {
         })
     }
 
+    pub fn cont_type<T: ?Sized + Reencode>(
+        reencoder: &mut T,
+        cont_ty: wasmparser::ContType,
+    ) -> Result<crate::ContType, Error<T::Error>> {
+        Ok(crate::ContType(
+            reencoder.type_index_unpacked(cont_ty.0.unpack())?,
+        ))
+    }
+
     pub fn val_type<T: ?Sized + Reencode>(
         reencoder: &mut T,
         val_ty: wasmparser::ValType,
@@ -1555,6 +1595,12 @@ pub mod utils {
             (map $arg:ident array_size) => ($arg);
             (map $arg:ident field_index) => ($arg);
             (map $arg:ident try_table) => ($arg);
+            (map $arg:ident argument_index) => (reencoder.type_index($arg));
+            (map $arg:ident result_index) => (reencoder.type_index($arg));
+            (map $arg:ident cont_type_index) => (reencoder.type_index($arg));
+            (map $arg:ident resume_table) => ((
+                $arg.handlers.into_iter().map(|h| reencoder.handle(h)).collect::<Vec<_>>().into()
+            ));
 
             // This case takes the arguments of a wasmparser instruction and creates
             // a wasm-encoder instruction. There are a few special cases for where
@@ -1785,6 +1831,12 @@ impl TryFrom<wasmparser::GlobalType> for crate::GlobalType {
     }
 }
 
+impl From<wasmparser::Handle> for crate::Handle {
+    fn from(arg: wasmparser::Handle) -> Self {
+        RoundtripReencoder.handle(arg)
+    }
+}
+
 impl TryFrom<wasmparser::TypeRef> for crate::EntityType {
     type Error = Error;
 

crates/wasm-smith/src/core.rs

@@ -797,7 +797,7 @@ impl Module {
                     Extern => {
                         choices.push(ht(NoExtern));
                     }
-                    Exn | NoExn | None | NoExtern | NoFunc => {}
+                    Exn | NoExn | None | NoExtern | NoFunc | Cont | NoCont => {}
                 }
             }
             HT::Concrete(idx) => {
@@ -918,7 +918,10 @@ impl Module {
                     Eq => {
                         choices.push(ht(Any));
                     }
-                    Exn | Any | Func | Extern => {}
+                    NoCont => {
+                        choices.push(ht(Cont));
+                    }
+                    Exn | Any | Func | Extern | Cont => {}
                 }
             }
             HT::Concrete(mut idx) => {

crates/wasmparser/src/binary_reader.rs

@@ -1089,6 +1089,15 @@ impl<'a> BinaryReader<'a> {
             0xd5 => visitor.visit_br_on_null(self.read_var_u32()?),
             0xd6 => visitor.visit_br_on_non_null(self.read_var_u32()?),
 
+            0xe0 => visitor.visit_cont_new(self.read_var_u32()?),
+            0xe1 => visitor.visit_cont_bind(self.read_var_u32()?, self.read_var_u32()?),
+            0xe2 => visitor.visit_suspend(self.read_var_u32()?),
+            0xe3 => visitor.visit_resume(self.read_var_u32()?, self.read()?),
+            0xe4 => {
+                visitor.visit_resume_throw(self.read_var_u32()?, self.read_var_u32()?, self.read()?)
+            }
+            0xe5 => visitor.visit_switch(self.read_var_u32()?, self.read_var_u32()?),
+
             0xfb => self.visit_0xfb_operator(pos, visitor)?,
             0xfc => self.visit_0xfc_operator(pos, visitor)?,
             0xfd => self.visit_0xfd_operator(pos, visitor)?,

crates/wasmparser/src/features.rs

@@ -225,6 +225,8 @@ define_wasm_features! {
         /// feature. Those are expected to not require a full garbage collector
         /// so are not gated by this.
         pub gc_types: GC_TYPES(1 << 26) = true;
+        /// The WebAssembly [stack-switching proposal](https://github.com/WebAssembly/stack-switching).
+        pub stack_switching: STACK_SWITCHING(1 << 27) = false;
     }
 }
 

crates/wasmparser/src/lib.rs

@@ -803,6 +803,14 @@ macro_rules! for_each_operator {
             @function_references RefAsNonNull => visit_ref_as_non_null
             @function_references BrOnNull { relative_depth: u32 } => visit_br_on_null
             @function_references BrOnNonNull { relative_depth: u32 } => visit_br_on_non_null
+
+            // Stack switching
+            @stack_switching ContNew { cont_type_index: u32 } => visit_cont_new
+            @stack_switching ContBind { argument_index: u32, result_index: u32 } => visit_cont_bind
+            @stack_switching Suspend { tag_index: u32 } => visit_suspend
+            @stack_switching Resume { cont_type_index: u32, resume_table: $crate::ResumeTable } => visit_resume
+            @stack_switching ResumeThrow { cont_type_index: u32, tag_index: u32, resume_table: $crate::ResumeTable } => visit_resume_throw
+            @stack_switching Switch { cont_type_index: u32, tag_index: u32 } => visit_switch
         }
     };
 }

crates/wasmparser/src/limits.rs

@@ -39,6 +39,7 @@ pub const MAX_WASM_BR_TABLE_SIZE: usize = MAX_WASM_FUNCTION_SIZE;
 pub const MAX_WASM_STRUCT_FIELDS: usize = 10_000;
 pub const MAX_WASM_CATCHES: usize = 10_000;
 pub const MAX_WASM_SUBTYPING_DEPTH: usize = 63;
+pub const MAX_WASM_HANDLERS: usize = 10_000;
 
 pub const DEFAULT_WASM_PAGE_SIZE: u64 = 1 << 16;