zen_expression/functions/
internal.rs

1use crate::functions::defs::{
2    CompositeFunction, FunctionDefinition, FunctionSignature, StaticFunction,
3};
4use std::rc::Rc;
5use strum_macros::{Display, EnumIter, EnumString, IntoStaticStr};
6
7#[derive(Debug, PartialEq, Eq, Hash, Display, EnumString, EnumIter, IntoStaticStr, Clone, Copy)]
8#[strum(serialize_all = "camelCase")]
9pub enum InternalFunction {
10    // General
11    Len,
12    Contains,
13    Flatten,
14
15    // String
16    Upper,
17    Lower,
18    Trim,
19    StartsWith,
20    EndsWith,
21    Matches,
22    Extract,
23    FuzzyMatch,
24    Split,
25
26    // Math
27    Abs,
28    Sum,
29    Avg,
30    Min,
31    Max,
32    Rand,
33    Median,
34    Mode,
35    Floor,
36    Ceil,
37    Round,
38
39    // Type
40    IsNumeric,
41    String,
42    Number,
43    Bool,
44    Type,
45
46    // Map
47    Keys,
48    Values,
49
50    #[strum(serialize = "d")]
51    Date,
52}
53
54impl From<&InternalFunction> for Rc<dyn FunctionDefinition> {
55    fn from(value: &InternalFunction) -> Self {
56        use crate::variable::VariableType as VT;
57        use InternalFunction as IF;
58
59        let s: Rc<dyn FunctionDefinition> = match value {
60            IF::Len => Rc::new(CompositeFunction {
61                implementation: Rc::new(imp::len),
62                signatures: vec![
63                    FunctionSignature::single(VT::String, VT::Number),
64                    FunctionSignature::single(VT::Any.array(), VT::Number),
65                ],
66            }),
67
68            IF::Contains => Rc::new(CompositeFunction {
69                implementation: Rc::new(imp::contains),
70                signatures: vec![
71                    FunctionSignature {
72                        parameters: vec![VT::String, VT::String],
73                        return_type: VT::Bool,
74                    },
75                    FunctionSignature {
76                        parameters: vec![VT::Any.array(), VT::Any],
77                        return_type: VT::Bool,
78                    },
79                ],
80            }),
81
82            IF::Flatten => Rc::new(StaticFunction {
83                implementation: Rc::new(imp::flatten),
84                signature: FunctionSignature::single(VT::Any.array(), VT::Any.array()),
85            }),
86
87            IF::Upper => Rc::new(StaticFunction {
88                implementation: Rc::new(imp::upper),
89                signature: FunctionSignature::single(VT::String, VT::String),
90            }),
91
92            IF::Lower => Rc::new(StaticFunction {
93                implementation: Rc::new(imp::lower),
94                signature: FunctionSignature::single(VT::String, VT::String),
95            }),
96
97            IF::Trim => Rc::new(StaticFunction {
98                implementation: Rc::new(imp::trim),
99                signature: FunctionSignature::single(VT::String, VT::String),
100            }),
101
102            IF::StartsWith => Rc::new(StaticFunction {
103                implementation: Rc::new(imp::starts_with),
104                signature: FunctionSignature {
105                    parameters: vec![VT::String, VT::String],
106                    return_type: VT::Bool,
107                },
108            }),
109
110            IF::EndsWith => Rc::new(StaticFunction {
111                implementation: Rc::new(imp::ends_with),
112                signature: FunctionSignature {
113                    parameters: vec![VT::String, VT::String],
114                    return_type: VT::Bool,
115                },
116            }),
117
118            IF::Matches => Rc::new(StaticFunction {
119                implementation: Rc::new(imp::matches),
120                signature: FunctionSignature {
121                    parameters: vec![VT::String, VT::String],
122                    return_type: VT::Bool,
123                },
124            }),
125
126            IF::Extract => Rc::new(StaticFunction {
127                implementation: Rc::new(imp::extract),
128                signature: FunctionSignature {
129                    parameters: vec![VT::String, VT::String],
130                    return_type: VT::String.array(),
131                },
132            }),
133
134            IF::Split => Rc::new(StaticFunction {
135                implementation: Rc::new(imp::split),
136                signature: FunctionSignature {
137                    parameters: vec![VT::String, VT::String],
138                    return_type: VT::String.array(),
139                },
140            }),
141
142            IF::FuzzyMatch => Rc::new(CompositeFunction {
143                implementation: Rc::new(imp::fuzzy_match),
144                signatures: vec![
145                    FunctionSignature {
146                        parameters: vec![VT::String, VT::String],
147                        return_type: VT::Number,
148                    },
149                    FunctionSignature {
150                        parameters: vec![VT::String.array(), VT::String],
151                        return_type: VT::Number.array(),
152                    },
153                ],
154            }),
155
156            IF::Abs => Rc::new(StaticFunction {
157                implementation: Rc::new(imp::abs),
158                signature: FunctionSignature::single(VT::Number, VT::Number),
159            }),
160
161            IF::Rand => Rc::new(StaticFunction {
162                implementation: Rc::new(imp::rand),
163                signature: FunctionSignature::single(VT::Number, VT::Number),
164            }),
165
166            IF::Floor => Rc::new(StaticFunction {
167                implementation: Rc::new(imp::floor),
168                signature: FunctionSignature::single(VT::Number, VT::Number),
169            }),
170
171            IF::Ceil => Rc::new(StaticFunction {
172                implementation: Rc::new(imp::ceil),
173                signature: FunctionSignature::single(VT::Number, VT::Number),
174            }),
175
176            IF::Round => Rc::new(StaticFunction {
177                implementation: Rc::new(imp::round),
178                signature: FunctionSignature::single(VT::Number, VT::Number),
179            }),
180
181            IF::Sum => Rc::new(StaticFunction {
182                implementation: Rc::new(imp::sum),
183                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
184            }),
185
186            IF::Avg => Rc::new(StaticFunction {
187                implementation: Rc::new(imp::avg),
188                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
189            }),
190
191            IF::Min => Rc::new(StaticFunction {
192                implementation: Rc::new(imp::min),
193                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
194            }),
195
196            IF::Max => Rc::new(StaticFunction {
197                implementation: Rc::new(imp::max),
198                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
199            }),
200
201            IF::Median => Rc::new(StaticFunction {
202                implementation: Rc::new(imp::median),
203                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
204            }),
205
206            IF::Mode => Rc::new(StaticFunction {
207                implementation: Rc::new(imp::mode),
208                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
209            }),
210
211            IF::Type => Rc::new(StaticFunction {
212                implementation: Rc::new(imp::to_type),
213                signature: FunctionSignature::single(VT::Any, VT::String),
214            }),
215
216            IF::String => Rc::new(StaticFunction {
217                implementation: Rc::new(imp::to_string),
218                signature: FunctionSignature::single(VT::Any, VT::String),
219            }),
220
221            IF::Bool => Rc::new(StaticFunction {
222                implementation: Rc::new(imp::to_bool),
223                signature: FunctionSignature::single(VT::Any, VT::Bool),
224            }),
225
226            IF::IsNumeric => Rc::new(StaticFunction {
227                implementation: Rc::new(imp::is_numeric),
228                signature: FunctionSignature::single(VT::Any, VT::Bool),
229            }),
230
231            IF::Number => Rc::new(StaticFunction {
232                implementation: Rc::new(imp::to_number),
233                signature: FunctionSignature::single(VT::Any, VT::String),
234            }),
235
236            IF::Keys => Rc::new(CompositeFunction {
237                implementation: Rc::new(imp::keys),
238                signatures: vec![
239                    FunctionSignature::single(VT::Object(Default::default()), VT::String.array()),
240                    FunctionSignature::single(VT::Any.array(), VT::Number.array()),
241                ],
242            }),
243
244            IF::Values => Rc::new(StaticFunction {
245                implementation: Rc::new(imp::values),
246                signature: FunctionSignature::single(
247                    VT::Object(Default::default()),
248                    VT::Any.array(),
249                ),
250            }),
251
252            IF::Date => Rc::new(CompositeFunction {
253                implementation: Rc::new(imp::date),
254                signatures: vec![
255                    FunctionSignature {
256                        parameters: vec![],
257                        return_type: VT::Date,
258                    },
259                    FunctionSignature {
260                        parameters: vec![VT::Any],
261                        return_type: VT::Date,
262                    },
263                    FunctionSignature {
264                        parameters: vec![VT::Any, VT::String],
265                        return_type: VT::Date,
266                    },
267                ],
268            }),
269        };
270
271        s
272    }
273}
274
275pub(crate) mod imp {
276    use crate::functions::arguments::Arguments;
277    use crate::vm::VmDate;
278    use crate::Variable as V;
279    use anyhow::{anyhow, Context};
280    use chrono_tz::Tz;
281    #[cfg(not(feature = "regex-lite"))]
282    use regex::Regex;
283    #[cfg(feature = "regex-lite")]
284    use regex_lite::Regex;
285    use rust_decimal::prelude::{FromPrimitive, ToPrimitive};
286    use rust_decimal::Decimal;
287    use rust_decimal_macros::dec;
288    use std::collections::BTreeMap;
289    use std::rc::Rc;
290    use std::str::FromStr;
291
292    fn __internal_number_array(args: &Arguments, pos: usize) -> anyhow::Result<Vec<Decimal>> {
293        let a = args.array(pos)?;
294        let arr = a.borrow();
295
296        arr.iter()
297            .map(|v| v.as_number())
298            .collect::<Option<Vec<_>>>()
299            .context("Expected a number array")
300    }
301
302    pub fn starts_with(args: Arguments) -> anyhow::Result<V> {
303        let a = args.str(0)?;
304        let b = args.str(1)?;
305
306        Ok(V::Bool(a.starts_with(b)))
307    }
308
309    pub fn ends_with(args: Arguments) -> anyhow::Result<V> {
310        let a = args.str(0)?;
311        let b = args.str(1)?;
312
313        Ok(V::Bool(a.ends_with(b)))
314    }
315
316    pub fn matches(args: Arguments) -> anyhow::Result<V> {
317        let a = args.str(0)?;
318        let b = args.str(1)?;
319
320        let regex = Regex::new(b.as_ref()).context("Invalid regular expression")?;
321
322        Ok(V::Bool(regex.is_match(a.as_ref())))
323    }
324
325    pub fn upper(args: Arguments) -> anyhow::Result<V> {
326        let a = args.str(0)?;
327        Ok(V::String(a.to_uppercase().into()))
328    }
329
330    pub fn lower(args: Arguments) -> anyhow::Result<V> {
331        let a = args.str(0)?;
332        Ok(V::String(a.to_lowercase().into()))
333    }
334
335    pub fn trim(args: Arguments) -> anyhow::Result<V> {
336        let a = args.str(0)?;
337        Ok(V::String(a.trim().into()))
338    }
339
340    pub fn extract(args: Arguments) -> anyhow::Result<V> {
341        let a = args.str(0)?;
342        let b = args.str(1)?;
343
344        let regex = Regex::new(b.as_ref()).context("Invalid regular expression")?;
345
346        let captures = regex
347            .captures(a.as_ref())
348            .map(|capture| {
349                capture
350                    .iter()
351                    .map(|c| c.map(|c| c.as_str()))
352                    .filter_map(|c| c)
353                    .map(|s| V::String(Rc::from(s)))
354                    .collect()
355            })
356            .unwrap_or_default();
357
358        Ok(V::from_array(captures))
359    }
360
361    pub fn split(args: Arguments) -> anyhow::Result<V> {
362        let a = args.str(0)?;
363        let b = args.str(1)?;
364
365        let arr = Vec::from_iter(
366            a.split(b)
367                .into_iter()
368                .map(|s| V::String(s.to_string().into())),
369        );
370
371        Ok(V::from_array(arr))
372    }
373
374    pub fn flatten(args: Arguments) -> anyhow::Result<V> {
375        let a = args.array(0)?;
376
377        let arr = a.borrow();
378        let mut flat_arr = Vec::with_capacity(arr.len());
379        arr.iter().for_each(|v| match v {
380            V::Array(b) => {
381                let arr = b.borrow();
382                arr.iter().for_each(|v| flat_arr.push(v.clone()))
383            }
384            _ => flat_arr.push(v.clone()),
385        });
386
387        Ok(V::from_array(flat_arr))
388    }
389
390    pub fn abs(args: Arguments) -> anyhow::Result<V> {
391        let a = args.number(0)?;
392        Ok(V::Number(a.abs()))
393    }
394
395    pub fn ceil(args: Arguments) -> anyhow::Result<V> {
396        let a = args.number(0)?;
397        Ok(V::Number(a.ceil()))
398    }
399
400    pub fn floor(args: Arguments) -> anyhow::Result<V> {
401        let a = args.number(0)?;
402        Ok(V::Number(a.floor()))
403    }
404
405    pub fn round(args: Arguments) -> anyhow::Result<V> {
406        let a = args.number(0)?;
407        Ok(V::Number(a.round()))
408    }
409
410    pub fn rand(args: Arguments) -> anyhow::Result<V> {
411        let a = args.number(0)?;
412        let upper_range = a.round().to_i64().context("Invalid upper range")?;
413
414        let random_number = fastrand::i64(0..=upper_range);
415        Ok(V::Number(Decimal::from(random_number)))
416    }
417
418    pub fn min(args: Arguments) -> anyhow::Result<V> {
419        let a = __internal_number_array(&args, 0)?;
420        let min = a.iter().min().context("Empty array")?;
421
422        Ok(V::Number(Decimal::from(*min)))
423    }
424
425    pub fn max(args: Arguments) -> anyhow::Result<V> {
426        let a = __internal_number_array(&args, 0)?;
427        let max = a.iter().max().context("Empty array")?;
428
429        Ok(V::Number(Decimal::from(*max)))
430    }
431
432    pub fn avg(args: Arguments) -> anyhow::Result<V> {
433        let a = __internal_number_array(&args, 0)?;
434        let sum = a.iter().fold(Decimal::ZERO, |acc, x| acc + x);
435
436        Ok(V::Number(Decimal::from(
437            sum.checked_div(Decimal::from(a.len()))
438                .context("Empty array")?,
439        )))
440    }
441
442    pub fn sum(args: Arguments) -> anyhow::Result<V> {
443        let a = __internal_number_array(&args, 0)?;
444        let sum = a.iter().fold(Decimal::ZERO, |acc, v| acc + v);
445
446        Ok(V::Number(Decimal::from(sum)))
447    }
448
449    pub fn median(args: Arguments) -> anyhow::Result<V> {
450        let mut a = __internal_number_array(&args, 0)?;
451        a.sort();
452
453        let center = a.len() / 2;
454        if a.len() % 2 == 1 {
455            let center_num = a.get(center).context("Index out of bounds")?;
456            Ok(V::Number(*center_num))
457        } else {
458            let center_left = a.get(center - 1).context("Index out of bounds")?;
459            let center_right = a.get(center).context("Index out of bounds")?;
460
461            let median = ((*center_left) + (*center_right)) / dec!(2);
462            Ok(V::Number(median))
463        }
464    }
465
466    pub fn mode(args: Arguments) -> anyhow::Result<V> {
467        let a = __internal_number_array(&args, 0)?;
468        let mut counts = BTreeMap::new();
469        for num in a {
470            *counts.entry(num).or_insert(0) += 1;
471        }
472
473        let most_common = counts
474            .into_iter()
475            .max_by_key(|&(_, count)| count)
476            .map(|(num, _)| num)
477            .context("Empty array")?;
478
479        Ok(V::Number(most_common))
480    }
481
482    pub fn to_type(args: Arguments) -> anyhow::Result<V> {
483        let a = args.var(0)?;
484        Ok(V::String(a.type_name().into()))
485    }
486
487    pub fn to_bool(args: Arguments) -> anyhow::Result<V> {
488        let a = args.var(0)?;
489        let val = match a {
490            V::Null => false,
491            V::Bool(v) => *v,
492            V::Number(n) => !n.is_zero(),
493            V::Array(_) | V::Object(_) | V::Dynamic(_) => true,
494            V::String(s) => match (*s).trim() {
495                "true" => true,
496                "false" => false,
497                _ => s.is_empty(),
498            },
499        };
500
501        Ok(V::Bool(val))
502    }
503
504    pub fn to_string(args: Arguments) -> anyhow::Result<V> {
505        let a = args.var(0)?;
506        let val = match a {
507            V::Null => Rc::from("null"),
508            V::Bool(v) => Rc::from(v.to_string().as_str()),
509            V::Number(n) => Rc::from(n.to_string().as_str()),
510            V::String(s) => s.clone(),
511            _ => return Err(anyhow!("Cannot convert type {} to string", a.type_name())),
512        };
513
514        Ok(V::String(val))
515    }
516
517    pub fn to_number(args: Arguments) -> anyhow::Result<V> {
518        let a = args.var(0)?;
519        let val = match a {
520            V::Number(n) => *n,
521            V::String(str) => Decimal::from_str_exact(str.trim()).context("Invalid number")?,
522            V::Bool(b) => match *b {
523                true => Decimal::ONE,
524                false => Decimal::ZERO,
525            },
526            _ => return Err(anyhow!("Cannot convert type {} to number", a.type_name())),
527        };
528
529        Ok(V::Number(val))
530    }
531
532    pub fn is_numeric(args: Arguments) -> anyhow::Result<V> {
533        let a = args.var(0)?;
534        let is_ok = match a {
535            V::Number(_) => true,
536            V::String(str) => Decimal::from_str_exact(str.trim()).is_ok(),
537            _ => false,
538        };
539
540        Ok(V::Bool(is_ok))
541    }
542
543    pub fn len(args: Arguments) -> anyhow::Result<V> {
544        let a = args.var(0)?;
545        let len = match a {
546            V::String(s) => s.len(),
547            V::Array(s) => {
548                let arr = s.borrow();
549                arr.len()
550            }
551            _ => {
552                return Err(anyhow!("Cannot determine len of type {}", a.type_name()));
553            }
554        };
555
556        Ok(V::Number(len.into()))
557    }
558
559    pub fn contains(args: Arguments) -> anyhow::Result<V> {
560        let a = args.var(0)?;
561        let b = args.var(1)?;
562
563        let val = match (a, b) {
564            (V::String(a), V::String(b)) => a.contains(b.as_ref()),
565            (V::Array(a), _) => {
566                let arr = a.borrow();
567
568                arr.iter().any(|a| match (a, b) {
569                    (V::Number(a), V::Number(b)) => a == b,
570                    (V::String(a), V::String(b)) => a == b,
571                    (V::Bool(a), V::Bool(b)) => a == b,
572                    (V::Null, V::Null) => true,
573                    _ => false,
574                })
575            }
576            _ => {
577                return Err(anyhow!(
578                    "Cannot determine contains for type {} and {}",
579                    a.type_name(),
580                    b.type_name()
581                ));
582            }
583        };
584
585        Ok(V::Bool(val))
586    }
587
588    pub fn fuzzy_match(args: Arguments) -> anyhow::Result<V> {
589        let a = args.var(0)?;
590        let b = args.str(1)?;
591
592        let val = match a {
593            V::String(a) => {
594                let sim = strsim::normalized_damerau_levenshtein(a.as_ref(), b.as_ref());
595                // This is okay, as NDL will return [0, 1]
596                V::Number(Decimal::from_f64(sim).unwrap_or(dec!(0)))
597            }
598            V::Array(_a) => {
599                let a = _a.borrow();
600                let mut sims = Vec::with_capacity(a.len());
601                for v in a.iter() {
602                    let s = v.as_str().context("Expected string array")?;
603
604                    let sim = Decimal::from_f64(strsim::normalized_damerau_levenshtein(
605                        s.as_ref(),
606                        b.as_ref(),
607                    ))
608                    .unwrap_or(dec!(0));
609                    sims.push(V::Number(sim));
610                }
611
612                V::from_array(sims)
613            }
614            _ => return Err(anyhow!("Fuzzy match not available for type")),
615        };
616
617        Ok(val)
618    }
619
620    pub fn keys(args: Arguments) -> anyhow::Result<V> {
621        let a = args.var(0)?;
622        let var = match a {
623            V::Array(a) => {
624                let arr = a.borrow();
625                let indices = arr
626                    .iter()
627                    .enumerate()
628                    .map(|(index, _)| V::Number(index.into()))
629                    .collect();
630
631                V::from_array(indices)
632            }
633            V::Object(a) => {
634                let obj = a.borrow();
635                let keys = obj.iter().map(|(key, _)| V::String(key.clone())).collect();
636
637                V::from_array(keys)
638            }
639            _ => {
640                return Err(anyhow!("Cannot determine keys of type {}", a.type_name()));
641            }
642        };
643
644        Ok(var)
645    }
646
647    pub fn values(args: Arguments) -> anyhow::Result<V> {
648        let a = args.object(0)?;
649        let obj = a.borrow();
650        let values: Vec<_> = obj.values().cloned().collect();
651
652        Ok(V::from_array(values))
653    }
654
655    pub fn date(args: Arguments) -> anyhow::Result<V> {
656        let provided = args.ovar(0);
657        let tz = args
658            .ostr(1)?
659            .map(|v| Tz::from_str(v).context("Invalid timezone"))
660            .transpose()?;
661
662        let date_time = match provided {
663            Some(v) => VmDate::new(v.clone(), tz),
664            None => VmDate::now(),
665        };
666
667        Ok(V::Dynamic(Rc::new(date_time)))
668    }
669}