Linking tags

cranelift/wasm/src/code_translator.rs

@@ -2579,9 +2579,8 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let params = state.peekn(param_types.len());
             let param_count = params.len();
 
-            let tag_index_val = builder.ins().iconst(I32, *tag_index as i64);
             let return_values =
-                environ.translate_suspend(builder, tag_index_val, params, &return_types);
+                environ.translate_suspend(builder, *tag_index, params, &return_types);
 
             state.popn(param_count);
             state.pushn(&return_values);

cranelift/wasm/src/environ/spec.rs

@@ -647,7 +647,7 @@ pub trait FuncEnvironment: TargetEnvironment {
     fn translate_suspend(
         &mut self,
         builder: &mut FunctionBuilder,
-        tag_index: ir::Value,
+        tag_index: u32,
         suspend_args: &[ir::Value],
         tag_return_types: &[wasmtime_types::WasmValType],
     ) -> Vec<ir::Value>;

cranelift/wasm/src/sections_translator.rs

@@ -27,9 +27,7 @@ use wasmtime_types::ConstExpr;
 
 fn tag(e: TagType) -> Tag {
     match e.kind {
-        wasmparser::TagKind::Exception => Tag {
-            ty: TypeIndex::from_u32(e.func_type_idx),
-        },
+        wasmparser::TagKind::Exception => Tag::partial(e.func_type_idx),
     }
 }
 

crates/c-api/src/extern.rs

@@ -21,6 +21,7 @@ pub extern "C" fn wasm_extern_kind(e: &wasm_extern_t) -> wasm_externkind_t {
         Extern::Table(_) => crate::WASM_EXTERN_TABLE,
         Extern::Memory(_) => crate::WASM_EXTERN_MEMORY,
         Extern::SharedMemory(_) => todo!(),
+        Extern::Tag(_) => todo!(),
     }
 }
 
@@ -141,6 +142,7 @@ impl From<Extern> for wasmtime_extern_t {
                     sharedmemory: ManuallyDrop::new(Box::new(sharedmemory)),
                 },
             },
+            Extern::Tag(_) => todo!(),
         }
     }
 }

crates/c-api/src/types/extern.rs

@@ -25,6 +25,7 @@ impl CExternType {
             ExternType::Global(f) => CExternType::Global(CGlobalType::new(f)),
             ExternType::Memory(f) => CExternType::Memory(CMemoryType::new(f)),
             ExternType::Table(f) => CExternType::Table(CTableType::new(f)),
+            ExternType::Tag(_) => todo!(),
         }
     }
 }

crates/continuations/src/lib.rs

@@ -41,15 +41,6 @@ pub mod types {
         /// Type of the entries in the actual buffer
         pub type DataEntries = u128;
     }
-
-    /// Types used by `wasmtime_fibre::SwitchDirection` struct
-    pub mod switch_reason {
-        /// Type of `discriminant` field
-        pub type Discriminant = u32;
-
-        /// Type of `data` field
-        pub type Data = u32;
-    }
 }
 
 /// Runtime configuration options for WasmFX that can be set via the command
@@ -165,114 +156,6 @@ impl From<State> for i32 {
     }
 }
 
-// Runtime representation of tags
-pub type TagId = u32;
-
-/// See SwitchDirection below for overall use of this type.
-#[repr(u32)]
-#[derive(Debug, Clone)]
-pub enum SwitchDirectionEnum {
-    // Used to indicate that the contination has returned normally.
-    Return = 0,
-
-    // Indicates that we are suspendinga continuation due to invoking suspend.
-    // The payload is the tag to suspend with
-    Suspend = 1,
-
-    // Indicates that we are resuming a continuation via resume.
-    Resume = 2,
-}
-
-impl SwitchDirectionEnum {
-    pub fn discriminant_val(&self) -> u32 {
-        // This is well-defined for an enum with repr(u32).
-        unsafe { *(self as *const SwitchDirectionEnum as *const u32) }
-    }
-}
-
-/// Values of this type are passed to `wasmtime_fibre_switch` to indicate why we
-/// are switching. A nicer way of representing this type would be the following
-/// enum:
-///
-///```
-///  #[repr(C, u32)]
-///  pub enum SwitchDirection {
-///      // Used to indicate that the contination has returned normally.
-///      Return = 0,
-///
-///      // Indicates that we are suspendinga continuation due to invoking suspend.
-///      // The payload is the tag to suspend with
-///      Suspend(u32) = 1,
-///
-///      // Indicates that we are resuming a continuation via resume.
-///      Resume = 2,
-///  }
-///```
-///
-/// However, we want to convert values of type `SwitchDirection` to and from u64
-/// easily, which is why we need to ensure that it contains no uninitialised
-/// memory, to avoid undefined behavior.
-///
-/// We allow converting values of this type to and from u64.
-/// In that representation, bits 0 to 31 (where 0 is the LSB) contain the
-/// discriminant (as u32), while bits 32 to 63 contain the `data`.
-#[repr(C)]
-#[derive(Debug, Clone)]
-pub struct SwitchDirection {
-    pub discriminant: SwitchDirectionEnum,
-
-    // Stores tag value if `discriminant` is `suspend`, 0 otherwise.
-    pub data: u32,
-}
-
-impl SwitchDirection {
-    pub fn return_() -> SwitchDirection {
-        SwitchDirection {
-            discriminant: SwitchDirectionEnum::Return,
-            data: 0,
-        }
-    }
-
-    pub fn resume() -> SwitchDirection {
-        SwitchDirection {
-            discriminant: SwitchDirectionEnum::Resume,
-            data: 0,
-        }
-    }
-
-    pub fn suspend(tag: u32) -> SwitchDirection {
-        SwitchDirection {
-            discriminant: SwitchDirectionEnum::Suspend,
-            data: tag,
-        }
-    }
-}
-
-impl From<SwitchDirection> for u64 {
-    fn from(val: SwitchDirection) -> u64 {
-        // TODO(frank-emrich) This assumes little endian data layout. Should
-        // make this more explicit.
-        unsafe { core::mem::transmute::<SwitchDirection, u64>(val) }
-    }
-}
-
-impl From<u64> for SwitchDirection {
-    fn from(val: u64) -> SwitchDirection {
-        #[cfg(debug_assertions)]
-        {
-            let discriminant = val as u32;
-            debug_assert!(discriminant <= 2);
-            if discriminant != SwitchDirectionEnum::Suspend.discriminant_val() {
-                let data = val >> 32;
-                debug_assert_eq!(data, 0);
-            }
-        }
-        // TODO(frank-emrich) This assumes little endian data layout. Should
-        // make this more explicit.
-        unsafe { core::mem::transmute::<u64, SwitchDirection>(val) }
-    }
-}
-
 /// Defines offsets of the fields in the continuation-related types
 pub mod offsets {
     /// Offsets of fields in `Payloads`
@@ -327,3 +210,92 @@ pub mod offsets {
     /// We test there that this value is correct.
     pub const STACK_CHAIN_SIZE: usize = 2 * core::mem::size_of::<usize>();
 }
+
+#[repr(transparent)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct TaggedPointer(usize);
+
+impl TaggedPointer {
+    const LOW_TAG_BITS: usize = 2;
+    const LOW_TAG_MASK: usize = (1 << Self::LOW_TAG_BITS) - 1;
+
+    pub fn untagged(val: usize) -> Self {
+        Self(val)
+    }
+
+    pub fn low_tag(self, tag: usize) -> Self {
+        assert!(tag <= Self::LOW_TAG_MASK);
+        Self(self.0 | tag)
+    }
+
+    pub fn get_low_tag(self) -> usize {
+        self.0 & Self::LOW_TAG_MASK
+    }
+
+    pub fn low_untag(self) -> Self {
+        Self(self.0 & !Self::LOW_TAG_MASK)
+    }
+}
+
+impl From<TaggedPointer> for usize {
+    fn from(val: TaggedPointer) -> usize {
+        val.0
+    }
+}
+
+impl From<usize> for TaggedPointer {
+    fn from(val: usize) -> TaggedPointer {
+        TaggedPointer::untagged(val)
+    }
+}
+
+/// Universal control effect. This structure encodes return signal,
+/// resume signal, suspension signal, and suspension tags into a
+/// pointer. This instance is used at runtime. There is a codegen
+/// counterpart in `cranelift/src/wasmfx/shared.rs`.
+#[repr(transparent)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct ControlEffect(TaggedPointer);
+
+impl ControlEffect {
+    pub fn suspend(ptr: *const u8) -> Self {
+        let tptr = TaggedPointer::untagged(ptr as usize);
+        Self(TaggedPointer::low_tag(tptr, 0b01_usize))
+    }
+
+    pub fn return_() -> Self {
+        Self((0b00_usize).into())
+    }
+
+    pub fn resume() -> Self {
+        Self((0b11_usize).into())
+    }
+
+    fn new(raw: usize) -> Self {
+        Self(TaggedPointer::untagged(raw))
+    }
+
+    pub fn is_suspend(self) -> bool {
+        TaggedPointer::get_low_tag(self.0) == 0b01
+    }
+}
+
+impl From<u64> for ControlEffect {
+    fn from(val: u64) -> ControlEffect {
+        ControlEffect::new(val as usize)
+    }
+}
+
+impl From<ControlEffect> for u64 {
+    fn from(val: ControlEffect) -> u64 {
+        let raw: usize = val.0.into();
+        raw as u64
+    }
+}
+
+impl From<ControlEffect> for *mut u8 {
+    fn from(val: ControlEffect) -> *mut u8 {
+        let raw: usize = val.0.into();
+        raw as *mut u8
+    }
+}

crates/cranelift/src/func_environ.rs

@@ -93,7 +93,7 @@ wasmtime_environ::foreach_builtin_function!(declare_function_signatures);
 pub struct FuncEnvironment<'module_environment> {
     /// NOTE(frank-emrich) pub for use in crate::wasmfx::* modules
     pub(crate) isa: &'module_environment (dyn TargetIsa + 'module_environment),
-    module: &'module_environment Module,
+    pub(crate) module: &'module_environment Module,
     types: &'module_environment ModuleTypesBuilder,
     wasm_func_ty: &'module_environment WasmFuncType,
     sig_ref_to_ty: SecondaryMap<ir::SigRef, Option<&'module_environment WasmFuncType>>,
@@ -2810,7 +2810,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
     fn translate_suspend(
         &mut self,
         builder: &mut FunctionBuilder,
-        tag_index: ir::Value,
+        tag_index: u32,
         suspend_args: &[ir::Value],
         tag_return_types: &[WasmValType],
     ) -> Vec<ir::Value> {
@@ -2833,12 +2833,16 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
 
     fn tag_params(&self, tag_index: u32) -> &[WasmValType] {
         let idx = self.module.tags[TagIndex::from_u32(tag_index)].signature;
-        self.types[idx].unwrap_func().params()
+        self.types[idx.unwrap_module_type_index()]
+            .unwrap_func()
+            .params()
     }
 
     fn tag_returns(&self, tag_index: u32) -> &[WasmValType] {
         let idx = self.module.tags[TagIndex::from_u32(tag_index)].signature;
-        self.types[idx].unwrap_func().returns()
+        self.types[idx.unwrap_module_type_index()]
+            .unwrap_func()
+            .returns()
     }
 
     fn use_x86_blendv_for_relaxed_laneselect(&self, ty: Type) -> bool {

crates/cranelift/src/wasmfx/baseline.rs

@@ -443,12 +443,13 @@ pub(crate) fn translate_cont_new<'a>(
 pub(crate) fn translate_suspend<'a>(
     env: &mut crate::func_environ::FuncEnvironment<'a>,
     builder: &mut FunctionBuilder,
-    tag_index: ir::Value,
+    tag_index: u32,
     suspend_args: &[ir::Value],
     tag_return_types: &[WasmValType],
 ) -> Vec<ir::Value> {
+    let tag_index_val = builder.ins().iconst(I32, tag_index as i64);
     typed_continuations_store_payloads(env, builder, suspend_args);
-    call_builtin!(builder, env, tc_baseline_suspend(tag_index));
+    call_builtin!(builder, env, tc_baseline_suspend(tag_index_val));
     let contref = typed_continuations_load_continuation_reference(env, builder);
 
     let return_values =