[suspend register] merged

crates/continuations/src/lib.rs

@@ -2,6 +2,8 @@ use num_enum::{IntoPrimitive, TryFromPrimitive};
 use std::ptr;
 use wasmtime_fibre::Fiber;
 
+pub use wasmtime_fibre::{SwitchReason, SwitchReasonEnum, TagId};
+
 /// TODO
 #[allow(dead_code)]
 pub const ENABLE_DEBUG_PRINTING: bool = false;
@@ -27,9 +29,18 @@ pub mod types {
         /// Type of the entries in the actual buffer
         pub type DataEntries = u128;
     }
+
+    /// Types used by `wasmtime_fibre::SwitchReason` struct
+    pub mod switch_reason {
+        /// Type of `discriminant` field
+        pub type Discriminant = u32;
+
+        /// Type of `data` field
+        pub type Data = u32;
+    }
 }
 
-pub type ContinuationFiber = Fiber<'static, (), u32, ()>;
+pub type ContinuationFiber = Fiber;
 
 pub struct Payloads {
     /// Number of currently occupied slots.

crates/cranelift/src/func_environ.rs

@@ -3502,7 +3502,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         // to resume objects other than `original_contobj`.
         // We make the continuation object that was actually resumed available via
         // `resumed_contobj`, so that subsequent blocks can refer to it.
-        let (tag, resumed_contobj) = {
+        let (resumed_contobj, resumed_contobj_parent, resume_result) = {
             builder.switch_to_block(resume_block);
             builder.append_block_param(resume_block, self.pointer_type());
 
@@ -3523,48 +3523,78 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
             let (_vmctx, result) =
                 generate_builtin_call!(self, builder, tc_resume, [resume_contobj]);
 
+            emit_debug_println!(
+                self,
+                builder,
+                "[resume] libcall finished, result is {:p}",
+                result
+            );
+
             // Now the parent contobj is active
             let parent_contobj = self.typed_continuations_load_parent(builder, resume_contobj);
             vmctx.set_active_continuation(self, builder, parent_contobj);
 
-            // The result encodes whether the return happens via ordinary
-            // means or via a suspend. If the high bit is set, then it is
-            // interpreted as the return happened via a suspend, and the
-            // remainder of the integer is to be interpreted as the index
-            // of the control tag that was supplied to the suspend.
-            let signal_mask = 0xf000_0000;
-            let inverse_signal_mask = 0x0fff_ffff;
-            let signal = builder.ins().band_imm(result, signal_mask);
-            let tag = builder.ins().band_imm(result, inverse_signal_mask);
-
-            // Description of results:
-            // * The `base_addr` is the base address of VM context.
-            // * The `signal` is an encoded boolean indicating whether
-            // the `resume` returned ordinarily or via a suspend
-            // instruction.
-            // * The `tag` is the index of the control tag supplied to
-            // suspend (only valid if `signal` is 1).
-
-            // Test the signal bit.
-            let is_zero = builder.ins().icmp_imm(IntCC::Equal, signal, 0);
-
-            // Jump to the return block if the signal is 0, otherwise
-            // jump to the suspend block.
+            // The `result` is a value of type wasmtime_fibre::SwitchReason,
+            // using the encoding described at its definition.
+            // Thus, the first 32 bit encode the discriminant, and the
+            // subsequent 32 bit encode the tag if suspending, or 0 otherwise.
+            // Thus, when returning, the overall u64 should be zero.
+            let return_discriminant =
+                wasmtime_continuations::SwitchReasonEnum::Return.discriminant_val();
+            debug_assert_eq!(return_discriminant, 0);
+
+            // If these two assumptions don't hold anymore, the code here becomes invalid.
+            debug_assert_eq!(
+                std::mem::size_of::<wasmtime_continuations::types::switch_reason::Discriminant>(),
+                4
+            );
+            debug_assert_eq!(
+                std::mem::size_of::<wasmtime_continuations::types::switch_reason::Data>(),
+                4
+            );
+
+            if cfg!(debug_assertions) {
+                let discriminant = builder.ins().ireduce(I32, result);
+                emit_debug_println!(self, builder, "[resume] discriminant is {}", discriminant);
+            }
+
+            // Jump to the return block if the result is 0, otherwise jump to
+            // the suspend block.
             builder
                 .ins()
-                .brif(is_zero, return_block, &[], suspend_block, &[]);
+                .brif(result, suspend_block, &[], return_block, &[]);
 
             // We do not seal this block, yet, because the effect forwarding block has a back edge to it
-            (tag, resume_contobj)
+            (resume_contobj, parent_contobj, result)
         };
 
         // Suspend block.
-        {
+        let tag = {
             builder.switch_to_block(suspend_block);
             builder.seal_block(suspend_block);
 
+            let discriminant = builder.ins().ireduce(I32, resume_result);
+            let discriminant_size_bytes =
+                std::mem::size_of::<wasmtime_continuations::types::switch_reason::Discriminant>();
+
+            if cfg!(debug_assertions) {
+                let suspend_discriminant =
+                    wasmtime_continuations::SwitchReasonEnum::Suspend.discriminant_val();
+                let suspend_discriminant = builder.ins().iconst(I32, suspend_discriminant as i64);
+                emit_debug_assert_eq!(self, builder, discriminant, suspend_discriminant);
+            }
+
+            let tag = builder
+                .ins()
+                .ushr_imm(resume_result, discriminant_size_bytes as i64 * 8);
+            let tag = builder.ins().ireduce(I32, tag);
+
+            emit_debug_println!(self, builder, "[resume] in suspend block, tag is {}", tag);
+
             // We need to terminate this block before being allowed to switch to another one
             builder.ins().jump(switch_block, &[]);
+
+            tag
         };
 
         // Now, construct blocks for the three continuations:
@@ -3634,11 +3664,9 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         {
             builder.switch_to_block(forwarding_block);
 
-            let parent_contobj = self.typed_continuations_load_parent(builder, resumed_contobj);
-
             builder
                 .ins()
-                .trapz(parent_contobj, ir::TrapCode::UnhandledTag);
+                .trapz(resumed_contobj_parent, ir::TrapCode::UnhandledTag);
 
             // We suspend, thus deferring handling to the parent.
             // We do nothing about tag *parameters*, these remain unchanged within the
@@ -3652,7 +3680,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
             // suspended.
             self.typed_continuations_forward_tag_return_values(
                 builder,
-                parent_contobj,
+                resumed_contobj_parent,
                 resumed_contobj,
             );
 

crates/environ/src/builtin.rs

@@ -54,8 +54,8 @@ macro_rules! foreach_builtin_function {
 
             /// Creates a new continuation from a funcref.
             tc_cont_new(vmctx: vmctx, r: pointer, param_count: i64, result_count: i64) -> pointer;
-            /// Resumes a continuation.
-            tc_resume(vmctx: vmctx, contobj: pointer) -> i32;
+            /// Resumes a continuation. The result value is of type wasmtime_fibre::SwitchReason.
+            tc_resume(vmctx: vmctx, contobj: pointer) -> i64;
             /// Suspends a continuation.
             tc_suspend(vmctx: vmctx, tag: i32);
             /// Returns the continuation object corresponding to the given continuation reference.