cranelift-wasm: Add support for translating Wasm tail calls to CLIF (#6760)

cranelift/filetests/src/test_wasm/env.rs

@@ -400,6 +400,27 @@ impl<'a> FuncEnvironment for FuncEnv<'a> {
         )
     }
 
+    fn translate_return_call_indirect(
+        &mut self,
+        builder: &mut cranelift_frontend::FunctionBuilder,
+        table_index: cranelift_wasm::TableIndex,
+        table: ir::Table,
+        sig_index: TypeIndex,
+        sig_ref: ir::SigRef,
+        callee: ir::Value,
+        call_args: &[ir::Value],
+    ) -> cranelift_wasm::WasmResult<()> {
+        self.inner.translate_return_call_indirect(
+            builder,
+            table_index,
+            table,
+            sig_index,
+            sig_ref,
+            callee,
+            call_args,
+        )
+    }
+
     fn translate_memory_grow(
         &mut self,
         pos: cranelift_codegen::cursor::FuncCursor,

cranelift/wasm/src/code_translator.rs

@@ -674,6 +674,69 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             state.popn(num_args);
             state.pushn(inst_results);
         }
+        /******************************* Tail Calls ******************************************
+         * The tail call instructions pop their arguments from the stack and
+         * then permanently transfer control to their callee. The indirect
+         * version requires environment support (while the direct version can
+         * optionally be hooked but doesn't require it) it interacts with the
+         * VM's runtime state via tables.
+         ************************************************************************************/
+        Operator::ReturnCall { function_index } => {
+            let (fref, num_args) = state.get_direct_func(builder.func, *function_index, environ)?;
+
+            // Bitcast any vector arguments to their default type, I8X16, before calling.
+            let args = state.peekn_mut(num_args);
+            bitcast_wasm_params(
+                environ,
+                builder.func.dfg.ext_funcs[fref].signature,
+                args,
+                builder,
+            );
+
+            environ.translate_return_call(
+                builder.cursor(),
+                FuncIndex::from_u32(*function_index),
+                fref,
+                args,
+            )?;
+
+            state.popn(num_args);
+            state.reachable = false;
+        }
+        Operator::ReturnCallIndirect {
+            type_index,
+            table_index,
+        } => {
+            // `type_index` is the index of the function's signature and
+            // `table_index` is the index of the table to search the function
+            // in.
+            let (sigref, num_args) = state.get_indirect_sig(builder.func, *type_index, environ)?;
+            let table = state.get_or_create_table(builder.func, *table_index, environ)?;
+            let callee = state.pop1();
+
+            // Bitcast any vector arguments to their default type, I8X16, before calling.
+            let args = state.peekn_mut(num_args);
+            bitcast_wasm_params(environ, sigref, args, builder);
+
+            environ.translate_return_call_indirect(
+                builder,
+                TableIndex::from_u32(*table_index),
+                table,
+                TypeIndex::from_u32(*type_index),
+                sigref,
+                callee,
+                state.peekn(num_args),
+            )?;
+
+            state.popn(num_args);
+            state.reachable = false;
+        }
+        Operator::ReturnCallRef { type_index: _ } => {
+            return Err(wasm_unsupported!(
+                "proposed tail-call operator for function references {:?}",
+                op
+            ));
+        }
         /******************************* Memory management ***********************************
          * Memory management is handled by environment. It is usually translated into calls to
          * special functions.
@@ -2158,9 +2221,6 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let b_high = builder.ins().uwiden_high(b);
             state.push1(builder.ins().imul(a_high, b_high));
         }
-        Operator::ReturnCall { .. } | Operator::ReturnCallIndirect { .. } => {
-            return Err(wasm_unsupported!("proposed tail-call operator {:?}", op));
-        }
         Operator::MemoryDiscard { .. } => {
             return Err(wasm_unsupported!(
                 "proposed memory-control operator {:?}",
@@ -2341,12 +2401,6 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             state.push1(builder.ins().iadd(dot32, c));
         }
 
-        Operator::ReturnCallRef { type_index: _ } => {
-            return Err(wasm_unsupported!(
-                "proposed tail-call operator for function references {:?}",
-                op
-            ));
-        }
         Operator::BrOnNull { relative_depth } => {
             let r = state.pop1();
             let (br_destination, inputs) = translate_br_if_args(*relative_depth, state);

cranelift/wasm/src/environ/dummy.rs

@@ -448,6 +448,19 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
             .0)
     }
 
+    fn translate_return_call_indirect(
+        &mut self,
+        _builder: &mut FunctionBuilder,
+        _table_index: TableIndex,
+        _table: ir::Table,
+        _sig_index: TypeIndex,
+        _sig_ref: ir::SigRef,
+        _callee: ir::Value,
+        _call_args: &[ir::Value],
+    ) -> WasmResult<()> {
+        unimplemented!()
+    }
+
     fn translate_call(
         &mut self,
         mut pos: FuncCursor,

cranelift/wasm/src/environ/spec.rs

@@ -170,6 +170,23 @@ pub trait FuncEnvironment: TargetEnvironment {
         index: FuncIndex,
     ) -> WasmResult<ir::FuncRef>;
 
+    /// Translate a `call` WebAssembly instruction at `pos`.
+    ///
+    /// Insert instructions at `pos` for a direct call to the function `callee_index`.
+    ///
+    /// The function reference `callee` was previously created by `make_direct_func()`.
+    ///
+    /// Return the call instruction whose results are the WebAssembly return values.
+    fn translate_call(
+        &mut self,
+        mut pos: FuncCursor,
+        _callee_index: FuncIndex,
+        callee: ir::FuncRef,
+        call_args: &[ir::Value],
+    ) -> WasmResult<ir::Inst> {
+        Ok(pos.ins().call(callee, call_args))
+    }
+
     /// Translate a `call_indirect` WebAssembly instruction at `pos`.
     ///
     /// Insert instructions at `pos` for an indirect call to the function `callee` in the table
@@ -191,23 +208,43 @@ pub trait FuncEnvironment: TargetEnvironment {
         call_args: &[ir::Value],
     ) -> WasmResult<ir::Inst>;
 
-    /// Translate a `call` WebAssembly instruction at `pos`.
+    /// Translate a `return_call` WebAssembly instruction at `pos`.
     ///
-    /// Insert instructions at `pos` for a direct call to the function `callee_index`.
+    /// Insert instructions at `pos` for a direct tail call to the function `callee_index`.
     ///
     /// The function reference `callee` was previously created by `make_direct_func()`.
     ///
     /// Return the call instruction whose results are the WebAssembly return values.
-    fn translate_call(
+    fn translate_return_call(
         &mut self,
         mut pos: FuncCursor,
         _callee_index: FuncIndex,
         callee: ir::FuncRef,
         call_args: &[ir::Value],
-    ) -> WasmResult<ir::Inst> {
-        Ok(pos.ins().call(callee, call_args))
+    ) -> WasmResult<()> {
+        pos.ins().return_call(callee, call_args);
+        Ok(())
     }
 
+    /// Translate a `return_call_indirect` WebAssembly instruction at `pos`.
+    ///
+    /// Insert instructions at `pos` for an indirect tail call to the function
+    /// `callee` in the table `table_index` with WebAssembly signature
+    /// `sig_index`. The `callee` value will have type `i32`.
+    ///
+    /// The signature `sig_ref` was previously created by `make_indirect_sig()`.
+    #[allow(clippy::too_many_arguments)]
+    fn translate_return_call_indirect(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        table_index: TableIndex,
+        table: ir::Table,
+        sig_index: TypeIndex,
+        sig_ref: ir::SigRef,
+        callee: ir::Value,
+        call_args: &[ir::Value],
+    ) -> WasmResult<()>;
+
     /// Translate a `call_ref` WebAssembly instruction at `pos`.
     ///
     /// Insert instructions at `pos` for an indirect call to the

crates/cranelift/src/func_environ.rs

@@ -1761,6 +1761,29 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         self.call_function_unchecked(builder, sig_ref, callee, call_args)
     }
 
+    fn translate_return_call(
+        &mut self,
+        _pos: FuncCursor,
+        _callee_index: FuncIndex,
+        _callee: ir::FuncRef,
+        _call_args: &[ir::Value],
+    ) -> WasmResult<()> {
+        unimplemented!()
+    }
+
+    fn translate_return_call_indirect(
+        &mut self,
+        _builder: &mut FunctionBuilder,
+        _table_index: TableIndex,
+        _table: ir::Table,
+        _sig_index: TypeIndex,
+        _sig_ref: ir::SigRef,
+        _callee: ir::Value,
+        _call_args: &[ir::Value],
+    ) -> WasmResult<()> {
+        unimplemented!()
+    }
+
     fn translate_memory_grow(
         &mut self,
         mut pos: FuncCursor<'_>,