Fix confusion caused by overloading of FuncRef

Prior to this change, the interpreter would use an incorrect `FuncRef` for accessing functions from the function store. This is now clarified and fixed by a new type--`FuncIndex`.

cranelift/interpreter/src/environment.rs

@@ -1,52 +1,74 @@
 //! Implements the function environment (e.g. a name-to-function mapping) for interpretation.
-
 use cranelift_codegen::ir::{FuncRef, Function};
+use cranelift_entity::{entity_impl, PrimaryMap};
 use std::collections::HashMap;
 
+/// A function store contains all of the functions that are accessible to an interpreter.
 #[derive(Default, Clone)]
 pub struct FunctionStore<'a> {
-    functions: HashMap<FuncRef, &'a Function>,
+    functions: PrimaryMap<FuncIndex, &'a Function>,
-    function_name_to_func_ref: HashMap<String, FuncRef>,
+    function_names: HashMap<String, FuncIndex>,
 }
 
+/// An opaque reference to a [`Function`](Function) stored in the [FunctionStore].
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub struct FuncIndex(u32);
+entity_impl!(FuncIndex, "fn");
+
+/// This is a helpful conversion for instantiating a store from a single [Function].
 impl<'a> From<&'a Function> for FunctionStore<'a> {
-    fn from(f: &'a Function) -> Self {
+    fn from(function: &'a Function) -> Self {
-        let func_ref = FuncRef::from_u32(0);
+        let mut store = FunctionStore::default();
-        let mut function_name_to_func_ref = HashMap::new();
+        store.add(function.name.to_string(), function);
-        function_name_to_func_ref.insert(f.name.to_string(), func_ref);
+        store
-        let mut functions = HashMap::new();
-        functions.insert(func_ref, f);
-        Self {
-            functions,
-            function_name_to_func_ref,
-        }
     }
 }
 
 impl<'a> FunctionStore<'a> {
     /// Add a function by name.
     pub fn add(&mut self, name: String, function: &'a Function) {
-        let func_ref = FuncRef::with_number(self.function_name_to_func_ref.len() as u32)
+        assert!(!self.function_names.contains_key(&name));
-            .expect("a valid function reference");
+        let index = self.functions.push(function);
-        self.function_name_to_func_ref.insert(name, func_ref);
+        self.function_names.insert(name, index);
-        self.functions.insert(func_ref, function);
     }
 
-    /// Retrieve a reference to a function in the environment by its name.
+    /// Retrieve the index of a function in the function store by its `name`.
-    pub fn index_of(&self, name: &str) -> Option<FuncRef> {
+    pub fn index_of(&self, name: &str) -> Option<FuncIndex> {
-        self.function_name_to_func_ref.get(name).cloned()
+        self.function_names.get(name).cloned()
     }
 
-    /// Retrieve a function by its function reference.
+    /// Retrieve a function by its index in the function store.
-    pub fn get_by_func_ref(&self, func_ref: FuncRef) -> Option<&'a Function> {
+    pub fn get_by_index(&self, index: FuncIndex) -> Option<&'a Function> {
-        self.functions.get(&func_ref).cloned()
+        self.functions.get(index).cloned()
     }
 
     /// Retrieve a function by its name.
     pub fn get_by_name(&self, name: &str) -> Option<&'a Function> {
-        let func_ref = self.index_of(name)?;
+        let index = self.index_of(name)?;
-        self.get_by_func_ref(func_ref)
+        self.get_by_index(index)
     }
+
+    /// Retrieve a function from a [FuncRef] within a [Function]. TODO this should be optimized, if possible, as
+    /// currently it retrieves the function name as a string and performs string matching.
+    pub fn get_from_func_ref(
+        &self,
+        func_ref: FuncRef,
+        function: &Function,
+    ) -> Option<&'a Function> {
+        self.get_by_name(&get_function_name(func_ref, function))
+    }
+}
+
+/// Retrieve a function name from a [FuncRef] within a [Function]. TODO this should be optimized, if possible, as
+/// currently it retrieves the function name as a string and performs string matching.
+fn get_function_name(func_ref: FuncRef, function: &Function) -> String {
+    function
+        .dfg
+        .ext_funcs
+        .get(func_ref)
+        .expect("function to exist")
+        .name
+        .to_string()
 }
 
 #[cfg(test)]
@@ -77,6 +99,6 @@ mod tests {
         let signature = Signature::new(CallConv::Fast);
         let func = &Function::with_name_signature(name, signature);
         let env: FunctionStore = func.into();
-        assert_eq!(env.index_of("%test"), FuncRef::with_number(0));
+        assert_eq!(env.index_of("%test"), Some(FuncIndex::from_u32(0)));
     }
 }

cranelift/interpreter/src/interpreter.rs

@@ -2,7 +2,7 @@
 //!
 //! This module partially contains the logic for interpreting Cranelift IR.
 
-use crate::environment::FunctionStore;
+use crate::environment::{FuncIndex, FunctionStore};
 use crate::frame::Frame;
 use crate::instruction::DfgInstructionContext;
 use crate::state::{MemoryError, State};
@@ -34,22 +34,22 @@ impl<'a> Interpreter<'a> {
         func_name: &str,
         arguments: &[DataValue],
     ) -> Result<ControlFlow<'a, DataValue>, InterpreterError> {
-        let func_ref = self
+        let index = self
             .state
             .functions
             .index_of(func_name)
             .ok_or_else(|| InterpreterError::UnknownFunctionName(func_name.to_string()))?;
-        self.call_by_index(func_ref, arguments)
+        self.call_by_index(index, arguments)
     }
 
     /// Call a function by its index in the [FunctionStore]; this is a proxy for [Interpreter::call].
     pub fn call_by_index(
         &mut self,
-        func_ref: FuncRef,
+        index: FuncIndex,
         arguments: &[DataValue],
     ) -> Result<ControlFlow<'a, DataValue>, InterpreterError> {
-        match self.state.get_function(func_ref) {
+        match self.state.functions.get_by_index(index) {
-            None => Err(InterpreterError::UnknownFunctionReference(func_ref)),
+            None => Err(InterpreterError::UnknownFunctionIndex(index)),
             Some(func) => self.call(func, arguments),
         }
     }
@@ -98,8 +98,9 @@ impl<'a> Interpreter<'a> {
                         .set_all(function.dfg.block_params(block), block_arguments.to_vec());
                     maybe_inst = layout.first_inst(block)
                 }
-                ControlFlow::Call(function, arguments) => {
+                ControlFlow::Call(called_function, arguments) => {
-                    let returned_arguments = self.call(function, &arguments)?.unwrap_return();
+                    let returned_arguments =
+                        self.call(called_function, &arguments)?.unwrap_return();
                     self.state
                         .current_frame_mut()
                         .set_all(function.dfg.inst_results(inst), returned_arguments);
@@ -123,8 +124,8 @@ pub enum InterpreterError {
     StepError(#[from] StepError),
     #[error("reached an unreachable statement")]
     Unreachable,
-    #[error("unknown function reference (has it been added to the function store?): {0}")]
+    #[error("unknown function index (has it been added to the function store?): {0}")]
-    UnknownFunctionReference(FuncRef),
+    UnknownFunctionIndex(FuncIndex),
     #[error("unknown function with name (has it been added to the function store?): {0}")]
     UnknownFunctionName(String),
     #[error("value error")]
@@ -176,7 +177,8 @@ impl<'a> InterpreterState<'a> {
 
 impl<'a> State<'a, DataValue> for InterpreterState<'a> {
     fn get_function(&self, func_ref: FuncRef) -> Option<&'a Function> {
-        self.functions.get_by_func_ref(func_ref)
+        self.functions
+            .get_from_func_ref(func_ref, self.frame_stack.last().unwrap().function)
     }
     fn push_frame(&mut self, function: &'a Function) {
         self.frame_stack.push(Frame::new(function));
@@ -273,6 +275,40 @@ mod tests {
         assert_eq!(result, vec![DataValue::B(true)])
     }
 
+    // This test verifies that functions can refer to each other using the function store. A double indirection is
+    // required, which is tricky to get right: a referenced function is a FuncRef when called but a FuncIndex inside the
+    // function store. This test would preferably be a CLIF filetest but the filetest infrastructure only looks at a
+    // single function at a time--we need more than one function in the store for this test.
+    #[test]
+    fn function_references() {
+        let code = "
+        function %child(i32) -> i32 {
+        block0(v0: i32):
+            v1 = iadd_imm v0, -1
+            return v1
+        }
+        
+        function %parent(i32) -> i32 {
+            fn42 = %child(i32) -> i32
+        block0(v0: i32):
+            v1 = iadd_imm v0, 1
+            v2 = call fn42(v1)
+            return v2
+        }";
+
+        let mut env = FunctionStore::default();
+        let funcs = parse_functions(code).unwrap().to_vec();
+        funcs.iter().for_each(|f| env.add(f.name.to_string(), f));
+
+        let state = InterpreterState::default().with_function_store(env);
+        let result = Interpreter::new(state)
+            .call_by_name("%parent", &[DataValue::I32(0)])
+            .unwrap()
+            .unwrap_return();
+
+        assert_eq!(result, vec![DataValue::I32(0)])
+    }
+
     #[test]
     fn state_heap_roundtrip() -> Result<(), MemoryError> {
         let mut state = InterpreterState::default();