Naive effect forwarding (#57)

This PR implements naive effect forwarding:

When generating code for `resume`, we ensure that when asked to handle an effect not considered by the `resume` instruction under consideration, we simply `suspend` and let our own parent handle the unknown effect. Once resuming from that we use a back edge to the code that invokes `resume` on the current fiber (i.e., after moving up in the handler chain to find a handler, we move down in the chain back to where the original `suspend` of the unknown effect happened).

Payloads are handled as follows:
1. When moving "up", we don't need to consider payloads: They are stored in the "payload buffer" of the `VMContext` and left untouched while moving upwards in the handler chain until the right handler is found and accesses the data.
2. When moving down, we need to deal with the "tag return payloads" that were provided via `resume` or `cont.bind` further up in the chain. These payloads are stored in continuation objects directly. Before resuming into our own child, we thus swap the payload pointers of the current continuation object and the child.

cranelift/filetests/src/test_wasm/env.rs

@@ -705,11 +705,8 @@ impl<'a> FuncEnvironment for FuncEnv<'a> {
         builder: &mut cranelift_frontend::FunctionBuilder,
         state: &cranelift_wasm::FuncTranslationState,
         cont: ir::Value,
-        call_arg_types: &[wasmtime_types::WasmType],
-        call_args: &[ir::Value],
     ) -> cranelift_wasm::WasmResult<(ir::Value, ir::Value, ir::Value)> {
-        self.inner
-            .translate_resume(builder, state, cont, call_arg_types, call_args)
+        self.inner.translate_resume(builder, state, cont)
     }
 
     /// TODO(dhil): write documentation.
@@ -729,7 +726,7 @@ impl<'a> FuncEnvironment for FuncEnv<'a> {
         &mut self,
         builder: &mut cranelift_frontend::FunctionBuilder,
         state: &cranelift_wasm::FuncTranslationState,
-        tag_index: u32,
+        tag_index: ir::Value,
     ) -> ir::Value {
         return self.inner.translate_suspend(builder, state, tag_index);
     }
@@ -850,4 +847,18 @@ impl<'a> FuncEnvironment for FuncEnv<'a> {
             .inner
             .typed_continuations_load_tag_return_values(builder, contobj, valtypes);
     }
+
+    /// TODO
+    fn typed_continuations_forward_tag_return_values(
+        &mut self,
+        builder: &mut cranelift_frontend::FunctionBuilder,
+        parent_contobj: ir::Value,
+        child_contobj: ir::Value,
+    ) {
+        return self.inner.typed_continuations_forward_tag_return_values(
+            builder,
+            parent_contobj,
+            child_contobj,
+        );
+    }
 }

cranelift/wasm/src/code_translator.rs

@@ -2507,43 +2507,62 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let (original_contref, call_args) = state.peekn(arity + 1).split_last().unwrap();
             let original_contobj =
                 environ.typed_continuations_cont_ref_get_cont_obj(builder, *original_contref);
-            let call_arg_types = environ.continuation_arguments(*type_index).to_vec();
-
-            // Now, we generate the call instruction.
-            let (base_addr, signal, tag) = environ.translate_resume(
-                builder,
-                state,
-                original_contobj,
-                &call_arg_types,
-                call_args,
-            )?;
-            // 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).
 
+            if call_args.len() > 0 {
+                let count = builder.ins().iconst(I32, call_args.len() as i64);
+                environ.typed_continuations_store_resume_args(
+                    builder,
+                    call_args,
+                    count,
+                    original_contobj,
+                );
+            }
             // Pop the `resume_args` off the stack.
             state.popn(arity + 1);
 
-            // Now, construct blocks for the three continuations:
-            // 1) `resume` returned normally.
-            // 2) `resume` returned via a suspend.
-            // 3) `resume` is forwarding (TODO)
-
-            // Test the signal bit.
-            let is_zero = builder.ins().icmp_imm(IntCC::Equal, signal, 0);
+            let resume_block = builder.create_block();
             let return_block = crate::translation_utils::return_block(builder, environ)?;
             let suspend_block = crate::translation_utils::suspend_block(builder, environ)?;
             let switch_block = builder.create_block();
-            // Jump to the return block if the signal is 0, otherwise
-            // jump to the suspend block.
-            canonicalise_brif(builder, is_zero, return_block, &[], suspend_block, &[]);
+
+            builder.ins().jump(resume_block, &[original_contobj]);
+
+            let (base_addr, tag, resumed_contobj) = {
+                builder.switch_to_block(resume_block);
+                builder.append_block_param(resume_block, environ.pointer_type());
+
+                // The continuation object to actually call resume on
+                let resume_contobj = builder.block_params(resume_block)[0];
+
+                // Now, we generate the call instruction.
+                let (base_addr, signal, tag) =
+                    environ.translate_resume(builder, state, resume_contobj)?;
+                // 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).
+
+                // Now, construct blocks for the three continuations:
+                // 1) `resume` returned normally.
+                // 2) `resume` returned via a suspend.
+                // 3) `resume` is forwarding
+
+                // 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.
+                canonicalise_brif(builder, is_zero, return_block, &[], suspend_block, &[]);
+
+                // We do not seal this block, yet, because the effect forwarding block has a back edge to it
+                (base_addr, tag, resume_contobj)
+            };
 
             // Next, build the suspend block.
-            let contref = {
+            let (contref, contobj) = {
                 builder.switch_to_block(suspend_block);
                 builder.seal_block(suspend_block);
 
@@ -2566,7 +2585,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
 
                 // We need to terminate this block before being allowed to switch to another one
                 builder.ins().jump(switch_block, &[]);
-                contref
+                (contref, contobj)
             };
 
             // Strategy:
@@ -2609,19 +2628,37 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             // the Switch structure and created the blocks it jumps
             // to.
 
-            let forwarding_case =
+            let forwarding_block =
                 crate::translation_utils::resumetable_forwarding_block(builder, environ)?;
+            {
+                builder.switch_to_block(forwarding_block);
+
+                // We suspend, thus deferring handling to the parent.
+                // We do nothing about tag *parameters, these remain unchanged within the
+                // payload buffer associcated with the whole VMContext.
+                environ.translate_suspend(builder, state, tag);
+
+                // When reaching this point, the parent handler has just invoked `resume`.
+                // We propagate the tag return values to the child (i.e., `contobj`).
+                let parent_contobj =
+                    environ.typed_continuations_load_continuation_object(builder, base_addr);
+                environ.typed_continuations_forward_tag_return_values(
+                    builder,
+                    parent_contobj,
+                    contobj,
+                );
+
+                builder.ins().jump(resume_block, &[contobj]);
+                builder.seal_block(resume_block);
+            }
 
             // Switch block (where the actual switching logic is
             // emitted to).
             {
                 builder.switch_to_block(switch_block);
-                switch.emit(builder, tag, forwarding_case);
+                switch.emit(builder, tag, forwarding_block);
                 builder.seal_block(switch_block);
-                builder.switch_to_block(forwarding_case);
-                builder.seal_block(forwarding_case);
-                // TODO: emit effect forwarding logic.
-                builder.ins().trap(ir::TrapCode::UnreachableCodeReached);
+                builder.seal_block(forwarding_block);
 
                 // We can only seal the blocks we generated for each
                 // tag now, after switch.emit ran.
@@ -2640,7 +2677,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                 let values = environ.typed_continuations_load_return_values(
                     builder,
                     &returns,
-                    original_contobj,
+                    resumed_contobj,
                 );
 
                 // The continuation has returned and all `ContinuationReferences`
@@ -2660,7 +2697,8 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             environ.typed_continuations_store_payloads(builder, &param_types, params);
             state.popn(param_count);
 
-            let vmctx = environ.translate_suspend(builder, state, *tag_index);
+            let tag_index_val = builder.ins().iconst(I32, *tag_index as i64);
+            let vmctx = environ.translate_suspend(builder, state, tag_index_val);
 
             let contobj = environ.typed_continuations_load_continuation_object(builder, vmctx);
 

cranelift/wasm/src/environ/dummy.rs

@@ -711,8 +711,6 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
         _builder: &mut FunctionBuilder,
         _state: &FuncTranslationState,
         _cont: ir::Value,
-        _call_arg_types: &[WasmType],
-        _call_args: &[ir::Value],
     ) -> WasmResult<(ir::Value, ir::Value, ir::Value)> {
         todo!()
     }
@@ -731,7 +729,7 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
         &mut self,
         _builder: &mut FunctionBuilder,
         _state: &FuncTranslationState,
-        _tag_index: u32,
+        _tag_index: ir::Value,
     ) -> ir::Value {
         todo!()
     }
@@ -769,6 +767,16 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
         todo!()
     }
 
+    /// TODO
+    fn typed_continuations_forward_tag_return_values(
+        &mut self,
+        _builder: &mut FunctionBuilder,
+        _parent_contobj: ir::Value,
+        _child_contobj: ir::Value,
+    ) {
+        todo!()
+    }
+
     fn typed_continuations_store_resume_args(
         &mut self,
         _builder: &mut FunctionBuilder,

cranelift/wasm/src/environ/spec.rs

@@ -644,8 +644,6 @@ pub trait FuncEnvironment: TargetEnvironment {
         builder: &mut FunctionBuilder,
         state: &FuncTranslationState,
         cont: ir::Value,
-        call_arg_types: &[wasmtime_types::WasmType],
-        call_args: &[ir::Value],
     ) -> WasmResult<(ir::Value, ir::Value, ir::Value)>;
 
     /// TODO(dhil): write documentation.
@@ -662,7 +660,7 @@ pub trait FuncEnvironment: TargetEnvironment {
         &mut self,
         builder: &mut FunctionBuilder,
         state: &FuncTranslationState,
-        tag_index: u32,
+        tag_index: ir::Value,
     ) -> ir::Value;
 
     /// TODO
@@ -694,6 +692,14 @@ pub trait FuncEnvironment: TargetEnvironment {
         valtypes: &[wasmtime_types::WasmType],
     ) -> std::vec::Vec<ir::Value>;
 
+    /// TODO
+    fn typed_continuations_forward_tag_return_values(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        parent_contobj: ir::Value,
+        child_contobj: ir::Value,
+    );
+
     /// TODO
     fn typed_continuations_store_payloads(
         &mut self,

crates/cranelift/src/func_environ.rs

@@ -2553,8 +2553,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         builder: &mut FunctionBuilder,
         _state: &FuncTranslationState,
         contobj: ir::Value,
-        _call_arg_types: &[WasmType],
-        call_args: &[ir::Value],
     ) -> WasmResult<(ir::Value, ir::Value, ir::Value)> {
         // Strategy:
         //
@@ -2564,11 +2562,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         //
         // Second: Call the `resume` builtin
 
-        if call_args.len() > 0 {
-            let count = builder.ins().iconst(I32, call_args.len() as i64);
-            self.typed_continuations_store_resume_args(builder, call_args, count, contobj);
-        }
-
         let (vmctx, result) = generate_builtin_call!(self, builder, resume, [contobj]);
 
         // The result encodes whether the return happens via ordinary
@@ -2601,10 +2594,8 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         &mut self,
         builder: &mut FunctionBuilder,
         _state: &FuncTranslationState,
-        tag_index: u32,
+        tag_index: ir::Value,
     ) -> ir::Value {
-        let tag_index = builder.ins().iconst(I32, tag_index as i64);
-
         // Returns the vmctx
         return generate_builtin_call_no_return_val!(self, builder, suspend, [tag_index]);
     }
@@ -2701,6 +2692,20 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         values
     }
 
+    fn typed_continuations_forward_tag_return_values(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        parent_contobj: ir::Value,
+        child_contobj: ir::Value,
+    ) {
+        generate_builtin_call_no_return_val!(
+            self,
+            builder,
+            cont_obj_forward_tag_return_values_buffer,
+            [parent_contobj, child_contobj]
+        );
+    }
+
     /// TODO
     fn typed_continuations_cont_ref_get_cont_obj(
         &mut self,

crates/environ/src/builtin.rs

@@ -106,6 +106,12 @@ macro_rules! foreach_builtin_function {
             cont_obj_get_tag_return_values_buffer(vmctx: vmctx, contobj: pointer, expected_count : i32) -> pointer;
             /// Deallocated the tag return value buffer within the continuation object.
             cont_obj_deallocate_tag_return_values_buffer(vmctx: vmctx, contobj: pointer);
+            /// Sets the tag return values of `child_contobj` to those of `parent_contobj`.
+            /// This is implemented by exchanging the pointers to the underlying buffers.
+            /// `child_contobj` must not currently have a tag return value buffer.
+            /// `parent_contobj` may or may not have one.
+            cont_obj_forward_tag_return_values_buffer(vmctx: vmctx, parent_contobj: pointer, child_contobj : pointer);
+
             /// TODO
             drop_cont_obj(vmctx: vmctx, contobj: pointer);
 

crates/runtime/src/continuation.rs

@@ -174,6 +174,21 @@ pub fn cont_obj_get_tag_return_values_buffer(
     return payloads.data;
 }
 
+/// TODO
+pub fn cont_obj_forward_tag_return_values_buffer(
+    parent: *mut ContinuationObject,
+    child: *mut ContinuationObject,
+) {
+    let parent = unsafe { parent.as_mut().unwrap() };
+    let child = unsafe { child.as_mut().unwrap() };
+    assert!(parent.state == State::Invoked);
+    assert!(child.state == State::Invoked);
+
+    assert!(child.tag_return_values.is_none());
+
+    child.tag_return_values = parent.tag_return_values.take()
+}
+
 /// TODO
 pub fn cont_obj_deallocate_tag_return_values_buffer(obj: *mut ContinuationObject) {
     let obj = unsafe { obj.as_mut().unwrap() };

crates/runtime/src/libcalls.rs

@@ -861,6 +861,17 @@ fn cont_obj_get_tag_return_values_buffer(
     ) as *mut u8
 }
 
+fn cont_obj_forward_tag_return_values_buffer(
+    _instance: &mut Instance,
+    parent_contobj: *mut u8,
+    child_contobj: *mut u8,
+) {
+    crate::continuation::cont_obj_forward_tag_return_values_buffer(
+        parent_contobj as *mut crate::continuation::ContinuationObject,
+        child_contobj as *mut crate::continuation::ContinuationObject,
+    );
+}
+
 fn cont_obj_deallocate_tag_return_values_buffer(_instance: &mut Instance, contobj: *mut u8) {
     crate::continuation::cont_obj_deallocate_tag_return_values_buffer(
         contobj as *mut crate::continuation::ContinuationObject,