reifydb_engine/function/
registry.rs1use std::{collections::HashMap, ops::Deref, sync::Arc};
5
6use crate::function::{AggregateFunction, GeneratorFunction, ScalarFunction};
7
8#[derive(Clone)]
9pub struct Functions(Arc<FunctionsInner>);
10
11impl Functions {
12 pub fn builder() -> FunctionsBuilder {
13 FunctionsBuilder(FunctionsInner {
14 scalars: HashMap::new(),
15 aggregates: HashMap::new(),
16 generators: HashMap::new(),
17 })
18 }
19}
20
21impl Deref for Functions {
22 type Target = FunctionsInner;
23
24 fn deref(&self) -> &Self::Target {
25 &self.0
26 }
27}
28
29#[derive(Clone)]
30pub struct FunctionsInner {
31 scalars: HashMap<String, Arc<dyn Fn() -> Box<dyn ScalarFunction> + Send + Sync>>,
32 aggregates: HashMap<String, Arc<dyn Fn() -> Box<dyn AggregateFunction> + Send + Sync>>,
33 generators: HashMap<String, Arc<dyn Fn() -> Box<dyn GeneratorFunction> + Send + Sync>>,
34}
35
36impl FunctionsInner {
37 pub fn get_aggregate(&self, name: &str) -> Option<Box<dyn AggregateFunction>> {
38 self.aggregates.get(name).map(|func| func())
39 }
40
41 pub fn get_scalar(&self, name: &str) -> Option<Box<dyn ScalarFunction>> {
42 self.scalars.get(name).map(|func| func())
43 }
44
45 pub fn get_generator(&self, name: &str) -> Option<Box<dyn GeneratorFunction>> {
46 self.generators.get(name).map(|func| func())
47 }
48}
49
50pub struct FunctionsBuilder(FunctionsInner);
51
52impl FunctionsBuilder {
53 pub fn register_scalar<F, A>(mut self, name: &str, init: F) -> Self
54 where
55 F: Fn() -> A + Send + Sync + 'static,
56 A: ScalarFunction + 'static,
57 {
58 self.0.scalars.insert(name.to_string(), Arc::new(move || Box::new(init()) as Box<dyn ScalarFunction>));
59
60 self
61 }
62
63 pub fn register_aggregate<F, A>(mut self, name: &str, init: F) -> Self
64 where
65 F: Fn() -> A + Send + Sync + 'static,
66 A: AggregateFunction + 'static,
67 {
68 self.0.aggregates
69 .insert(name.to_string(), Arc::new(move || Box::new(init()) as Box<dyn AggregateFunction>));
70
71 self
72 }
73
74 pub fn register_generator<F, G>(mut self, name: &str, init: F) -> Self
75 where
76 F: Fn() -> G + Send + Sync + 'static,
77 G: GeneratorFunction + 'static,
78 {
79 self.0.generators
80 .insert(name.to_string(), Arc::new(move || Box::new(init()) as Box<dyn GeneratorFunction>));
81
82 self
83 }
84
85 pub fn build(self) -> Functions {
86 Functions(Arc::new(self.0))
87 }
88}