reifydb_engine/function/
registry.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4use 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}