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 Len,
12 Contains,
13 Flatten,
14
15 Upper,
17 Lower,
18 Trim,
19 StartsWith,
20 EndsWith,
21 Matches,
22 Extract,
23 FuzzyMatch,
24 Split,
25
26 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 IsNumeric,
42 String,
43 Number,
44 Bool,
45 Type,
46
47 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::date::DynamicVariableExt;
308 use crate::vm::VmDate;
309 use crate::{Variable as V, Variable};
310 use anyhow::{anyhow, Context};
311 use chrono_tz::Tz;
312 #[cfg(not(feature = "regex-lite"))]
313 use regex::Regex;
314 #[cfg(feature = "regex-lite")]
315 use regex_lite::Regex;
316 use rust_decimal::prelude::{FromPrimitive, ToPrimitive};
317 use rust_decimal::{Decimal, RoundingStrategy};
318 use rust_decimal_macros::dec;
319 use std::collections::BTreeMap;
320 use std::rc::Rc;
321 use std::str::FromStr;
322
323 fn __internal_number_array(args: &Arguments, pos: usize) -> anyhow::Result<Vec<Decimal>> {
324 let a = args.array(pos)?;
325 let arr = a.borrow();
326
327 arr.iter()
328 .map(|v| v.as_number())
329 .collect::<Option<Vec<_>>>()
330 .context("Expected a number array")
331 }
332
333 enum Either<A, B> {
334 Left(A),
335 Right(B),
336 }
337
338 fn __internal_number_or_date_array(
339 args: &Arguments,
340 pos: usize,
341 ) -> anyhow::Result<Either<Vec<Decimal>, Vec<VmDate>>> {
342 let a = args.array(pos)?;
343 let arr = a.borrow();
344
345 let is_number = arr.first().map(|v| v.as_number()).flatten().is_some();
346 if is_number {
347 Ok(Either::Left(
348 arr.iter()
349 .map(|v| v.as_number())
350 .collect::<Option<Vec<_>>>()
351 .context("Expected a number array")?,
352 ))
353 } else {
354 Ok(Either::Right(
355 arr.iter()
356 .map(|v| match v {
357 Variable::Dynamic(d) => d.as_date().cloned(),
358 _ => None,
359 })
360 .collect::<Option<Vec<_>>>()
361 .context("Expected a number array")?,
362 ))
363 }
364 }
365
366 pub fn starts_with(args: Arguments) -> anyhow::Result<V> {
367 let a = args.str(0)?;
368 let b = args.str(1)?;
369
370 Ok(V::Bool(a.starts_with(b)))
371 }
372
373 pub fn ends_with(args: Arguments) -> anyhow::Result<V> {
374 let a = args.str(0)?;
375 let b = args.str(1)?;
376
377 Ok(V::Bool(a.ends_with(b)))
378 }
379
380 pub fn matches(args: Arguments) -> anyhow::Result<V> {
381 let a = args.str(0)?;
382 let b = args.str(1)?;
383
384 let regex = Regex::new(b.as_ref()).context("Invalid regular expression")?;
385
386 Ok(V::Bool(regex.is_match(a.as_ref())))
387 }
388
389 pub fn upper(args: Arguments) -> anyhow::Result<V> {
390 let a = args.str(0)?;
391 Ok(V::String(a.to_uppercase().into()))
392 }
393
394 pub fn lower(args: Arguments) -> anyhow::Result<V> {
395 let a = args.str(0)?;
396 Ok(V::String(a.to_lowercase().into()))
397 }
398
399 pub fn trim(args: Arguments) -> anyhow::Result<V> {
400 let a = args.str(0)?;
401 Ok(V::String(a.trim().into()))
402 }
403
404 pub fn extract(args: Arguments) -> anyhow::Result<V> {
405 let a = args.str(0)?;
406 let b = args.str(1)?;
407
408 let regex = Regex::new(b.as_ref()).context("Invalid regular expression")?;
409
410 let captures = regex
411 .captures(a.as_ref())
412 .map(|capture| {
413 capture
414 .iter()
415 .map(|c| c.map(|c| c.as_str()))
416 .filter_map(|c| c)
417 .map(|s| V::String(Rc::from(s)))
418 .collect()
419 })
420 .unwrap_or_default();
421
422 Ok(V::from_array(captures))
423 }
424
425 pub fn split(args: Arguments) -> anyhow::Result<V> {
426 let a = args.str(0)?;
427 let b = args.str(1)?;
428
429 let arr = Vec::from_iter(
430 a.split(b)
431 .into_iter()
432 .map(|s| V::String(s.to_string().into())),
433 );
434
435 Ok(V::from_array(arr))
436 }
437
438 pub fn flatten(args: Arguments) -> anyhow::Result<V> {
439 let a = args.array(0)?;
440
441 let arr = a.borrow();
442 let mut flat_arr = Vec::with_capacity(arr.len());
443 arr.iter().for_each(|v| match v {
444 V::Array(b) => {
445 let arr = b.borrow();
446 arr.iter().for_each(|v| flat_arr.push(v.clone()))
447 }
448 _ => flat_arr.push(v.clone()),
449 });
450
451 Ok(V::from_array(flat_arr))
452 }
453
454 pub fn abs(args: Arguments) -> anyhow::Result<V> {
455 let a = args.number(0)?;
456 Ok(V::Number(a.abs()))
457 }
458
459 pub fn ceil(args: Arguments) -> anyhow::Result<V> {
460 let a = args.number(0)?;
461 Ok(V::Number(a.ceil()))
462 }
463
464 pub fn floor(args: Arguments) -> anyhow::Result<V> {
465 let a = args.number(0)?;
466 Ok(V::Number(a.floor()))
467 }
468
469 pub fn round(args: Arguments) -> anyhow::Result<V> {
470 let a = args.number(0)?;
471 let dp = args
472 .onumber(1)?
473 .map(|v| v.to_u32().context("Invalid number of decimal places"))
474 .transpose()?
475 .unwrap_or(0);
476
477 Ok(V::Number(a.round_dp_with_strategy(
478 dp,
479 RoundingStrategy::MidpointAwayFromZero,
480 )))
481 }
482
483 pub fn trunc(args: Arguments) -> anyhow::Result<V> {
484 let a = args.number(0)?;
485 let dp = args
486 .onumber(1)?
487 .map(|v| v.to_u32().context("Invalid number of decimal places"))
488 .transpose()?
489 .unwrap_or(0);
490
491 Ok(V::Number(a.trunc_with_scale(dp)))
492 }
493
494 pub fn rand(args: Arguments) -> anyhow::Result<V> {
495 let a = args.number(0)?;
496 let upper_range = a.round().to_i64().context("Invalid upper range")?;
497
498 let random_number = fastrand::i64(0..=upper_range);
499 Ok(V::Number(Decimal::from(random_number)))
500 }
501
502 pub fn min(args: Arguments) -> anyhow::Result<V> {
503 let a = __internal_number_or_date_array(&args, 0)?;
504
505 match a {
506 Either::Left(arr) => {
507 let max = arr.into_iter().min().context("Empty array")?;
508 Ok(V::Number(Decimal::from(max)))
509 }
510 Either::Right(arr) => {
511 let max = arr.into_iter().min().context("Empty array")?;
512 Ok(V::Dynamic(Rc::new(max)))
513 }
514 }
515 }
516
517 pub fn max(args: Arguments) -> anyhow::Result<V> {
518 let a = __internal_number_or_date_array(&args, 0)?;
519
520 match a {
521 Either::Left(arr) => {
522 let max = arr.into_iter().max().context("Empty array")?;
523 Ok(V::Number(Decimal::from(max)))
524 }
525 Either::Right(arr) => {
526 let max = arr.into_iter().max().context("Empty array")?;
527 Ok(V::Dynamic(Rc::new(max)))
528 }
529 }
530 }
531
532 pub fn avg(args: Arguments) -> anyhow::Result<V> {
533 let a = __internal_number_array(&args, 0)?;
534 let sum = a.iter().fold(Decimal::ZERO, |acc, x| acc + x);
535
536 Ok(V::Number(Decimal::from(
537 sum.checked_div(Decimal::from(a.len()))
538 .context("Empty array")?,
539 )))
540 }
541
542 pub fn sum(args: Arguments) -> anyhow::Result<V> {
543 let a = __internal_number_array(&args, 0)?;
544 let sum = a.iter().fold(Decimal::ZERO, |acc, v| acc + v);
545
546 Ok(V::Number(Decimal::from(sum)))
547 }
548
549 pub fn median(args: Arguments) -> anyhow::Result<V> {
550 let mut a = __internal_number_array(&args, 0)?;
551 a.sort();
552
553 let center = a.len() / 2;
554 if a.len() % 2 == 1 {
555 let center_num = a.get(center).context("Index out of bounds")?;
556 Ok(V::Number(*center_num))
557 } else {
558 let center_left = a.get(center - 1).context("Index out of bounds")?;
559 let center_right = a.get(center).context("Index out of bounds")?;
560
561 let median = ((*center_left) + (*center_right)) / dec!(2);
562 Ok(V::Number(median))
563 }
564 }
565
566 pub fn mode(args: Arguments) -> anyhow::Result<V> {
567 let a = __internal_number_array(&args, 0)?;
568 let mut counts = BTreeMap::new();
569 for num in a {
570 *counts.entry(num).or_insert(0) += 1;
571 }
572
573 let most_common = counts
574 .into_iter()
575 .max_by_key(|&(_, count)| count)
576 .map(|(num, _)| num)
577 .context("Empty array")?;
578
579 Ok(V::Number(most_common))
580 }
581
582 pub fn to_type(args: Arguments) -> anyhow::Result<V> {
583 let a = args.var(0)?;
584 Ok(V::String(a.type_name().into()))
585 }
586
587 pub fn to_bool(args: Arguments) -> anyhow::Result<V> {
588 let a = args.var(0)?;
589 let val = match a {
590 V::Null => false,
591 V::Bool(v) => *v,
592 V::Number(n) => !n.is_zero(),
593 V::Array(_) | V::Object(_) | V::Dynamic(_) => true,
594 V::String(s) => match (*s).trim() {
595 "true" => true,
596 "false" => false,
597 _ => s.is_empty(),
598 },
599 };
600
601 Ok(V::Bool(val))
602 }
603
604 pub fn to_string(args: Arguments) -> anyhow::Result<V> {
605 let a = args.var(0)?;
606 let val = match a {
607 V::Null => Rc::from("null"),
608 V::Bool(v) => Rc::from(v.to_string().as_str()),
609 V::Number(n) => Rc::from(n.to_string().as_str()),
610 V::String(s) => s.clone(),
611 _ => return Err(anyhow!("Cannot convert type {} to string", a.type_name())),
612 };
613
614 Ok(V::String(val))
615 }
616
617 pub fn to_number(args: Arguments) -> anyhow::Result<V> {
618 let a = args.var(0)?;
619 let val = match a {
620 V::Number(n) => *n,
621 V::String(str) => {
622 let s = str.trim();
623 Decimal::from_str_exact(s)
624 .or_else(|_| Decimal::from_scientific(s))
625 .context("Invalid number")?
626 }
627 V::Bool(b) => match *b {
628 true => Decimal::ONE,
629 false => Decimal::ZERO,
630 },
631 _ => return Err(anyhow!("Cannot convert type {} to number", a.type_name())),
632 };
633
634 Ok(V::Number(val))
635 }
636
637 pub fn is_numeric(args: Arguments) -> anyhow::Result<V> {
638 let a = args.var(0)?;
639 let is_ok = match a {
640 V::Number(_) => true,
641 V::String(str) => {
642 let s = str.trim();
643 Decimal::from_str_exact(s)
644 .or_else(|_| Decimal::from_scientific(s))
645 .is_ok()
646 }
647 _ => false,
648 };
649
650 Ok(V::Bool(is_ok))
651 }
652
653 pub fn len(args: Arguments) -> anyhow::Result<V> {
654 let a = args.var(0)?;
655 let len = match a {
656 V::String(s) => s.len(),
657 V::Array(s) => {
658 let arr = s.borrow();
659 arr.len()
660 }
661 _ => {
662 return Err(anyhow!("Cannot determine len of type {}", a.type_name()));
663 }
664 };
665
666 Ok(V::Number(len.into()))
667 }
668
669 pub fn contains(args: Arguments) -> anyhow::Result<V> {
670 let a = args.var(0)?;
671 let b = args.var(1)?;
672
673 let val = match (a, b) {
674 (V::String(a), V::String(b)) => a.contains(b.as_ref()),
675 (V::Array(a), _) => {
676 let arr = a.borrow();
677
678 arr.iter().any(|a| match (a, b) {
679 (V::Number(a), V::Number(b)) => a == b,
680 (V::String(a), V::String(b)) => a == b,
681 (V::Bool(a), V::Bool(b)) => a == b,
682 (V::Null, V::Null) => true,
683 _ => false,
684 })
685 }
686 _ => {
687 return Err(anyhow!(
688 "Cannot determine contains for type {} and {}",
689 a.type_name(),
690 b.type_name()
691 ));
692 }
693 };
694
695 Ok(V::Bool(val))
696 }
697
698 pub fn fuzzy_match(args: Arguments) -> anyhow::Result<V> {
699 let a = args.var(0)?;
700 let b = args.str(1)?;
701
702 let val = match a {
703 V::String(a) => {
704 let sim = strsim::normalized_damerau_levenshtein(a.as_ref(), b.as_ref());
705 V::Number(Decimal::from_f64(sim).unwrap_or(dec!(0)))
707 }
708 V::Array(_a) => {
709 let a = _a.borrow();
710 let mut sims = Vec::with_capacity(a.len());
711 for v in a.iter() {
712 let s = v.as_str().context("Expected string array")?;
713
714 let sim = Decimal::from_f64(strsim::normalized_damerau_levenshtein(
715 s.as_ref(),
716 b.as_ref(),
717 ))
718 .unwrap_or(dec!(0));
719 sims.push(V::Number(sim));
720 }
721
722 V::from_array(sims)
723 }
724 _ => return Err(anyhow!("Fuzzy match not available for type")),
725 };
726
727 Ok(val)
728 }
729
730 pub fn keys(args: Arguments) -> anyhow::Result<V> {
731 let a = args.var(0)?;
732 let var = match a {
733 V::Array(a) => {
734 let arr = a.borrow();
735 let indices = arr
736 .iter()
737 .enumerate()
738 .map(|(index, _)| V::Number(index.into()))
739 .collect();
740
741 V::from_array(indices)
742 }
743 V::Object(a) => {
744 let obj = a.borrow();
745 let keys = obj.iter().map(|(key, _)| V::String(key.clone())).collect();
746
747 V::from_array(keys)
748 }
749 _ => {
750 return Err(anyhow!("Cannot determine keys of type {}", a.type_name()));
751 }
752 };
753
754 Ok(var)
755 }
756
757 pub fn values(args: Arguments) -> anyhow::Result<V> {
758 let a = args.object(0)?;
759 let obj = a.borrow();
760 let values: Vec<_> = obj.values().cloned().collect();
761
762 Ok(V::from_array(values))
763 }
764
765 pub fn date(args: Arguments) -> anyhow::Result<V> {
766 let provided = args.ovar(0);
767 let tz = args
768 .ostr(1)?
769 .map(|v| Tz::from_str(v).context("Invalid timezone"))
770 .transpose()?;
771
772 let date_time = match provided {
773 Some(v) => VmDate::new(v.clone(), tz),
774 None => VmDate::now(),
775 };
776
777 Ok(V::Dynamic(Rc::new(date_time)))
778 }
779}