ocas_eval/
function_map.rs

1//! Function registry for user-defined external functions.
2//!
3//! The [`FunctionMap`] allows registering custom functions that can be
4//! called during expression evaluation via [`ExternalFun`](crate::instruction::Instr::ExternalFun)
5//! instructions.
6
7use ocas_core::FastHashMap as HashMap;
8
9use crate::domain::EvaluationDomain;
10
11type ExternalFn<T> = Box<dyn Fn(&[T]) -> T + Send + Sync>;
12
13/// A map of named functions that can be called during evaluation.
14pub struct FunctionMap<T: EvaluationDomain> {
15    entries: Vec<(String, FunctionEntry<T>)>,
16    name_to_idx: HashMap<String, usize>,
17    aliases: HashMap<String, String>,
18}
19
20/// A registered external function.
21pub struct FunctionEntry<T: EvaluationDomain> {
22    /// Number of arguments the function expects.
23    pub arity: usize,
24    func: ExternalFn<T>,
25}
26
27impl<T: EvaluationDomain> FunctionMap<T> {
28    /// Create an empty function map.
29    pub fn new() -> Self {
30        Self {
31            entries: Vec::new(),
32            name_to_idx: HashMap::default(),
33            aliases: HashMap::default(),
34        }
35    }
36
37    /// Register a function with the given name and arity.
38    pub fn register(&mut self, name: &str, arity: usize, func: ExternalFn<T>) {
39        let idx = self.entries.len();
40        self.entries
41            .push((name.to_string(), FunctionEntry { arity, func }));
42        self.name_to_idx.insert(name.to_string(), idx);
43    }
44
45    /// Register an alias for a function name.
46    pub fn register_alias(&mut self, alias: &str, canonical: &str) {
47        self.aliases
48            .insert(alias.to_string(), canonical.to_string());
49    }
50
51    /// Look up a function by name (resolving aliases and case).
52    pub fn resolve(&self, name: &str) -> Option<&FunctionEntry<T>> {
53        self.resolve_idx(name).map(|idx| &self.entries[idx].1)
54    }
55
56    /// Get the index of a function by name.
57    pub fn index_of(&self, name: &str) -> Option<usize> {
58        self.resolve_idx(name)
59    }
60
61    fn resolve_idx(&self, name: &str) -> Option<usize> {
62        if let Some(idx) = self.name_to_idx.get(name) {
63            return Some(*idx);
64        }
65        let lower = name.to_lowercase();
66        if let Some(idx) = self.name_to_idx.get(&lower) {
67            return Some(*idx);
68        }
69        if let Some(canonical) = self.aliases.get(name) {
70            return self.name_to_idx.get(canonical.as_str()).copied();
71        }
72        if let Some(canonical) = self.aliases.get(&lower) {
73            return self.name_to_idx.get(canonical.as_str()).copied();
74        }
75        None
76    }
77
78    /// Call a function by its index in the map.
79    pub fn call_by_index(&self, idx: usize, args: &[T]) -> Option<T> {
80        self.entries.get(idx).map(|(_, entry)| (entry.func)(args))
81    }
82
83    /// Return the number of registered functions.
84    pub fn len(&self) -> usize {
85        self.entries.len()
86    }
87
88    /// Return true if no functions are registered.
89    pub fn is_empty(&self) -> bool {
90        self.entries.is_empty()
91    }
92}
93
94impl<T: EvaluationDomain> Default for FunctionMap<T> {
95    fn default() -> Self {
96        Self::new()
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn register_and_resolve() {
106        let mut map = FunctionMap::<f64>::new();
107        map.register("square", 1, Box::new(|args| args[0] * args[0]));
108        assert!(map.resolve("square").is_some());
109        assert!(map.resolve("Square").is_some());
110        assert!(map.resolve("unknown").is_none());
111    }
112
113    #[test]
114    fn alias_resolution() {
115        let mut map = FunctionMap::<f64>::new();
116        map.register("log", 1, Box::new(|args| args[0].ln()));
117        map.register_alias("ln", "log");
118        assert!(map.resolve("ln").is_some());
119        assert!(map.resolve("Ln").is_some());
120    }
121
122    #[test]
123    fn call_by_index() {
124        let mut map = FunctionMap::<f64>::new();
125        map.register("square", 1, Box::new(|args| args[0] * args[0]));
126        let result = map.call_by_index(0, &[3.0]).unwrap();
127        assert!((result - 9.0).abs() < 1e-10);
128    }
129
130    #[test]
131    fn empty_map() {
132        let map = FunctionMap::<f64>::new();
133        assert!(map.is_empty());
134        assert_eq!(map.len(), 0);
135    }
136}
ocas_eval/function_map.rs

ocas_eval/
function_map.rs
