cranelift: Add LibCalls to the interpreter (#4782)

* cranelift: Add libcall handlers to interpreter

* cranelift: Fuzz IshlI64 libcall

* cranelift: Revert back to fuzzing udivi64

* cranelift: Use sdiv as a fuzz libcall

* cranelift: Register Sdiv in fuzzgen

* cranelift: Add multiple libcalls to fuzzer

* cranelift: Register a single libcall handler

* cranelift: Simplify args checking in interpreter

* cranelift: Remove unused LibCalls

* cranelift: Cleanup interpreter libcall types

* cranelift: Fix Interpreter Docs

cranelift/interpreter/src/interpreter.rs

@@ -12,10 +12,11 @@ use crate::value::{Value, ValueError};
 use cranelift_codegen::data_value::DataValue;
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::{
-    ArgumentPurpose, Block, FuncRef, Function, GlobalValue, GlobalValueData, Heap, StackSlot, Type,
-    Value as ValueRef,
+    ArgumentPurpose, Block, FuncRef, Function, GlobalValue, GlobalValueData, Heap, LibCall,
+    StackSlot, TrapCode, Type, Value as ValueRef,
 };
 use log::trace;
+use smallvec::SmallVec;
 use std::collections::HashSet;
 use std::convert::{TryFrom, TryInto};
 use std::fmt::Debug;
@@ -191,9 +192,13 @@ pub enum HeapInit {
     FromBacking(HeapBacking),
 }
 
+pub type LibCallValues<V> = SmallVec<[V; 1]>;
+pub type LibCallHandler<V> = fn(LibCall, LibCallValues<V>) -> Result<LibCallValues<V>, TrapCode>;
+
 /// Maintains the [Interpreter]'s state, implementing the [State] trait.
 pub struct InterpreterState<'a> {
     pub functions: FunctionStore<'a>,
+    pub libcall_handler: LibCallHandler<DataValue>,
     pub frame_stack: Vec<Frame<'a>>,
     /// Number of bytes from the bottom of the stack where the current frame's stack space is
     pub frame_offset: usize,
@@ -208,6 +213,7 @@ impl Default for InterpreterState<'_> {
     fn default() -> Self {
         Self {
             functions: FunctionStore::default(),
+            libcall_handler: |_, _| Err(TrapCode::UnreachableCodeReached),
             frame_stack: vec![],
             frame_offset: 0,
             stack: Vec::with_capacity(1024),
@@ -224,6 +230,12 @@ impl<'a> InterpreterState<'a> {
         Self { functions, ..self }
     }
 
+    /// Registers a libcall handler
+    pub fn with_libcall_handler(mut self, handler: LibCallHandler<DataValue>) -> Self {
+        self.libcall_handler = handler;
+        self
+    }
+
     /// Registers a static heap and returns a reference to it
     ///
     /// This heap reference can be used to generate a heap pointer, which
@@ -301,6 +313,10 @@ impl<'a> State<'a, DataValue> for InterpreterState<'a> {
         self.current_frame().function
     }
 
+    fn get_libcall_handler(&self) -> LibCallHandler<DataValue> {
+        self.libcall_handler
+    }
+
     fn push_frame(&mut self, function: &'a Function) {
         if let Some(frame) = self.frame_stack.iter().last() {
             self.frame_offset += frame.function.fixed_stack_size() as usize;
@@ -612,9 +628,11 @@ impl<'a> State<'a, DataValue> for InterpreterState<'a> {
 mod tests {
     use super::*;
     use crate::step::CraneliftTrap;
+    use cranelift_codegen::ir::immediates::Ieee32;
     use cranelift_codegen::ir::types::I64;
     use cranelift_codegen::ir::TrapCode;
     use cranelift_reader::parse_functions;
+    use smallvec::smallvec;
 
     // Most interpreter tests should use the more ergonomic `test interpret` filetest but this
     // unit test serves as a sanity check that the interpreter still works without all of the
@@ -1045,4 +1063,34 @@ mod tests {
 
         assert_eq!(trap, CraneliftTrap::User(TrapCode::IntegerOverflow));
     }
+
+    #[test]
+    fn libcall() {
+        let code = "function %test() -> i64 {
+            fn0 = colocated %CeilF32 (f32) -> f32 fast
+        block0:
+            v1 = f32const 0x0.5
+            v2 = call fn0(v1)
+            return v2
+        }";
+
+        let func = parse_functions(code).unwrap().into_iter().next().unwrap();
+        let mut env = FunctionStore::default();
+        env.add(func.name.to_string(), &func);
+        let state = InterpreterState::default()
+            .with_function_store(env)
+            .with_libcall_handler(|libcall, args| {
+                Ok(smallvec![match (libcall, &args[..]) {
+                    (LibCall::CeilF32, [DataValue::F32(a)]) => DataValue::F32(a.ceil()),
+                    _ => panic!("Unexpected args"),
+                }])
+            });
+
+        let result = Interpreter::new(state)
+            .call_by_name("%test", &[])
+            .unwrap()
+            .unwrap_return();
+
+        assert_eq!(result, vec![DataValue::F32(Ieee32::with_float(1.0))])
+    }
 }

cranelift/interpreter/src/state.rs

@@ -1,6 +1,7 @@
 //! Cranelift instructions modify the state of the machine; the [State] trait describes these
 //! ways this can happen.
 use crate::address::{Address, AddressSize};
+use crate::interpreter::LibCallHandler;
 use cranelift_codegen::data_value::DataValue;
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::{FuncRef, Function, GlobalValue, Heap, StackSlot, Type, Value};
@@ -23,6 +24,8 @@ pub trait State<'a, V> {
     fn get_function(&self, func_ref: FuncRef) -> Option<&'a Function>;
     /// Retrieve a reference to the currently executing [Function].
     fn get_current_function(&self) -> &'a Function;
+    /// Retrieve the handler callback for a [LibCall](cranelift_codegen::ir::LibCall)
+    fn get_libcall_handler(&self) -> LibCallHandler<V>;
     /// Record that an interpreter has called into a new [Function].
     fn push_frame(&mut self, function: &'a Function);
     /// Record that an interpreter has returned from a called [Function].
@@ -129,6 +132,10 @@ where
         unimplemented!()
     }
 
+    fn get_libcall_handler(&self) -> LibCallHandler<V> {
+        unimplemented!()
+    }
+
     fn push_frame(&mut self, _function: &'a Function) {
         unimplemented!()
     }

cranelift/interpreter/src/step.rs

@@ -7,7 +7,8 @@ use crate::value::{Value, ValueConversionKind, ValueError, ValueResult};
 use cranelift_codegen::data_value::DataValue;
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::{
-    types, Block, FuncRef, Function, InstructionData, Opcode, TrapCode, Type, Value as ValueRef,
+    types, AbiParam, Block, ExternalName, FuncRef, Function, InstructionData, Opcode, TrapCode,
+    Type, Value as ValueRef,
 };
 use log::trace;
 use smallvec::{smallvec, SmallVec};
@@ -16,6 +17,14 @@ use std::fmt::Debug;
 use std::ops::RangeFrom;
 use thiserror::Error;
 
+/// Ensures that all types in args are the same as expected by the signature
+fn validate_signature_params(sig: &[AbiParam], args: &[impl Value]) -> bool {
+    args.iter()
+        .map(|r| r.ty())
+        .zip(sig.iter().map(|r| r.value_type))
+        .all(|(a, b)| a == b)
+}
+
 /// 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(...)`.
@@ -25,7 +34,7 @@ pub fn step<'a, V, I>(
     inst_context: I,
 ) -> Result<ControlFlow<'a, V>, StepError>
 where
-    V: Value,
+    V: Value + Debug,
     I: InstructionContext,
 {
     let inst = inst_context.data();
@@ -295,13 +304,65 @@ where
         ),
         Opcode::Return => ControlFlow::Return(args()?),
         Opcode::Call => {
-            if let InstructionData::Call { func_ref, .. } = inst {
-                let function = state
-                    .get_function(func_ref)
-                    .ok_or(StepError::UnknownFunction(func_ref))?;
-                ControlFlow::Call(function, args()?)
+            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"),