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(StaticFunction {
216                implementation: Rc::new(imp::min),
217                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
218            }),
219
220            IF::Max => Rc::new(StaticFunction {
221                implementation: Rc::new(imp::max),
222                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
223            }),
224
225            IF::Median => Rc::new(StaticFunction {
226                implementation: Rc::new(imp::median),
227                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
228            }),
229
230            IF::Mode => Rc::new(StaticFunction {
231                implementation: Rc::new(imp::mode),
232                signature: FunctionSignature::single(VT::Number.array(), VT::Number),
233            }),
234
235            IF::Type => Rc::new(StaticFunction {
236                implementation: Rc::new(imp::to_type),
237                signature: FunctionSignature::single(VT::Any, VT::String),
238            }),
239
240            IF::String => Rc::new(StaticFunction {
241                implementation: Rc::new(imp::to_string),
242                signature: FunctionSignature::single(VT::Any, VT::String),
243            }),
244
245            IF::Bool => Rc::new(StaticFunction {
246                implementation: Rc::new(imp::to_bool),
247                signature: FunctionSignature::single(VT::Any, VT::Bool),
248            }),
249
250            IF::IsNumeric => Rc::new(StaticFunction {
251                implementation: Rc::new(imp::is_numeric),
252                signature: FunctionSignature::single(VT::Any, VT::Bool),
253            }),
254
255            IF::Number => Rc::new(StaticFunction {
256                implementation: Rc::new(imp::to_number),
257                signature: FunctionSignature::single(VT::Any, VT::String),
258            }),
259
260            IF::Keys => Rc::new(CompositeFunction {
261                implementation: Rc::new(imp::keys),
262                signatures: vec![
263                    FunctionSignature::single(VT::Object(Default::default()), VT::String.array()),
264                    FunctionSignature::single(VT::Any.array(), VT::Number.array()),
265                ],
266            }),
267
268            IF::Values => Rc::new(StaticFunction {
269                implementation: Rc::new(imp::values),
270                signature: FunctionSignature::single(
271                    VT::Object(Default::default()),
272                    VT::Any.array(),
273                ),
274            }),
275
276            IF::Date => Rc::new(CompositeFunction {
277                implementation: Rc::new(imp::date),
278                signatures: vec![
279                    FunctionSignature {
280                        parameters: vec![],
281                        return_type: VT::Date,
282                    },
283                    FunctionSignature {
284                        parameters: vec![VT::Any],
285                        return_type: VT::Date,
286                    },
287                    FunctionSignature {
288                        parameters: vec![VT::Any, VT::String],
289                        return_type: VT::Date,
290                    },
291                ],
292            }),
293        };
294
295        s
296    }
297}
298
299pub(crate) mod imp {
300    use crate::functions::arguments::Arguments;
301    use crate::vm::VmDate;
302    use crate::Variable as V;
303    use anyhow::{anyhow, Context};
304    use chrono_tz::Tz;
305    #[cfg(not(feature = "regex-lite"))]
306    use regex::Regex;
307    #[cfg(feature = "regex-lite")]
308    use regex_lite::Regex;
309    use rust_decimal::prelude::{FromPrimitive, ToPrimitive};
310    use rust_decimal::Decimal;
311    use rust_decimal_macros::dec;
312    use std::collections::BTreeMap;
313    use std::rc::Rc;
314    use std::str::FromStr;
315
316    fn __internal_number_array(args: &Arguments, pos: usize) -> anyhow::Result<Vec<Decimal>> {
317        let a = args.array(pos)?;
318        let arr = a.borrow();
319
320        arr.iter()
321            .map(|v| v.as_number())
322            .collect::<Option<Vec<_>>>()
323            .context("Expected a number array")
324    }
325
326    pub fn starts_with(args: Arguments) -> anyhow::Result<V> {
327        let a = args.str(0)?;
328        let b = args.str(1)?;
329
330        Ok(V::Bool(a.starts_with(b)))
331    }
332
333    pub fn ends_with(args: Arguments) -> anyhow::Result<V> {
334        let a = args.str(0)?;
335        let b = args.str(1)?;
336
337        Ok(V::Bool(a.ends_with(b)))
338    }
339
340    pub fn matches(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        Ok(V::Bool(regex.is_match(a.as_ref())))
347    }
348
349    pub fn upper(args: Arguments) -> anyhow::Result<V> {
350        let a = args.str(0)?;
351        Ok(V::String(a.to_uppercase().into()))
352    }
353
354    pub fn lower(args: Arguments) -> anyhow::Result<V> {
355        let a = args.str(0)?;
356        Ok(V::String(a.to_lowercase().into()))
357    }
358
359    pub fn trim(args: Arguments) -> anyhow::Result<V> {
360        let a = args.str(0)?;
361        Ok(V::String(a.trim().into()))
362    }
363
364    pub fn extract(args: Arguments) -> anyhow::Result<V> {
365        let a = args.str(0)?;
366        let b = args.str(1)?;
367
368        let regex = Regex::new(b.as_ref()).context("Invalid regular expression")?;
369
370        let captures = regex
371            .captures(a.as_ref())
372            .map(|capture| {
373                capture
374                    .iter()
375                    .map(|c| c.map(|c| c.as_str()))
376                    .filter_map(|c| c)
377                    .map(|s| V::String(Rc::from(s)))
378                    .collect()
379            })
380            .unwrap_or_default();
381
382        Ok(V::from_array(captures))
383    }
384
385    pub fn split(args: Arguments) -> anyhow::Result<V> {
386        let a = args.str(0)?;
387        let b = args.str(1)?;
388
389        let arr = Vec::from_iter(
390            a.split(b)
391                .into_iter()
392                .map(|s| V::String(s.to_string().into())),
393        );
394
395        Ok(V::from_array(arr))
396    }
397
398    pub fn flatten(args: Arguments) -> anyhow::Result<V> {
399        let a = args.array(0)?;
400
401        let arr = a.borrow();
402        let mut flat_arr = Vec::with_capacity(arr.len());
403        arr.iter().for_each(|v| match v {
404            V::Array(b) => {
405                let arr = b.borrow();
406                arr.iter().for_each(|v| flat_arr.push(v.clone()))
407            }
408            _ => flat_arr.push(v.clone()),
409        });
410
411        Ok(V::from_array(flat_arr))
412    }
413
414    pub fn abs(args: Arguments) -> anyhow::Result<V> {
415        let a = args.number(0)?;
416        Ok(V::Number(a.abs()))
417    }
418
419    pub fn ceil(args: Arguments) -> anyhow::Result<V> {
420        let a = args.number(0)?;
421        Ok(V::Number(a.ceil()))
422    }
423
424    pub fn floor(args: Arguments) -> anyhow::Result<V> {
425        let a = args.number(0)?;
426        Ok(V::Number(a.floor()))
427    }
428
429    pub fn round(args: Arguments) -> anyhow::Result<V> {
430        let a = args.number(0)?;
431        let dp = args
432            .onumber(1)?
433            .map(|v| v.to_u32().context("Invalid number of decimal places"))
434            .transpose()?
435            .unwrap_or(0);
436
437        Ok(V::Number(a.round_dp(dp)))
438    }
439
440    pub fn trunc(args: Arguments) -> anyhow::Result<V> {
441        let a = args.number(0)?;
442        let dp = args
443            .onumber(1)?
444            .map(|v| v.to_u32().context("Invalid number of decimal places"))
445            .transpose()?
446            .unwrap_or(0);
447
448        Ok(V::Number(a.trunc_with_scale(dp)))
449    }
450
451    pub fn rand(args: Arguments) -> anyhow::Result<V> {
452        let a = args.number(0)?;
453        let upper_range = a.round().to_i64().context("Invalid upper range")?;
454
455        let random_number = fastrand::i64(0..=upper_range);
456        Ok(V::Number(Decimal::from(random_number)))
457    }
458
459    pub fn min(args: Arguments) -> anyhow::Result<V> {
460        let a = __internal_number_array(&args, 0)?;
461        let min = a.iter().min().context("Empty array")?;
462
463        Ok(V::Number(Decimal::from(*min)))
464    }
465
466    pub fn max(args: Arguments) -> anyhow::Result<V> {
467        let a = __internal_number_array(&args, 0)?;
468        let max = a.iter().max().context("Empty array")?;
469
470        Ok(V::Number(Decimal::from(*max)))
471    }
472
473    pub fn avg(args: Arguments) -> anyhow::Result<V> {
474        let a = __internal_number_array(&args, 0)?;
475        let sum = a.iter().fold(Decimal::ZERO, |acc, x| acc + x);
476
477        Ok(V::Number(Decimal::from(
478            sum.checked_div(Decimal::from(a.len()))
479                .context("Empty array")?,
480        )))
481    }
482
483    pub fn sum(args: Arguments) -> anyhow::Result<V> {
484        let a = __internal_number_array(&args, 0)?;
485        let sum = a.iter().fold(Decimal::ZERO, |acc, v| acc + v);
486
487        Ok(V::Number(Decimal::from(sum)))
488    }
489
490    pub fn median(args: Arguments) -> anyhow::Result<V> {
491        let mut a = __internal_number_array(&args, 0)?;
492        a.sort();
493
494        let center = a.len() / 2;
495        if a.len() % 2 == 1 {
496            let center_num = a.get(center).context("Index out of bounds")?;
497            Ok(V::Number(*center_num))
498        } else {
499            let center_left = a.get(center - 1).context("Index out of bounds")?;
500            let center_right = a.get(center).context("Index out of bounds")?;
501
502            let median = ((*center_left) + (*center_right)) / dec!(2);
503            Ok(V::Number(median))
504        }
505    }
506
507    pub fn mode(args: Arguments) -> anyhow::Result<V> {
508        let a = __internal_number_array(&args, 0)?;
509        let mut counts = BTreeMap::new();
510        for num in a {
511            *counts.entry(num).or_insert(0) += 1;
512        }
513
514        let most_common = counts
515            .into_iter()
516            .max_by_key(|&(_, count)| count)
517            .map(|(num, _)| num)
518            .context("Empty array")?;
519
520        Ok(V::Number(most_common))
521    }
522
523    pub fn to_type(args: Arguments) -> anyhow::Result<V> {
524        let a = args.var(0)?;
525        Ok(V::String(a.type_name().into()))
526    }
527
528    pub fn to_bool(args: Arguments) -> anyhow::Result<V> {
529        let a = args.var(0)?;
530        let val = match a {
531            V::Null => false,
532            V::Bool(v) => *v,
533            V::Number(n) => !n.is_zero(),
534            V::Array(_) | V::Object(_) | V::Dynamic(_) => true,
535            V::String(s) => match (*s).trim() {
536                "true" => true,
537                "false" => false,
538                _ => s.is_empty(),
539            },
540        };
541
542        Ok(V::Bool(val))
543    }
544
545    pub fn to_string(args: Arguments) -> anyhow::Result<V> {
546        let a = args.var(0)?;
547        let val = match a {
548            V::Null => Rc::from("null"),
549            V::Bool(v) => Rc::from(v.to_string().as_str()),
550            V::Number(n) => Rc::from(n.to_string().as_str()),
551            V::String(s) => s.clone(),
552            _ => return Err(anyhow!("Cannot convert type {} to string", a.type_name())),
553        };
554
555        Ok(V::String(val))
556    }
557
558    pub fn to_number(args: Arguments) -> anyhow::Result<V> {
559        let a = args.var(0)?;
560        let val = match a {
561            V::Number(n) => *n,
562            V::String(str) => Decimal::from_str_exact(str.trim()).context("Invalid number")?,
563            V::Bool(b) => match *b {
564                true => Decimal::ONE,
565                false => Decimal::ZERO,
566            },
567            _ => return Err(anyhow!("Cannot convert type {} to number", a.type_name())),
568        };
569
570        Ok(V::Number(val))
571    }
572
573    pub fn is_numeric(args: Arguments) -> anyhow::Result<V> {
574        let a = args.var(0)?;
575        let is_ok = match a {
576            V::Number(_) => true,
577            V::String(str) => Decimal::from_str_exact(str.trim()).is_ok(),
578            _ => false,
579        };
580
581        Ok(V::Bool(is_ok))
582    }
583
584    pub fn len(args: Arguments) -> anyhow::Result<V> {
585        let a = args.var(0)?;
586        let len = match a {
587            V::String(s) => s.len(),
588            V::Array(s) => {
589                let arr = s.borrow();
590                arr.len()
591            }
592            _ => {
593                return Err(anyhow!("Cannot determine len of type {}", a.type_name()));
594            }
595        };
596
597        Ok(V::Number(len.into()))
598    }
599
600    pub fn contains(args: Arguments) -> anyhow::Result<V> {
601        let a = args.var(0)?;
602        let b = args.var(1)?;
603
604        let val = match (a, b) {
605            (V::String(a), V::String(b)) => a.contains(b.as_ref()),
606            (V::Array(a), _) => {
607                let arr = a.borrow();
608
609                arr.iter().any(|a| match (a, b) {
610                    (V::Number(a), V::Number(b)) => a == b,
611                    (V::String(a), V::String(b)) => a == b,
612                    (V::Bool(a), V::Bool(b)) => a == b,
613                    (V::Null, V::Null) => true,
614                    _ => false,
615                })
616            }
617            _ => {
618                return Err(anyhow!(
619                    "Cannot determine contains for type {} and {}",
620                    a.type_name(),
621                    b.type_name()
622                ));
623            }
624        };
625
626        Ok(V::Bool(val))
627    }
628
629    pub fn fuzzy_match(args: Arguments) -> anyhow::Result<V> {
630        let a = args.var(0)?;
631        let b = args.str(1)?;
632
633        let val = match a {
634            V::String(a) => {
635                let sim = strsim::normalized_damerau_levenshtein(a.as_ref(), b.as_ref());
636                // This is okay, as NDL will return [0, 1]
637                V::Number(Decimal::from_f64(sim).unwrap_or(dec!(0)))
638            }
639            V::Array(_a) => {
640                let a = _a.borrow();
641                let mut sims = Vec::with_capacity(a.len());
642                for v in a.iter() {
643                    let s = v.as_str().context("Expected string array")?;
644
645                    let sim = Decimal::from_f64(strsim::normalized_damerau_levenshtein(
646                        s.as_ref(),
647                        b.as_ref(),
648                    ))
649                    .unwrap_or(dec!(0));
650                    sims.push(V::Number(sim));
651                }
652
653                V::from_array(sims)
654            }
655            _ => return Err(anyhow!("Fuzzy match not available for type")),
656        };
657
658        Ok(val)
659    }
660
661    pub fn keys(args: Arguments) -> anyhow::Result<V> {
662        let a = args.var(0)?;
663        let var = match a {
664            V::Array(a) => {
665                let arr = a.borrow();
666                let indices = arr
667                    .iter()
668                    .enumerate()
669                    .map(|(index, _)| V::Number(index.into()))
670                    .collect();
671
672                V::from_array(indices)
673            }
674            V::Object(a) => {
675                let obj = a.borrow();
676                let keys = obj.iter().map(|(key, _)| V::String(key.clone())).collect();
677
678                V::from_array(keys)
679            }
680            _ => {
681                return Err(anyhow!("Cannot determine keys of type {}", a.type_name()));
682            }
683        };
684
685        Ok(var)
686    }
687
688    pub fn values(args: Arguments) -> anyhow::Result<V> {
689        let a = args.object(0)?;
690        let obj = a.borrow();
691        let values: Vec<_> = obj.values().cloned().collect();
692
693        Ok(V::from_array(values))
694    }
695
696    pub fn date(args: Arguments) -> anyhow::Result<V> {
697        let provided = args.ovar(0);
698        let tz = args
699            .ostr(1)?
700            .map(|v| Tz::from_str(v).context("Invalid timezone"))
701            .transpose()?;
702
703        let date_time = match provided {
704            Some(v) => VmDate::new(v.clone(), tz),
705            None => VmDate::now(),
706        };
707
708        Ok(V::Dynamic(Rc::new(date_time)))
709    }
710}