Improve handling of nested continuations

crates/cranelift/src/func_environ.rs

@@ -2697,7 +2697,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         builder: &mut FunctionBuilder,
         base_addr: ir::Value,
     ) -> ir::Value {
-        let memflags = ir::MemFlags::trusted().with_readonly();
+        let memflags = ir::MemFlags::trusted();
         let offset = i32::try_from(self.offsets.vmctx_typed_continuations_store()).unwrap();
         builder
             .ins()

crates/runtime/src/continuation.rs

@@ -11,6 +11,18 @@ use wasmtime_fibre::{Fiber, FiberStack, Suspend};
 type ContinuationFiber = Fiber<'static, (), u32, ()>;
 type Yield = Suspend<(), u32, ()>;
 
+#[allow(dead_code)]
+const ENABLE_DEBUG_PRINTING: bool = false;
+
+macro_rules! debug_println {
+    ($( $args:expr ),+ ) => {
+        #[cfg(debug_assertions)]
+        if ENABLE_DEBUG_PRINTING {
+            println!($($args),*);
+        }
+    }
+}
+
 struct Payloads {
     length: usize,
     capacity: usize,
@@ -311,9 +323,12 @@ pub fn cont_new(
         tag_return_values: None,
         state: State::Allocated,
     });
+
     // TODO(dhil): we need memory clean up of
     // continuation reference objects.
-    return Box::into_raw(contobj);
+    let pointer = Box::into_raw(contobj);
+    debug_println!("Created contobj @ {:p}", pointer);
+    return pointer;
 }
 
 /// TODO
@@ -328,6 +343,12 @@ pub fn resume(
     let tsp = TopOfStackPointer::as_raw(instance.tsp());
     unsafe { fiber_stack.write_parent(tsp) };
     instance.set_tsp(TopOfStackPointer::from_raw(fiber_stack.top().unwrap()));
+    debug_println!(
+        "Resuming contobj @ {:p}, tsp is {:p}, setting it to {:p}",
+        contobj,
+        tsp,
+        fiber_stack.top().unwrap()
+    );
     unsafe {
         (*(*(*instance.store()).vmruntime_limits())
             .stack_limit
@@ -347,10 +368,24 @@ pub fn resume(
             // entry of the payload store by virtue of using the array
             // calling trampoline to execute it.
 
+            // Restore tsp pointer in instance
+            let _tsp = TopOfStackPointer::as_raw(instance.tsp());
+            let parent = unsafe { (*(*contobj).fiber).stack().parent() };
+            instance.set_tsp(TopOfStackPointer::from_raw(parent));
+
+            debug_println!(
+                "Continuation @ {:p} returned normally, setting tsp from {:p} to {:p}",
+                contobj,
+                _tsp,
+                parent
+            );
+
             unsafe { (*contobj).state = State::Returned };
             Ok(0) // zero value = return normally.
         }
         Err(tag) => {
+            debug_println!("Continuation {:p} suspended", contobj);
+
             // We set the high bit to signal a return via suspend. We
             // encode the tag into the remainder of the integer.
             let signal_mask = 0xf000_0000;
@@ -370,6 +405,11 @@ pub fn resume(
 pub fn suspend(instance: &mut Instance, tag_index: u32) {
     let stack_ptr = TopOfStackPointer::as_raw(instance.tsp());
     let parent = unsafe { stack_ptr.cast::<*mut u8>().offset(-2).read() };
+    debug_println!(
+        "Suspending, setting tsp from {:p} to {:p}",
+        stack_ptr,
+        parent
+    );
     instance.set_tsp(TopOfStackPointer::from_raw(parent));
     let suspend = wasmtime_fibre::unix::Suspend::from_top_ptr(stack_ptr);
     suspend.switch::<(), u32, ()>(wasmtime_fibre::RunResult::Yield(tag_index))

tests/misc_testsuite/typed-continuations/cont_nested1.wast

@@ -0,0 +1,52 @@
+;; test using continuations from within a continuation
+
+(module
+
+  (type $unit_to_unit (func))
+  (type $ct (cont $unit_to_unit))
+
+  (type $g2_res_type (func (result (ref $ct))))
+  (type $g2_res_type_ct (cont $g2_res_type))
+
+  (tag $e1)
+  (tag $e2 (param (ref $ct)))
+
+  (global $marker (mut i32) (i32.const 0))
+
+  (func $update_marker (param $x i32)
+    (i32.add (global.get $marker) (i32.const 1))
+    (i32.mul (local.get $x))
+    (global.set $marker))
+
+  (func $g1
+    (call $update_marker (i32.const 2))
+    (suspend $e1)
+    (call $update_marker (i32.const 3))
+    )
+  (elem declare func $g1)
+
+  (func $g2
+    (local $k1 (ref $ct))
+    (local $k2 (ref $ct))
+    (call $update_marker (i32.const 5))
+
+    (block $on_e1 (result (ref $ct))
+      (resume $ct (tag $e1 $on_e1) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (local.set $k1)
+    (call $update_marker (i32.const 7))
+    (block $on_e1_2 (result (ref $ct))
+      (resume $ct (tag $e1 $on_e1_2) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (local.set $k2)
+    (call $update_marker (i32.const 11))
+    (resume $ct (local.get $k1)))
+  (elem declare func $g2)
+
+
+
+  (func $test (export "test") (result i32)
+    (resume $ct (cont.new $ct (ref.func $g2)))
+    (global.get $marker)))
+
+(assert_return (invoke "test") (i32.const 6_108))

tests/misc_testsuite/typed-continuations/cont_nested2.wast

@@ -0,0 +1,64 @@
+;; Similar to cont_nested1, but with payloads
+
+(module
+
+  (type $int_to_int (func (param i32) (result i32)))
+  (type $ct (cont $int_to_int))
+
+  (type $unit_to_int (func (result i32)))
+  (type $ct_unit_to_int (cont $unit_to_int))
+
+  (tag $e1 (param i32) (result i32))
+
+  (global $marker (mut i32) (i32.const 0))
+
+  ;; (func $update_marker (param $x i32) (result i32)
+  ;;   (i32.add (global.get $marker) (i32.const 1))
+  ;;   (i32.mul (local.get $x))
+  ;;   (global.set $marker)
+  ;;   (global.get $marker))
+
+  (func $scramble (param $x i32) (param $y i32) (result i32)
+    (i32.add (local.get $y) (i32.const 1))
+    (i32.mul (local.get $x))
+  )
+
+  (func $g1 (param $x i32) (result i32)
+    (call $scramble (i32.const 3) (local.get $x))
+    (suspend $e1)
+    (call $scramble (i32.const 5))
+    (i32.add (local.get $x))
+    (global.set $marker)
+    (global.get $marker))
+  (elem declare func $g1)
+
+  (func $g2 (result i32)
+    (local $k1 (ref $ct))
+    (local $v i32)
+
+    (block $on_e1 (result i32 (ref $ct))
+      (resume $ct (tag $e1 $on_e1) (i32.const 7) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (local.set $k1)
+    (call $scramble (i32.const 11)) ;; scramble the value received via $e1 from $g1
+    (local.set $v)
+
+    (block $on_e1_2 (result i32 (ref $ct))
+      (resume $ct (tag $e1 $on_e1_2) (local.get $v) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (drop) ;; drop continuation, we don't intend to resume the second invocation of g1
+    (call $scramble (i32.const 13))
+
+    (resume $ct (local.get $k1))
+    (i32.add (global.get $marker)))
+  (elem declare func $g2)
+
+
+
+  (func $test (export "test") (result i32)
+    (resume $ct_unit_to_int (cont.new $ct_unit_to_int (ref.func $g2)))
+
+    )
+  )
+
+(assert_return (invoke "test") (i32.const 145_670))

tests/misc_testsuite/typed-continuations/cont_nested3.wast

@@ -0,0 +1,40 @@
+;; Minimal test for resuming continuation after its original parent is gone
+
+(module
+
+  (type $unit_to_unit (func))
+  (type $ct (cont $unit_to_unit))
+
+  (type $g2 (func (result (ref $ct))))
+  (type $g2_ct (cont $g2))
+
+  (tag $e1)
+  ;;(tag $e2 (param (ref $ct)))
+
+  (global $marker (mut i32) (i32.const 0))
+
+  ;;(global $orphan (mut (ref $ct)))
+
+  (func $g1
+    (suspend $e1)
+    (global.set $marker (i32.const 100))
+    )
+  (elem declare func $g1)
+
+  (func $g2 (result (ref $ct))
+    (block $on_e1 (result (ref $ct))
+      (resume $ct (tag $e1 $on_e1) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    ;; continuation becomes return value
+    )
+
+  (elem declare func $g2)
+
+
+  (func $test (export "test") (result i32)
+    (resume $g2_ct (cont.new $g2_ct (ref.func $g2)))
+    (resume $ct) ;; resume return value of $g2
+    (global.get $marker))
+  )
+
+(assert_return (invoke "test") (i32.const 100))

tests/misc_testsuite/typed-continuations/cont_nested4.wast

@@ -0,0 +1,44 @@
+;; Minimal test for resuming continuation after its original parent is suspended
+
+(module
+
+  (type $unit_to_unit (func))
+  (type $ct (cont $unit_to_unit))
+
+  ;;(type $g2 (func (result (ref $ct))))
+  ;;(type $g2_ct (cont $g2))
+
+  (tag $e1)
+  (tag $e2 (param (ref $ct)))
+
+  (global $marker (mut i32) (i32.const 0))
+
+  ;;(global $orphan (mut (ref $ct)))
+
+  (func $g1
+    (suspend $e1)
+    (global.set $marker (i32.const 100))
+    )
+  (elem declare func $g1)
+
+  (func $g2
+    (block $on_e1 (result (ref $ct))
+      (resume $ct (tag $e1 $on_e1) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (suspend $e2)
+    ;; continuation becomes return value
+    (unreachable))
+
+  (elem declare func $g2)
+
+
+  (func $test (export "test") (result i32)
+    (block $on_e2 (result (ref $ct) (ref $ct))
+      (resume $ct (tag $e2 $on_e2) (cont.new $ct (ref.func $g2)))
+      (unreachable))
+    (drop) ;; drop the continuation (i.e., for resuming g2)
+    (resume $ct) ;; resume continuation received as payload of $e2 (i.e., continuing execution of $g1)
+    (global.get $marker))
+)
+
+(assert_return (invoke "test") (i32.const 100))

tests/misc_testsuite/typed-continuations/cont_nested5.wast

@@ -0,0 +1,68 @@
+;; test using continuations from within a continuation
+
+
+(module
+
+  (type $unit_to_unit (func))
+  (type $ct (cont $unit_to_unit))
+
+  (type $g2_res_type (func (result (ref $ct))))
+  (type $g2_res_type_ct (cont $g2_res_type))
+
+  (tag $e1)
+  (tag $e2 (param (ref $ct)))
+
+  (global $marker (mut i32) (i32.const 0))
+
+  (func $update_marker (param $x i32)
+    (i32.add (global.get $marker) (i32.const 1))
+    (i32.mul (local.get $x))
+    (global.set $marker))
+
+  (func $g1
+    (call $update_marker (i32.const 2))
+    (suspend $e1)
+    (call $update_marker (i32.const 3))
+    )
+  (elem declare func $g1)
+
+  (func $g2 (result (ref $ct))
+    (local $k1 (ref $ct))
+    (local $k2 (ref $ct))
+    (call $update_marker (i32.const 5))
+
+    (block $on_e1 (result (ref $ct))
+      (resume $ct (tag $e1 $on_e1) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (local.set $k1)
+    (call $update_marker (i32.const 7))
+    (block $on_e1_2 (result (ref $ct))
+      (resume $ct (tag $e1 $on_e1_2) (cont.new $ct (ref.func $g1)))
+      (unreachable))
+    (local.set $k2)
+    (call $update_marker (i32.const 11))
+    (resume $ct (local.get $k1))
+    (call $update_marker (i32.const 13))
+    (local.get $k2)
+    )
+  (elem declare func $g2)
+
+  (func $g3
+    (call $update_marker (i32.const 17))
+    (resume $g2_res_type_ct (cont.new $g2_res_type_ct (ref.func $g2)))
+    (call $update_marker (i32.const 19))
+    (suspend $e2))
+  (elem declare func $g3)
+
+
+  (func $test (export "test") (result i32)
+    (call $update_marker (i32.const 23))
+    (block $on_e2 (result (ref $ct) (ref $ct))
+      (resume $ct (tag $e2 $on_e2) (cont.new $ct (ref.func $g3)))
+      (unreachable))
+    (drop) ;; we won't resume g3, but want the payload
+    (call $update_marker (i32.const 31))
+    (resume $ct)
+    (global.get $marker)))
+
+(assert_return (invoke "test") (i32.const 490_074_902))