zen_expression/functions/
internal.rs

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