Naive effect forwarding

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,