Cranelift: Introduce support for return_call in the interpreter (#5697)

Co-authored-by: Jamey Sharp <jsharp@fastly.com>

cranelift/filetests/filetests/runtests/return-call.clif

@@ -0,0 +1,68 @@
+test interpret
+;; test run
+;; target x86_64
+;; target aarch64
+;; target aarch64 sign_return_address
+;; target aarch64 has_pauth sign_return_address
+;; target s390x
+
+;;;; Test passing `i64`s ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+function %callee_i64(i64) -> i64 tail {
+block0(v0: i64):
+    v1 = iadd_imm.i64 v0, 10
+    return v1
+}
+
+function %call_i64(i64) -> i64 tail {
+    fn0 = %callee_i64(i64) -> i64 tail
+
+block0(v0: i64):
+    return_call fn0(v0)
+}
+; run: %call_i64(10) == 20
+
+;;;; Test colocated tail calls ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+function %colocated_i64(i64) -> i64 tail {
+    fn0 = colocated %callee_i64(i64) -> i64 tail
+
+block0(v0: i64):
+    return_call fn0(v0)
+}
+; run: %colocated_i64(10) == 20
+
+;;;; Test passing `f64`s ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+function %callee_f64(f64) -> f64 tail {
+block0(v0: f64):
+    v1 = f64const 0x10.0
+    v2 = fadd.f64 v0, v1
+    return v2
+}
+
+function %call_f64(f64) -> f64 tail {
+    fn0 = %callee_f64(f64) -> f64 tail
+
+block0(v0: f64):
+    return_call fn0(v0)
+}
+; run: %call_f64(0x10.0) == 0x20.0
+
+;;;; Test passing `i8`s ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+function %callee_i8(i8) -> i8 tail {
+block0(v0: i8):
+    v1 = iconst.i8 0
+    v2 = icmp eq v0, v1
+    return v2
+}
+
+function %call_i8(i8) -> i8 tail {
+    fn0 = %callee_i8(i8) -> i8 tail
+
+block0(v0: i8):
+    return_call fn0(v0)
+}
+; run: %call_i8(1) == 0
+; run: %call_i8(0) == 1

cranelift/interpreter/src/interpreter.rs

@@ -125,6 +125,11 @@ impl<'a> Interpreter<'a> {
                         .set_all(function.dfg.inst_results(inst), returned_arguments);
                     maybe_inst = layout.next_inst(inst)
                 }
+                ControlFlow::ReturnCall(callee, args) => {
+                    self.state.pop_frame();
+                    let rets = self.call(callee, &args)?.unwrap_return();
+                    return Ok(ControlFlow::Return(rets.into()));
+                }
                 ControlFlow::Return(returned_values) => {
                     self.state.pop_frame();
                     return Ok(ControlFlow::Return(returned_values));

cranelift/interpreter/src/step.rs

@@ -34,6 +34,15 @@ fn validate_signature_params(sig: &[AbiParam], args: &[impl Value]) -> bool {
         })
 }
 
+// Helper for summing a sequence of values.
+fn sum<V: Value>(head: V, tail: SmallVec<[V; 1]>) -> ValueResult<i128> {
+    let mut acc = head;
+    for t in tail {
+        acc = Value::add(acc, t)?;
+    }
+    acc.into_int()
+}
+
 /// Interpret a single Cranelift instruction. Note that program traps and interpreter errors are
 /// distinct: a program trap results in `Ok(Flow::Trap(...))` whereas an interpretation error (e.g.
 /// the types of two values are incompatible) results in `Err(...)`.
@@ -267,14 +276,75 @@ where
         }
     };
 
-    // Helper for summing a sequence of values.
+    // Perform a call operation.
-    fn sum<V: Value>(head: V, tail: SmallVec<[V; 1]>) -> ValueResult<i128> {
+    //
-        let mut acc = head;
+    // The returned `ControlFlow` variant is determined by the given function
-        for t in tail {
+    // argument, which should make either a `ControlFlow::Call` or a
-            acc = Value::add(acc, t)?;
+    // `ControlFlow::ReturnCall`.
+    let do_call = |make_ctrl_flow: fn(&'a Function, SmallVec<[V; 1]>) -> ControlFlow<'a, V>|
+     -> Result<ControlFlow<'a, V>, StepError> {
+        let func_ref = if let InstructionData::Call { func_ref, .. } = inst {
+            func_ref
+        } else {
+            unreachable!()
+        };
+
+        let curr_func = state.get_current_function();
+        let ext_data = curr_func
+            .dfg
+            .ext_funcs
+            .get(func_ref)
+            .ok_or(StepError::UnknownFunction(func_ref))?;
+
+        let signature = if let Some(sig) = curr_func.dfg.signatures.get(ext_data.signature) {
+            sig
+        } else {
+            return Ok(ControlFlow::Trap(CraneliftTrap::User(
+                TrapCode::BadSignature,
+            )));
+        };
+
+        let args = args()?;
+
+        // Check the types of the arguments. This is usually done by the verifier, but nothing
+        // guarantees that the user has ran that.
+        let args_match = validate_signature_params(&signature.params[..], &args[..]);
+        if !args_match {
+            return Ok(ControlFlow::Trap(CraneliftTrap::User(
+                TrapCode::BadSignature,
+            )));
         }
-        acc.into_int()
+
-    }
+        Ok(match ext_data.name {
+            // These functions should be registered in the regular function store
+            ExternalName::User(_) | ExternalName::TestCase(_) => {
+                let function = state
+                    .get_function(func_ref)
+                    .ok_or(StepError::UnknownFunction(func_ref))?;
+
+                make_ctrl_flow(function, args)
+            }
+            ExternalName::LibCall(libcall) => {
+                debug_assert_ne!(inst.opcode(), Opcode::ReturnCall, "Cannot tail call to libcalls");
+                let libcall_handler = state.get_libcall_handler();
+
+                // We don't transfer control to a libcall, we just execute it and return the results
+                let res = libcall_handler(libcall, args);
+                let res = match res {
+                    Err(trap) => return Ok(ControlFlow::Trap(CraneliftTrap::User(trap))),
+                    Ok(rets) => rets,
+                };
+
+                // Check that what the handler returned is what we expect.
+                if validate_signature_params(&signature.returns[..], &res[..]) {
+                    ControlFlow::Assign(res)
+                } else {
+                    ControlFlow::Trap(CraneliftTrap::User(TrapCode::BadSignature))
+                }
+            }
+            ExternalName::KnownSymbol(_) => unimplemented!(),
+        })
+    };
 
     // Interpret a Cranelift instruction.
     Ok(match inst.opcode() {
@@ -328,70 +398,9 @@ where
         Opcode::Trapnz => trap_when(arg(0)?.into_bool()?, CraneliftTrap::User(trap_code())),
         Opcode::ResumableTrapnz => trap_when(arg(0)?.into_bool()?, CraneliftTrap::Resumable),
         Opcode::Return => ControlFlow::Return(args()?),
-        Opcode::Call => {
+        Opcode::Call => do_call(ControlFlow::Call)?,
-            let func_ref = if let InstructionData::Call { func_ref, .. } = inst {
-                func_ref
-            } else {
-                unreachable!()
-            };
-
-            let curr_func = state.get_current_function();
-            let ext_data = curr_func
-                .dfg
-                .ext_funcs
-                .get(func_ref)
-                .ok_or(StepError::UnknownFunction(func_ref))?;
-
-            let signature = if let Some(sig) = curr_func.dfg.signatures.get(ext_data.signature) {
-                sig
-            } else {
-                return Ok(ControlFlow::Trap(CraneliftTrap::User(
-                    TrapCode::BadSignature,
-                )));
-            };
-
-            let args = args()?;
-
-            // Check the types of the arguments. This is usually done by the verifier, but nothing
-            // guarantees that the user has ran that.
-            let args_match = validate_signature_params(&signature.params[..], &args[..]);
-            if !args_match {
-                return Ok(ControlFlow::Trap(CraneliftTrap::User(
-                    TrapCode::BadSignature,
-                )));
-            }
-
-            match ext_data.name {
-                // These functions should be registered in the regular function store
-                ExternalName::User(_) | ExternalName::TestCase(_) => {
-                    let function = state
-                        .get_function(func_ref)
-                        .ok_or(StepError::UnknownFunction(func_ref))?;
-
-                    ControlFlow::Call(function, args)
-                }
-                ExternalName::LibCall(libcall) => {
-                    let libcall_handler = state.get_libcall_handler();
-
-                    // We don't transfer control to a libcall, we just execute it and return the results
-                    let res = libcall_handler(libcall, args);
-                    let res = match res {
-                        Err(trap) => return Ok(ControlFlow::Trap(CraneliftTrap::User(trap))),
-                        Ok(rets) => rets,
-                    };
-
-                    // Check that what the handler returned is what we expect.
-                    if validate_signature_params(&signature.returns[..], &res[..]) {
-                        ControlFlow::Assign(res)
-                    } else {
-                        ControlFlow::Trap(CraneliftTrap::User(TrapCode::BadSignature))
-                    }
-                }
-                ExternalName::KnownSymbol(_) => unimplemented!(),
-            }
-        }
         Opcode::CallIndirect => unimplemented!("CallIndirect"),
-        Opcode::ReturnCall => unimplemented!("ReturnCall"),
+        Opcode::ReturnCall => do_call(ControlFlow::ReturnCall)?,
         Opcode::ReturnCallIndirect => unimplemented!("ReturnCallIndirect"),
         Opcode::FuncAddr => unimplemented!("FuncAddr"),
         Opcode::Load
@@ -1288,6 +1297,8 @@ pub enum ControlFlow<'a, V> {
     ContinueAt(Block, SmallVec<[V; 1]>),
     /// Indicates a call the given [Function] with the supplied arguments.
     Call(&'a Function, SmallVec<[V; 1]>),
+    /// Indicates a tail call to the given [Function] with the supplied arguments.
+    ReturnCall(&'a Function, SmallVec<[V; 1]>),
     /// Return from the current function with the given parameters, e.g.: `return [v1, v2]`.
     Return(SmallVec<[V; 1]>),
     /// Stop with a program-generated trap; note that these are distinct from errors that may occur