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    /// Return the function's bare name (no module qualification).
150    #[must_use]
151    pub fn name(&self) -> &str {
152        &self.0
153    }
154
155    pub(crate) fn new(name: &str) -> Self {
156        FunctionId(name.to_string())
157    }
158}
159
160impl Borrow<str> for FunctionId {
161    fn borrow(&self) -> &str {
162        &self.0
163    }
164}
165
166pub(crate) struct Function {
167    pub(crate) module_path: Vec<ModuleId>,
168    pub(crate) name: FunctionId,
169}
170
171impl std::fmt::Display for Function {
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        for m in &self.module_path {
174            write!(f, "{m}::")?;
175        }
176        write!(f, "{}", self.name)
177    }
178}
179
180/// Module definition
181pub struct Module {
182    pub(crate) name: ModuleId,
183    pub(crate) doc: &'static str,
184    pub(crate) align_functions: HashMap<FunctionId, AlignFunction>,
185    pub(crate) mapping_functions: HashMap<FunctionId, MapFunction>,
186    pub(crate) group_functions: HashMap<FunctionId, GroupFunction>,
187    pub(crate) bucket_functions: HashMap<FunctionId, BucketType>,
188    pub(crate) compute_functions: HashMap<FunctionId, ComputeFunction>,
189    pub(crate) submodules: HashMap<ModuleId, Module>,
190}
191
192impl std::fmt::Display for Module {
193    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
194        write!(f, "Module: {}", self.name)
195    }
196}
197
198struct FunctionSerializer<'a, F: FunctionTrait>(&'a F);
199
200impl<F> Serialize for FunctionSerializer<'_, F>
201where
202    F: FunctionTrait,
203{
204    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
205    where
206        S: serde::Serializer,
207    {
208        let mut state = serializer.serialize_struct("Function", 2)?;
209        state.serialize_field("doc", self.0.doc())?;
210        state.serialize_field("args", &self.0.args())?;
211        state.end()
212    }
213}
214
215struct FunctionMapSerializer<'a, F: FunctionTrait>(&'a HashMap<FunctionId, F>);
216
217impl<F> Serialize for FunctionMapSerializer<'_, F>
218where
219    F: FunctionTrait,
220{
221    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
222    where
223        S: serde::Serializer,
224    {
225        let mut state = serializer.serialize_map(Some(self.0.len()))?;
226        for (id, func) in self.0 {
227            state.serialize_entry(&id.to_string(), &FunctionSerializer(func))?;
228        }
229        state.end()
230    }
231}
232
233impl Serialize for Module {
234    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
235    where
236        S: serde::Serializer,
237    {
238        let mut state = serializer.serialize_struct("Module", 8)?;
239        state.serialize_field("name", &self.name)?;
240        state.serialize_field("doc", &self.doc)?;
241        state.serialize_field(
242            "align_functions",
243            &FunctionMapSerializer(&self.align_functions),
244        )?;
245        state.serialize_field(
246            "mapping_functions",
247            &FunctionMapSerializer(&self.mapping_functions),
248        )?;
249        state.serialize_field(
250            "group_functions",
251            &FunctionMapSerializer(&self.group_functions),
252        )?;
253        state.serialize_field(
254            "compute_functions",
255            &FunctionMapSerializer(&self.compute_functions),
256        )?;
257        state.serialize_field(
258            "bucket_functions",
259            &FunctionMapSerializer(&self.bucket_functions),
260        )?;
261        state.serialize_field("submodules", &self.submodules)?;
262        state.end()
263    }
264}
265
266impl Module {
267    /// Iterate over this module's submodules, yielding `(qualified_name, module)`
268    /// pairs. Used by language-server crates that walk the stdlib tree to
269    /// build completion lists and function lookups.
270    pub fn submodule_iter(&self) -> impl Iterator<Item = (&str, &Module)> {
271        self.submodules.iter().map(|(k, v)| (k.0.as_str(), v))
272    }
273
274    /// Resolve a submodule by name (single segment, not `::`-qualified).
275    #[must_use]
276    pub fn submodule(&self, name: &str) -> Option<&Module> {
277        self.submodules.get(name)
278    }
279
280    /// Iterate this module's align functions.
281    pub fn align_function_iter(&self) -> impl Iterator<Item = (&str, &AlignFunction)> {
282        self.align_functions.iter().map(|(k, v)| (k.0.as_str(), v))
283    }
284
285    /// Iterate this module's map functions.
286    pub fn mapping_function_iter(&self) -> impl Iterator<Item = (&str, &MapFunction)> {
287        self.mapping_functions
288            .iter()
289            .map(|(k, v)| (k.0.as_str(), v))
290    }
291
292    /// Iterate this module's group functions.
293    pub fn group_function_iter(&self) -> impl Iterator<Item = (&str, &GroupFunction)> {
294        self.group_functions.iter().map(|(k, v)| (k.0.as_str(), v))
295    }
296
297    /// Iterate this module's bucket types.
298    pub fn bucket_function_iter(&self) -> impl Iterator<Item = (&str, &BucketType)> {
299        self.bucket_functions.iter().map(|(k, v)| (k.0.as_str(), v))
300    }
301
302    /// Iterate this module's compute functions.
303    pub fn compute_function_iter(&self) -> impl Iterator<Item = (&str, &ComputeFunction)> {
304        self.compute_functions
305            .iter()
306            .map(|(k, v)| (k.0.as_str(), v))
307    }
308
309    /// Look up an align function by bare name.
310    #[must_use]
311    pub fn align_function(&self, name: &str) -> Option<&AlignFunction> {
312        self.align_functions.get(name)
313    }
314
315    /// Look up a map function by bare name.
316    #[must_use]
317    pub fn mapping_function(&self, name: &str) -> Option<&MapFunction> {
318        self.mapping_functions.get(name)
319    }
320
321    /// Look up a group function by bare name.
322    #[must_use]
323    pub fn group_function(&self, name: &str) -> Option<&GroupFunction> {
324        self.group_functions.get(name)
325    }
326
327    /// Look up a bucket type by bare name.
328    #[must_use]
329    pub fn bucket_function(&self, name: &str) -> Option<&BucketType> {
330        self.bucket_functions.get(name)
331    }
332
333    /// Look up a compute function by bare name.
334    #[must_use]
335    pub fn compute_function(&self, name: &str) -> Option<&ComputeFunction> {
336        self.compute_functions.get(name)
337    }
338
339    /// Generates the markdown style documentation for the module
340    pub fn documentation(&self, level: usize) -> Result<String, std::fmt::Error> {
341        let header = "#".repeat(level + 1);
342        let mut functions = String::new();
343        let mut align_functions: Vec<_> = self.align_functions.iter().collect();
344
345        align_functions.sort_by_key(|(i, _)| *i);
346        if !align_functions.is_empty() {
347            writeln!(&mut functions, "#{header} Align Functions")?;
348        }
349        for (n, f) in &align_functions {
350            writeln!(&mut functions, "{}", f.documentation(n))?;
351            writeln!(&mut functions)?;
352        }
353        let mut mapping_functions: Vec<_> = self.mapping_functions.iter().collect();
354        mapping_functions.sort_by_key(|(i, _)| *i);
355        if !mapping_functions.is_empty() {
356            writeln!(&mut functions, "#{header} Map Functions")?;
357        }
358        for (n, f) in &mapping_functions {
359            writeln!(&mut functions, "{}", f.documentation(n))?;
360            writeln!(&mut functions)?;
361        }
362        let mut group_functions: Vec<_> = self.group_functions.iter().collect();
363        group_functions.sort_by_key(|(i, _)| *i);
364        if !group_functions.is_empty() {
365            writeln!(&mut functions, "#{header} Group Functions")?;
366        }
367        for (n, f) in &group_functions {
368            writeln!(&mut functions, "{}", f.documentation(n))?;
369            writeln!(&mut functions)?;
370        }
371        let mut compute_functions: Vec<_> = self.compute_functions.iter().collect();
372        compute_functions.sort_by_key(|(i, _)| *i);
373        if !compute_functions.is_empty() {
374            writeln!(&mut functions, "#{header} Compute Functions")?;
375        }
376        for (n, f) in &compute_functions {
377            writeln!(&mut functions, "{}", f.documentation(n))?;
378            writeln!(&mut functions)?;
379        }
380        let mut bucket_functions: Vec<_> = self.bucket_functions.iter().collect();
381        bucket_functions.sort_by_key(|(i, _)| *i);
382        if !bucket_functions.is_empty() {
383            writeln!(&mut functions, "#{header} Bucket Functions")?;
384        }
385        for (n, f) in &bucket_functions {
386            writeln!(&mut functions, "{}", f.documentation(n))?;
387            writeln!(&mut functions)?;
388        }
389
390        let mut submodule_list: Vec<_> = self.submodules.iter().collect();
391        submodule_list.sort_by_key(|(i, _)| *i);
392        let mut submodules = String::new();
393        for (_, m) in submodule_list {
394            writeln!(&mut submodules, "{}", m.documentation(level + 1)?)?;
395        }
396        Ok(format!(
397            r"{header} {name}
398{doc}
399{functions}{submodules}",
400            name = self.name,
401            doc = self.doc,
402        ))
403    }
404    pub(crate) fn map_fn(&self, id: &Function) -> Option<&MapFunction> {
405        self.map_fn_(&id.module_path, &id.name)
406    }
407    fn map_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&MapFunction> {
408        if let Some((first, rest)) = modules.split_first() {
409            self.submodules.get(first)?.map_fn_(rest, id)
410        } else {
411            self.mapping_functions.get(id)
412        }
413    }
414
415    pub(crate) fn align_fn(&self, id: &Function) -> Option<&AlignFunction> {
416        self.align_fn_(&id.module_path, &id.name)
417    }
418    fn align_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&AlignFunction> {
419        if let Some((first, rest)) = modules.split_first() {
420            self.submodules.get(first)?.align_fn_(rest, id)
421        } else {
422            self.align_functions.get(id)
423        }
424    }
425    pub(crate) fn group_fn(&self, id: &Function) -> Option<&GroupFunction> {
426        self.group_fn_(&id.module_path, &id.name)
427    }
428    fn group_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&GroupFunction> {
429        if let Some((first, rest)) = modules.split_first() {
430            self.submodules.get(first)?.group_fn_(rest, id)
431        } else {
432            self.group_functions.get(id)
433        }
434    }
435
436    pub(crate) fn compute_fn(&self, id: &Function) -> Option<&ComputeFunction> {
437        self.compute_fn_(&id.module_path, &id.name)
438    }
439    fn compute_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&ComputeFunction> {
440        if let Some((first, rest)) = modules.split_first() {
441            self.submodules.get(first)?.compute_fn_(rest, id)
442        } else {
443            self.compute_functions.get(id)
444        }
445    }
446}
447
448/// User supplied Mapping function
449pub trait MapFunctionTrait:
450    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
451{
452    /// calls the function
453    #[must_use]
454    fn call(&self, input: &str) -> String;
455    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
456    #[must_use]
457    fn box_clone(&self) -> Box<dyn MapFunctionTrait>;
458}
459
460#[derive(Debug, serde::Serialize, serde::Deserialize)]
461/// A map functio wrapper
462pub enum MapFunction {
463    /// A builtin function
464    Builtin(MapType),
465    #[serde(skip)]
466    /// A use defined function
467    UserDefined(Box<dyn MapFunctionTrait>),
468}
469
470impl FunctionTrait for MapFunction {
471    fn doc(&self) -> &str {
472        match self {
473            MapFunction::Builtin(t) => t.doc(),
474            MapFunction::UserDefined(func) => func.doc(),
475        }
476    }
477    fn args(&self) -> Vec<Arg> {
478        match self {
479            MapFunction::Builtin(t) => t.args(),
480            MapFunction::UserDefined(func) => func.args(),
481        }
482    }
483}
484
485impl From<MapType> for MapFunction {
486    fn from(t: MapType) -> Self {
487        MapFunction::Builtin(t)
488    }
489}
490impl std::fmt::Display for MapFunction {
491    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
492        match self {
493            MapFunction::Builtin(t) => write!(f, "{t}"),
494            MapFunction::UserDefined(func) => write!(f, "{func}"),
495        }
496    }
497}
498impl Clone for MapFunction {
499    fn clone(&self) -> Self {
500        match self {
501            MapFunction::Builtin(t) => MapFunction::Builtin(*t),
502            MapFunction::UserDefined(func) => MapFunction::UserDefined(func.box_clone()),
503        }
504    }
505}
506
507/// User supplied Mapping function
508pub trait AlignFunctionTrait:
509    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
510{
511    /// calls the function
512    #[must_use]
513    fn call(&self, input: &str) -> String;
514    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
515    #[must_use]
516    fn box_clone(&self) -> Box<dyn AlignFunctionTrait>;
517}
518
519#[derive(Debug, serde::Serialize, serde::Deserialize)]
520/// A align function wrapper
521pub enum AlignFunction {
522    /// A builtin function
523    Builtin(TimeType),
524    #[serde(skip)]
525    /// A use defined function
526    UserDefined(Box<dyn AlignFunctionTrait>),
527}
528impl FunctionTrait for AlignFunction {
529    fn doc(&self) -> &str {
530        match self {
531            AlignFunction::Builtin(t) => t.doc(),
532            AlignFunction::UserDefined(func) => func.doc(),
533        }
534    }
535    fn args(&self) -> Vec<Arg> {
536        match self {
537            AlignFunction::Builtin(t) => t.args(),
538            AlignFunction::UserDefined(func) => func.args(),
539        }
540    }
541}
542impl From<TimeType> for AlignFunction {
543    fn from(t: TimeType) -> Self {
544        AlignFunction::Builtin(t)
545    }
546}
547impl std::fmt::Display for AlignFunction {
548    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
549        match self {
550            AlignFunction::Builtin(t) => write!(f, "{t}"),
551            AlignFunction::UserDefined(func) => write!(f, "{func}"),
552        }
553    }
554}
555impl Clone for AlignFunction {
556    fn clone(&self) -> Self {
557        match self {
558            AlignFunction::Builtin(t) => AlignFunction::Builtin(*t),
559            AlignFunction::UserDefined(func) => AlignFunction::UserDefined(func.box_clone()),
560        }
561    }
562}
563
564/// User supplied Mapping function
565pub trait GroupFunctionTrait:
566    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
567{
568    /// calls the function
569    #[must_use]
570    fn call(&self, input: &str) -> String;
571    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
572    #[must_use]
573    fn box_clone(&self) -> Box<dyn GroupFunctionTrait>;
574}
575
576#[derive(Debug, serde::Serialize, serde::Deserialize)]
577/// A group-by function wrapper
578pub enum GroupFunction {
579    /// A builtin function
580    Builtin(TagsType),
581    #[serde(skip)]
582    /// A use defined function
583    UserDefined(Box<dyn GroupFunctionTrait>),
584}
585impl FunctionTrait for GroupFunction {
586    fn doc(&self) -> &str {
587        match self {
588            GroupFunction::Builtin(t) => t.doc(),
589            GroupFunction::UserDefined(func) => func.doc(),
590        }
591    }
592    fn args(&self) -> Vec<Arg> {
593        match self {
594            GroupFunction::Builtin(t) => t.args(),
595            GroupFunction::UserDefined(func) => func.args(),
596        }
597    }
598}
599impl From<TagsType> for GroupFunction {
600    fn from(t: TagsType) -> Self {
601        GroupFunction::Builtin(t)
602    }
603}
604impl std::fmt::Display for GroupFunction {
605    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
606        match self {
607            GroupFunction::Builtin(t) => write!(f, "{t}"),
608            GroupFunction::UserDefined(func) => write!(f, "{func}"),
609        }
610    }
611}
612impl Clone for GroupFunction {
613    fn clone(&self) -> Self {
614        match self {
615            GroupFunction::Builtin(t) => GroupFunction::Builtin(*t),
616            GroupFunction::UserDefined(func) => GroupFunction::UserDefined(func.box_clone()),
617        }
618    }
619}
620
621/// User supplied Compute function
622pub trait ComputeFunctionTrait:
623    Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
624{
625    /// calls the function
626    #[must_use]
627    fn call(&self, input: &str) -> String;
628    /// Creates a boxed clone of the function object (clone doesn't work with boxed traites)
629    #[must_use]
630    fn box_clone(&self) -> Box<dyn ComputeFunctionTrait>;
631}
632#[derive(Debug, serde::Serialize, serde::Deserialize)]
633/// A compute function wrapper
634pub enum ComputeFunction {
635    /// A builtin function
636    Builtin(ComputeType),
637    #[serde(skip)]
638    /// A use defined function
639    UserDefined(Box<dyn ComputeFunctionTrait>),
640}
641impl FunctionTrait for ComputeFunction {
642    fn doc(&self) -> &str {
643        match self {
644            ComputeFunction::Builtin(t) => t.doc(),
645            ComputeFunction::UserDefined(func) => func.doc(),
646        }
647    }
648    fn args(&self) -> Vec<Arg> {
649        match self {
650            ComputeFunction::Builtin(t) => t.args(),
651            ComputeFunction::UserDefined(func) => func.args(),
652        }
653    }
654}
655impl From<ComputeType> for ComputeFunction {
656    fn from(c: ComputeType) -> Self {
657        ComputeFunction::Builtin(c)
658    }
659}
660impl std::fmt::Display for ComputeFunction {
661    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
662        match self {
663            ComputeFunction::Builtin(t) => write!(f, "{t}"),
664            ComputeFunction::UserDefined(func) => write!(f, "{func}"),
665        }
666    }
667}
668impl Clone for ComputeFunction {
669    fn clone(&self) -> Self {
670        match self {
671            ComputeFunction::Builtin(c) => ComputeFunction::Builtin(*c),
672            ComputeFunction::UserDefined(func) => ComputeFunction::UserDefined(func.box_clone()),
673        }
674    }
675}