1 //! Implements the function environment (e.g. a name-to-function mapping) for interpretation.
2 use cranelift_codegen::ir::{FuncRef, Function};
3 use cranelift_entity::{PrimaryMap, entity_impl};
4 use std::collections::HashMap;
5 
6 /// A function store contains all of the functions that are accessible to an interpreter.
7 #[derive(Default, Clone)]
8 pub struct FunctionStore<'a> {
9     functions: PrimaryMap<FuncIndex, &'a Function>,
10     function_names: HashMap<String, FuncIndex>,
11 }
12 
13 /// An opaque reference to a [`Function`] stored in the [FunctionStore].
14 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
15 pub struct FuncIndex(u32);
16 entity_impl!(FuncIndex, "fn");
17 
18 /// This is a helpful conversion for instantiating a store from a single [Function].
19 impl<'a> From<&'a Function> for FunctionStore<'a> {
20     fn from(function: &'a Function) -> Self {
21         let mut store = FunctionStore::default();
22         store.add(function.name.to_string(), function);
23         store
24     }
25 }
26 
27 impl<'a> FunctionStore<'a> {
28     /// Add a function by name.
29     pub fn add(&mut self, name: String, function: &'a Function) {
30         assert!(!self.function_names.contains_key(&name));
31         let index = self.functions.push(function);
32         self.function_names.insert(name, index);
33     }
34 
35     /// Retrieve the index of a function in the function store by its `name`.
36     pub fn index_of(&self, name: &str) -> Option<FuncIndex> {
37         self.function_names.get(name).cloned()
38     }
39 
40     /// Retrieve a function by its index in the function store.
41     pub fn get_by_index(&self, index: FuncIndex) -> Option<&'a Function> {
42         self.functions.get(index).cloned()
43     }
44 
45     /// Retrieve a function by its name.
46     pub fn get_by_name(&self, name: &str) -> Option<&'a Function> {
47         let index = self.index_of(name)?;
48         self.get_by_index(index)
49     }
50 
51     /// Retrieve a function from a [FuncRef] within a [Function]. TODO this should be optimized, if possible, as
52     /// currently it retrieves the function name as a string and performs string matching.
53     pub fn get_from_func_ref(
54         &self,
55         func_ref: FuncRef,
56         function: &Function,
57     ) -> Option<&'a Function> {
58         self.get_by_name(&get_function_name(func_ref, function))
59     }
60 }
61 
62 /// Retrieve a function name from a [FuncRef] within a [Function]. TODO this should be optimized, if possible, as
63 /// currently it retrieves the function name as a string and performs string matching.
64 fn get_function_name(func_ref: FuncRef, function: &Function) -> String {
65     function
66         .stencil
67         .dfg
68         .ext_funcs
69         .get(func_ref)
70         .expect("function to exist")
71         .name
72         .display(Some(&function.params))
73         .to_string()
74 }
75 
76 #[cfg(test)]
77 mod tests {
78     use super::*;
79     use cranelift_codegen::ir::{Signature, UserFuncName};
80     use cranelift_codegen::isa::CallConv;
81 
82     #[test]
83     fn addition() {
84         let mut env = FunctionStore::default();
85         let a = "a";
86         let f = Function::new();
87 
88         env.add(a.to_string(), &f);
89         assert!(env.get_by_name(a).is_some());
90     }
91 
92     #[test]
93     fn nonexistence() {
94         let env = FunctionStore::default();
95         assert!(env.get_by_name("a").is_none());
96     }
97 
98     #[test]
99     fn from() {
100         let name = UserFuncName::testcase("test");
101         let signature = Signature::new(CallConv::Fast);
102         let func = &Function::with_name_signature(name, signature);
103         let env: FunctionStore = func.into();
104         assert_eq!(env.index_of("%test"), Some(FuncIndex::from_u32(0)));
105     }
106 }
107