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

zen_expression/functions/
internal.rs
