Skip to main content

mpl_lang/
linker.rs

1//! Builtin and custom function linking for `MPL`
2use std::{
3    borrow::Borrow,
4    collections::HashMap,
5    fmt::{Display, Write as _},
6};
7
8use serde::{
9    Serialize,
10    ser::{SerializeMap, SerializeStruct as _},
11};
12
13use crate::types::{BucketType, ComputeType, MapType, TagsType, TimeType};
14
15#[derive(Debug, Clone, serde::Serialize)]
16/// A function argument
17pub enum ArgType {
18    /// A floating point argument
19    Float,
20    /// A enum argument, the value can be any of the values
21    Enum(&'static [&'static str]),
22    /// A repeated argument
23    Repeated {
24        /// Type of the repeated argument
25        typ: Box<ArgType>,
26        /// Minimum number of repetitions
27        min: usize,
28        /// Maximum number of repetitions
29        max: Option<usize>,
30    },
31    /// The argument can be one of the following types
32    OneOf(Vec<ArgType>),
33    /// Optional argument
34    Optional(Box<ArgType>),
35}
36impl Display for ArgType {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        match self {
39            ArgType::Float => write!(f, "float"),
40            ArgType::Enum(values) => write!(f, "enum({})", values.join(", ")),
41            ArgType::Repeated { typ, min, max } => {
42                write!(f, "repeated({typ}")?;
43                if *min > 0 {
44                    write!(f, ", min={min}")?;
45                }
46                if let Some(max) = max {
47                    write!(f, ", max={max}")?;
48                }
49                write!(f, ")")
50            }
51            ArgType::OneOf(types) => write!(
52                f,
53                "one_of({})",
54                types
55                    .iter()
56                    .map(ToString::to_string)
57                    .collect::<Vec<String>>()
58                    .join(", ")
59            ),
60            ArgType::Optional(typ) => write!(f, "[{typ}]"),
61        }
62    }
63}
64
65#[derive(Debug, Clone, serde::Serialize)]
66/// A argument to a function
67pub struct Arg {
68    /// Name of the argument
69    pub name: &'static str,
70    /// Type of the argument
71    pub typ: ArgType,
72}
73impl Display for Arg {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        write!(f, "{}: {}", self.name, self.typ)
76    }
77}
78impl Arg {
79    /// Creates a new argument
80    #[must_use]
81    pub const fn new(name: &'static str, typ: ArgType) -> Self {
82        Self { name, typ }
83    }
84}
85
86/// Trait for functions
87pub trait FunctionTrait {
88    /// Documentation of the function
89    fn doc(&self) -> &str;
90    /// Arguments to the function
91    fn args(&self) -> Vec<Arg>;
92    /// Creates the description for the function
93    fn documentation(&self, name: &FunctionId) -> String {
94        let args = self
95            .args()
96            .iter()
97            .map(ToString::to_string)
98            .collect::<Vec<_>>()
99            .join(", ");
100        let doc = self.doc();
101        let e = if name.0.contains('*') { "__" } else { "**" };
102        if args.is_empty() {
103            format!(
104                r"{e}{name}{e}:
105{doc}"
106            )
107        } else {
108            format!(
109                r"{e}{name}{e}({args}):
110
111{doc}"
112            )
113        }
114    }
115}
116
117/// Module identifier
118#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize)]
119pub(crate) struct ModuleId(pub(crate) String);
120
121impl std::fmt::Display for ModuleId {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        write!(f, "{}", self.0)
124    }
125}
126
127impl ModuleId {
128    pub(crate) fn new(name: &str) -> Self {
129        ModuleId(name.to_string())
130    }
131}
132
133impl Borrow<str> for ModuleId {
134    fn borrow(&self) -> &str {
135        &self.0
136    }
137}
138
139/// Function identifier
140#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize)]
141pub struct FunctionId(pub(crate) String);
142
143impl std::fmt::Display for FunctionId {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        write!(f, "{}", self.0)
146    }
147}
148impl FunctionId {
149    pub(crate) fn new(name: &str) -> Self {
150        FunctionId(name.to_string())
151    }
152}
153
154impl Borrow<str> for FunctionId {
155    fn borrow(&self) -> &str {
156        &self.0
157    }
158}
159
160pub(crate) struct Function {
161    pub(crate) module_path: Vec<ModuleId>,
162    pub(crate) name: FunctionId,
163}
164
165impl std::fmt::Display for Function {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        for m in &self.module_path {
168            write!(f, "{m}::")?;
169        }
170        write!(f, "{}", self.name)
171    }
172}
173
174/// Module definition
175pub struct Module {
176    pub(crate) name: ModuleId,
177    pub(crate) doc: &'static str,
178    pub(crate) align_functions: HashMap<FunctionId, AlignFunction>,
179    pub(crate) mapping_functions: HashMap<FunctionId, MapFunction>,
180    pub(crate) group_functions: HashMap<FunctionId, GroupFunction>,
181    pub(crate) bucket_functions: HashMap<FunctionId, BucketType>,
182    pub(crate) compute_functions: HashMap<FunctionId, ComputeFunction>,
183    pub(crate) submodules: HashMap<ModuleId, Module>,
184}
185
186impl std::fmt::Display for Module {
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        write!(f, "Module: {}", self.name)
189    }
190}
191
192struct FunctionSerializer<'a, F: FunctionTrait>(&'a F);
193
194impl<F> Serialize for FunctionSerializer<'_, F>
195where
196    F: FunctionTrait,
197{
198    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
199    where
200        S: serde::Serializer,
201    {
202        let mut state = serializer.serialize_struct("Function", 2)?;
203        state.serialize_field("doc", self.0.doc())?;
204        state.serialize_field("args", &self.0.args())?;
205        state.end()
206    }
207}
208
209struct FunctionMapSerializer<'a, F: FunctionTrait>(&'a HashMap<FunctionId, F>);
210
211impl<F> Serialize for FunctionMapSerializer<'_, F>
212where
213    F: FunctionTrait,
214{
215    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
216    where
217        S: serde::Serializer,
218    {
219        let mut state = serializer.serialize_map(Some(self.0.len()))?;
220        for (id, func) in self.0 {
221            state.serialize_entry(&id.to_string(), &FunctionSerializer(func))?;
222        }
223        state.end()
224    }
225}
226
227impl Serialize for Module {
228    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
229    where
230        S: serde::Serializer,
231    {
232        let mut state = serializer.serialize_struct("Module", 8)?;
233        state.serialize_field("name", &self.name)?;
234        state.serialize_field("doc", &self.doc)?;
235        state.serialize_field(
236            "align_functions",
237            &FunctionMapSerializer(&self.align_functions),
238        )?;
239        state.serialize_field(
240            "mapping_functions",
241            &FunctionMapSerializer(&self.mapping_functions),
242        )?;
243        state.serialize_field(
244            "group_functions",
245            &FunctionMapSerializer(&self.group_functions),
246        )?;
247        state.serialize_field(
248            "compute_functions",
249            &FunctionMapSerializer(&self.compute_functions),
250        )?;
251        state.serialize_field(
252            "bucket_functions",
253            &FunctionMapSerializer(&self.bucket_functions),
254        )?;
255        state.serialize_field("submodules", &self.submodules)?;
256        state.end()
257    }
258}
259
260impl Module {
261    /// Generates tye markdown style documentation for hte module
262    pub fn documentation(&self, level: usize) -> Result<String, std::fmt::Error> {
263        let header = "#".repeat(level + 1);
264        let mut functions = String::new();
265        let mut align_functions: Vec<_> = self.align_functions.iter().collect();
266
267        align_functions.sort_by_key(|(i, _)| *i);
268        if !align_functions.is_empty() {
269            writeln!(&mut functions, "#{header} Align Functions")?;
270        }
271        for (n, f) in &align_functions {
272            writeln!(&mut functions, "{}", f.documentation(n))?;
273            writeln!(&mut functions)?;
274        }
275        let mut mapping_functions: Vec<_> = self.mapping_functions.iter().collect();
276        mapping_functions.sort_by_key(|(i, _)| *i);
277        if !mapping_functions.is_empty() {
278            writeln!(&mut functions, "#{header} Map Functions")?;
279        }
280        for (n, f) in &mapping_functions {
281            writeln!(&mut functions, "{}", f.documentation(n))?;
282            writeln!(&mut functions)?;
283        }
284        let mut group_functions: Vec<_> = self.group_functions.iter().collect();
285        group_functions.sort_by_key(|(i, _)| *i);
286        if !group_functions.is_empty() {
287            writeln!(&mut functions, "#{header} Group Functions")?;
288        }
289        for (n, f) in &group_functions {
290            writeln!(&mut functions, "{}", f.documentation(n))?;
291            writeln!(&mut functions)?;
292        }
293        let mut compute_functions: Vec<_> = self.compute_functions.iter().collect();
294        compute_functions.sort_by_key(|(i, _)| *i);
295        if !compute_functions.is_empty() {
296            writeln!(&mut functions, "#{header} Compute Functions")?;
297        }
298        for (n, f) in &compute_functions {
299            writeln!(&mut functions, "{}", f.documentation(n))?;
300            writeln!(&mut functions)?;
301        }
302        let mut bucket_functions: Vec<_> = self.bucket_functions.iter().collect();
303        bucket_functions.sort_by_key(|(i, _)| *i);
304        if !bucket_functions.is_empty() {
305            writeln!(&mut functions, "#{header} Bucket Functions")?;
306        }
307        for (n, f) in &bucket_functions {
308            writeln!(&mut functions, "{}", f.documentation(n))?;
309            writeln!(&mut functions)?;
310        }
311
312        let mut submodule_list: Vec<_> = self.submodules.iter().collect();
313        submodule_list.sort_by_key(|(i, _)| *i);
314        let mut submodules = String::new();
315        for (_, m) in submodule_list {
316            writeln!(&mut submodules, "{}", m.documentation(level + 1)?)?;
317        }
318        Ok(format!(
319            r"{header} {name}
320{doc}
321{functions}{submodules}",
322            name = self.name,
323            doc = self.doc,
324        ))
325    }
326    pub(crate) fn map_fn(&self, id: &Function) -> Option<&MapFunction> {
327        self.map_fn_(&id.module_path, &id.name)
328    }
329    fn map_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&MapFunction> {
330        if let Some((first, rest)) = modules.split_first() {
331            self.submodules.get(first)?.map_fn_(rest, id)
332        } else {
333            self.mapping_functions.get(id)
334        }
335    }
336
337    pub(crate) fn align_fn(&self, id: &Function) -> Option<&AlignFunction> {
338        self.align_fn_(&id.module_path, &id.name)
339    }
340    fn align_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&AlignFunction> {
341        if let Some((first, rest)) = modules.split_first() {
342            self.submodules.get(first)?.align_fn_(rest, id)
343        } else {
344            self.align_functions.get(id)
345        }
346    }
347    pub(crate) fn group_fn(&self, id: &Function) -> Option<&GroupFunction> {
348        self.group_fn_(&id.module_path, &id.name)
349    }
350    fn group_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&GroupFunction> {
351        if let Some((first, rest)) = modules.split_first() {
352            self.submodules.get(first)?.group_fn_(rest, id)
353        } else {
354            self.group_functions.get(id)
355        }
356    }
357
358    pub(crate) fn compute_fn(&self, id: &Function) -> Option<&ComputeFunction> {
359        self.compute_fn_(&id.module_path, &id.name)
360    }
361    fn compute_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&ComputeFunction> {
362        if let Some((first, rest)) = modules.split_first() {
363            self.submodules.get(first)?.compute_fn_(rest, id)
364        } else {
365            self.compute_functions.get(id)
366        }
367    }
368}
369
370/// User supplied Mapping function
371pub trait MapFunctionTrait:
372    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
373{
374    /// calls the function
375    #[must_use]
376    fn call(&self, input: &str) -> String;
377    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
378    #[must_use]
379    fn box_clone(&self) -> Box<dyn MapFunctionTrait>;
380}
381
382#[derive(Debug, serde::Serialize, serde::Deserialize)]
383#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
384#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
385/// A map functio wrapper
386pub enum MapFunction {
387    /// A builtin function
388    Builtin(MapType),
389    #[serde(skip)]
390    /// A use defined function
391    UserDefined(Box<dyn MapFunctionTrait>),
392}
393
394impl FunctionTrait for MapFunction {
395    fn doc(&self) -> &str {
396        match self {
397            MapFunction::Builtin(t) => t.doc(),
398            MapFunction::UserDefined(func) => func.doc(),
399        }
400    }
401    fn args(&self) -> Vec<Arg> {
402        match self {
403            MapFunction::Builtin(t) => t.args(),
404            MapFunction::UserDefined(func) => func.args(),
405        }
406    }
407}
408
409impl From<MapType> for MapFunction {
410    fn from(t: MapType) -> Self {
411        MapFunction::Builtin(t)
412    }
413}
414impl std::fmt::Display for MapFunction {
415    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
416        match self {
417            MapFunction::Builtin(t) => write!(f, "{t}"),
418            MapFunction::UserDefined(func) => write!(f, "{func}"),
419        }
420    }
421}
422impl Clone for MapFunction {
423    fn clone(&self) -> Self {
424        match self {
425            MapFunction::Builtin(t) => MapFunction::Builtin(*t),
426            MapFunction::UserDefined(func) => MapFunction::UserDefined(func.box_clone()),
427        }
428    }
429}
430
431/// User supplied Mapping function
432pub trait AlignFunctionTrait:
433    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
434{
435    /// calls the function
436    #[must_use]
437    fn call(&self, input: &str) -> String;
438    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
439    #[must_use]
440    fn box_clone(&self) -> Box<dyn AlignFunctionTrait>;
441}
442
443#[derive(Debug, serde::Serialize, serde::Deserialize)]
444#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
445#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
446/// A align function wrapper
447pub enum AlignFunction {
448    /// A builtin function
449    Builtin(TimeType),
450    #[serde(skip)]
451    /// A use defined function
452    UserDefined(Box<dyn AlignFunctionTrait>),
453}
454impl FunctionTrait for AlignFunction {
455    fn doc(&self) -> &str {
456        match self {
457            AlignFunction::Builtin(t) => t.doc(),
458            AlignFunction::UserDefined(func) => func.doc(),
459        }
460    }
461    fn args(&self) -> Vec<Arg> {
462        match self {
463            AlignFunction::Builtin(t) => t.args(),
464            AlignFunction::UserDefined(func) => func.args(),
465        }
466    }
467}
468impl From<TimeType> for AlignFunction {
469    fn from(t: TimeType) -> Self {
470        AlignFunction::Builtin(t)
471    }
472}
473impl std::fmt::Display for AlignFunction {
474    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
475        match self {
476            AlignFunction::Builtin(t) => write!(f, "{t}"),
477            AlignFunction::UserDefined(func) => write!(f, "{func}"),
478        }
479    }
480}
481impl Clone for AlignFunction {
482    fn clone(&self) -> Self {
483        match self {
484            AlignFunction::Builtin(t) => AlignFunction::Builtin(*t),
485            AlignFunction::UserDefined(func) => AlignFunction::UserDefined(func.box_clone()),
486        }
487    }
488}
489
490/// User supplied Mapping function
491pub trait GroupFunctionTrait:
492    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
493{
494    /// calls the function
495    #[must_use]
496    fn call(&self, input: &str) -> String;
497    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
498    #[must_use]
499    fn box_clone(&self) -> Box<dyn GroupFunctionTrait>;
500}
501
502#[derive(Debug, serde::Serialize, serde::Deserialize)]
503#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
504#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
505/// A group-by function wrapper
506pub enum GroupFunction {
507    /// A builtin function
508    Builtin(TagsType),
509    #[serde(skip)]
510    /// A use defined function
511    UserDefined(Box<dyn GroupFunctionTrait>),
512}
513impl FunctionTrait for GroupFunction {
514    fn doc(&self) -> &str {
515        match self {
516            GroupFunction::Builtin(t) => t.doc(),
517            GroupFunction::UserDefined(func) => func.doc(),
518        }
519    }
520    fn args(&self) -> Vec<Arg> {
521        match self {
522            GroupFunction::Builtin(t) => t.args(),
523            GroupFunction::UserDefined(func) => func.args(),
524        }
525    }
526}
527impl From<TagsType> for GroupFunction {
528    fn from(t: TagsType) -> Self {
529        GroupFunction::Builtin(t)
530    }
531}
532impl std::fmt::Display for GroupFunction {
533    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
534        match self {
535            GroupFunction::Builtin(t) => write!(f, "{t}"),
536            GroupFunction::UserDefined(func) => write!(f, "{func}"),
537        }
538    }
539}
540impl Clone for GroupFunction {
541    fn clone(&self) -> Self {
542        match self {
543            GroupFunction::Builtin(t) => GroupFunction::Builtin(*t),
544            GroupFunction::UserDefined(func) => GroupFunction::UserDefined(func.box_clone()),
545        }
546    }
547}
548
549/// User supplied Compute function
550pub trait ComputeFunctionTrait:
551    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
552{
553    /// calls the function
554    #[must_use]
555    fn call(&self, input: &str) -> String;
556    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
557    #[must_use]
558    fn box_clone(&self) -> Box<dyn ComputeFunctionTrait>;
559}
560#[derive(Debug, serde::Serialize, serde::Deserialize)]
561#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
562#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
563/// A compute function wrapper
564pub enum ComputeFunction {
565    /// A builtin function
566    Builtin(ComputeType),
567    #[serde(skip)]
568    /// A use defined function
569    UserDefined(Box<dyn ComputeFunctionTrait>),
570}
571impl FunctionTrait for ComputeFunction {
572    fn doc(&self) -> &str {
573        match self {
574            ComputeFunction::Builtin(t) => t.doc(),
575            ComputeFunction::UserDefined(func) => func.doc(),
576        }
577    }
578    fn args(&self) -> Vec<Arg> {
579        match self {
580            ComputeFunction::Builtin(t) => t.args(),
581            ComputeFunction::UserDefined(func) => func.args(),
582        }
583    }
584}
585impl From<ComputeType> for ComputeFunction {
586    fn from(c: ComputeType) -> Self {
587        ComputeFunction::Builtin(c)
588    }
589}
590impl std::fmt::Display for ComputeFunction {
591    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
592        match self {
593            ComputeFunction::Builtin(t) => write!(f, "{t}"),
594            ComputeFunction::UserDefined(func) => write!(f, "{func}"),
595        }
596    }
597}
598impl Clone for ComputeFunction {
599    fn clone(&self) -> Self {
600        match self {
601            ComputeFunction::Builtin(c) => ComputeFunction::Builtin(*c),
602            ComputeFunction::UserDefined(func) => ComputeFunction::UserDefined(func.box_clone()),
603        }
604    }
605}