Skip to main content

akar_function/scalar/
arithmetic.rs

1use super::{gamma_func, log_gamma, rng_next, set_rng_seed};
2use crate::registry::*;
3use akar_common::types::Value;
4
5// ==================== Arithmetic ====================
6
7fn gcd_impl(mut a: i64, mut b: i64) -> i64 {
8    while b != 0 {
9        let t = b;
10        b = a % b;
11        a = t;
12    }
13    a.abs()
14}
15
16fn lcm_impl(a: i64, b: i64) -> i64 {
17    if a == 0 || b == 0 {
18        0
19    } else {
20        (a / gcd_impl(a, b)).abs() * b.abs()
21    }
22}
23
24pub(crate) fn evaluate_arithmetic(op: ArithmeticOp, args: &[Value]) -> Result<Value, String> {
25    // Allow empty args for ops that take no arguments (Pi, Rand)
26    let needs_args = !matches!(op, ArithmeticOp::Pi | ArithmeticOp::Rand);
27    if args.is_empty() && needs_args {
28        return Err("Arithmetic requires at least one argument".into());
29    }
30
31    match op {
32        ArithmeticOp::Negate => {
33            let v = args[0].clone();
34            match v {
35                Value::Int64(x) => Ok(Value::Int64(-x)),
36                Value::Double(x) => Ok(Value::Double(-x)),
37                _ => Err("Cannot negate non-numeric".into()),
38            }
39        }
40        ArithmeticOp::Abs => {
41            let v = args[0].clone();
42            match v {
43                Value::Int64(x) => Ok(Value::Int64(x.checked_abs().unwrap_or(i64::MAX))),
44                Value::UInt64(x) => Ok(Value::UInt64(x)),
45                Value::Double(x) => Ok(Value::Double(x.abs())),
46                _ => Err("Abs requires numeric".into()),
47            }
48        }
49        ArithmeticOp::Ceil => match args[0] {
50            Value::Double(x) => Ok(Value::Double(x.ceil())),
51            Value::Int64(x) => Ok(Value::Double((x as f64).ceil())),
52            _ => Err("Ceil requires numeric".into()),
53        },
54        ArithmeticOp::Floor => match args[0] {
55            Value::Double(x) => Ok(Value::Double(x.floor())),
56            Value::Int64(x) => Ok(Value::Double((x as f64).floor())),
57            _ => Err("Floor requires numeric".into()),
58        },
59        ArithmeticOp::Round => match args[0] {
60            Value::Double(x) => Ok(Value::Double(x.round())),
61            _ => Err("Round requires numeric".into()),
62        },
63        ArithmeticOp::Sqrt => {
64            let v = numeric_to_f64(&args[0])?;
65            Ok(Value::Double(v.sqrt()))
66        }
67        ArithmeticOp::Log => {
68            let v = numeric_to_f64(&args[0])?;
69            Ok(Value::Double(v.ln()))
70        }
71        ArithmeticOp::Exp => {
72            let v = numeric_to_f64(&args[0])?;
73            Ok(Value::Double(v.exp()))
74        }
75        ArithmeticOp::Sin => {
76            let v = numeric_to_f64(&args[0])?;
77            Ok(Value::Double(v.sin()))
78        }
79        ArithmeticOp::Cos => {
80            let v = numeric_to_f64(&args[0])?;
81            Ok(Value::Double(v.cos()))
82        }
83        ArithmeticOp::Tan => {
84            let v = numeric_to_f64(&args[0])?;
85            Ok(Value::Double(v.tan()))
86        }
87        ArithmeticOp::Sinh => {
88            let v = numeric_to_f64(&args[0])?;
89            Ok(Value::Double(v.sinh()))
90        }
91        ArithmeticOp::Cosh => {
92            let v = numeric_to_f64(&args[0])?;
93            Ok(Value::Double(v.cosh()))
94        }
95        ArithmeticOp::Tanh => {
96            let v = numeric_to_f64(&args[0])?;
97            Ok(Value::Double(v.tanh()))
98        }
99        ArithmeticOp::Asin => {
100            let v = numeric_to_f64(&args[0])?;
101            Ok(Value::Double(v.asin()))
102        }
103        ArithmeticOp::Acos => {
104            let v = numeric_to_f64(&args[0])?;
105            Ok(Value::Double(v.acos()))
106        }
107        ArithmeticOp::Atan => {
108            let v = numeric_to_f64(&args[0])?;
109            Ok(Value::Double(v.atan()))
110        }
111        ArithmeticOp::Atan2 => {
112            if args.len() < 2 {
113                return Err("Atan2 requires 2 arguments".into());
114            }
115            let y = numeric_to_f64(&args[0])?;
116            let x = numeric_to_f64(&args[1])?;
117            Ok(Value::Double(y.atan2(x)))
118        }
119        ArithmeticOp::Degrees => {
120            let v = numeric_to_f64(&args[0])?;
121            Ok(Value::Double(v.to_degrees()))
122        }
123        ArithmeticOp::Radians => {
124            let v = numeric_to_f64(&args[0])?;
125            Ok(Value::Double(v.to_radians()))
126        }
127        ArithmeticOp::Sign => {
128            let v = numeric_to_f64(&args[0])?;
129            Ok(Value::Int64(if v > 0.0 {
130                1
131            } else if v < 0.0 {
132                -1
133            } else {
134                0
135            }))
136        }
137        ArithmeticOp::Pi => Ok(Value::Double(std::f64::consts::PI)),
138        ArithmeticOp::Rand => Ok(Value::Double(rng_next())),
139        ArithmeticOp::Power => {
140            if args.len() < 2 {
141                return Err("Power requires 2 arguments".into());
142            }
143            let base = numeric_to_f64(&args[0])?;
144            let exp = numeric_to_f64(&args[1])?;
145            Ok(Value::Double(base.powf(exp)))
146        }
147        // Math functions (f64-based, single argument)
148        ArithmeticOp::Cbrt => {
149            let v = numeric_to_f64(&args[0])?;
150            Ok(Value::Double(v.cbrt()))
151        }
152        ArithmeticOp::Cot => {
153            let v = numeric_to_f64(&args[0])?;
154            Ok(Value::Double(1.0 / v.tan()))
155        }
156        ArithmeticOp::Log2 => {
157            let v = numeric_to_f64(&args[0])?;
158            Ok(Value::Double(v.log2()))
159        }
160        ArithmeticOp::Even => {
161            let v = args[0].clone();
162            match v {
163                Value::Int64(x) => {
164                    // Round up to nearest even integer
165                    let rounded = if x % 2 == 0 { x } else { x + 1 };
166                    Ok(Value::Int64(rounded))
167                }
168                Value::Double(x) => {
169                    let rounded = x.ceil() as i64;
170                    let result = if rounded % 2 == 0 { rounded } else { rounded + 1 };
171                    Ok(Value::Int64(result))
172                }
173                _ => Err("Even requires numeric argument".into()),
174            }
175        }
176        // Heavy math functions (C++ port)
177        ArithmeticOp::Factorial => {
178            let n = match &args[0] {
179                Value::Int64(x) if *x >= 0 => *x,
180                Value::Int64(_) => return Err("Factorial requires non-negative integer".into()),
181                _ => return Err("Factorial requires integer argument".into()),
182            };
183            let mut result: i64 = 1;
184            for i in 2..=n {
185                result = result.wrapping_mul(i);
186            }
187            Ok(Value::Int64(result))
188        }
189        ArithmeticOp::Gamma => {
190            let v = numeric_to_f64(&args[0])?;
191            // Poles at non-positive integers
192            if v <= 0.0 && (v - v.round()).abs() < 1e-12 {
193                return Ok(Value::Double(f64::INFINITY));
194            }
195            Ok(Value::Double(gamma_func(v)))
196        }
197        ArithmeticOp::Lgamma => {
198            let v = numeric_to_f64(&args[0])?;
199            // Poles at non-positive integers: return infinity
200            if v <= 0.0 && (v - v.round()).abs() < 1e-12 {
201                return Ok(Value::Double(f64::INFINITY));
202            }
203            Ok(Value::Double(log_gamma(v)))
204        }
205        ArithmeticOp::SetSeed => {
206            let v = match &args[0] {
207                Value::Double(x) => x,
208                Value::Int64(x) => &(*x as f64),
209                _ => return Err("SetSeed requires numeric argument".into()),
210            };
211            let seed = (v * (u64::MAX as f64)) as u64;
212            set_rng_seed(seed);
213            // Return INT32(0) to match C++ semantics
214            Ok(Value::Int32(0))
215        }
216        // Bitwise operations (int64-only, matching C++ hardcoded int64_t)
217        ArithmeticOp::BitwiseAnd => {
218            if args.len() < 2 {
219                return Err("Bitwise AND requires 2 arguments".into());
220            }
221            match (&args[0], &args[1]) {
222                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x & y)),
223                _ => Err("Bitwise AND requires integer arguments".into()),
224            }
225        }
226        ArithmeticOp::BitwiseOr => {
227            if args.len() < 2 {
228                return Err("Bitwise OR requires 2 arguments".into());
229            }
230            match (&args[0], &args[1]) {
231                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x | y)),
232                _ => Err("Bitwise OR requires integer arguments".into()),
233            }
234        }
235        ArithmeticOp::BitwiseXor => {
236            if args.len() < 2 {
237                return Err("Bitwise XOR requires 2 arguments".into());
238            }
239            match (&args[0], &args[1]) {
240                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x ^ y)),
241                _ => Err("Bitwise XOR requires integer arguments".into()),
242            }
243        }
244        ArithmeticOp::BitShiftLeft => {
245            if args.len() < 2 {
246                return Err("Bit shift left requires 2 arguments".into());
247            }
248            match (&args[0], &args[1]) {
249                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x << y)),
250                _ => Err("Bit shift left requires integer arguments".into()),
251            }
252        }
253        ArithmeticOp::BitShiftRight => {
254            if args.len() < 2 {
255                return Err("Bit shift right requires 2 arguments".into());
256            }
257            match (&args[0], &args[1]) {
258                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x >> y)),
259                _ => Err("Bit shift right requires integer arguments".into()),
260            }
261        }
262        ArithmeticOp::Gcd => {
263            if args.len() < 2 {
264                return Err("Gcd requires 2 arguments".into());
265            }
266            match (&args[0], &args[1]) {
267                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(gcd_impl(*x, *y))),
268                _ => Err("Gcd requires integer arguments".into()),
269            }
270        }
271        ArithmeticOp::Lcm => {
272            if args.len() < 2 {
273                return Err("Lcm requires 2 arguments".into());
274            }
275            match (&args[0], &args[1]) {
276                (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(lcm_impl(*x, *y))),
277                _ => Err("Lcm requires integer arguments".into()),
278            }
279        }
280        // Binary arithmetic ops
281        _ => {
282            if args.len() < 2 {
283                return Err("Binary arithmetic requires 2 arguments".into());
284            }
285            let a = args[0].clone();
286            let b = args[1].clone();
287            match op {
288                ArithmeticOp::Add => add_values(a, b),
289                ArithmeticOp::Sub => sub_values(a, b),
290                ArithmeticOp::Mul => mul_values(a, b),
291                ArithmeticOp::Div => div_values(a, b),
292                ArithmeticOp::Mod => mod_values(a, b),
293                _ => Err(format!("Unimplemented arithmetic op: {:?}", op)),
294            }
295        }
296    }
297}
298
299pub(crate) fn numeric_to_f64(v: &Value) -> Result<f64, String> {
300    match v {
301        Value::Int64(x) => Ok(*x as f64),
302        Value::Double(x) => Ok(*x),
303        Value::Float(x) => Ok(*x as f64),
304        Value::Int32(x) => Ok(*x as f64),
305        Value::UInt64(x) => Ok(*x as f64),
306        Value::UInt32(x) => Ok(*x as f64),
307        Value::UInt16(x) => Ok(*x as f64),
308        Value::UInt8(x) => Ok(*x as f64),
309        _ => Err("Expected numeric value".into()),
310    }
311}
312
313fn add_values(a: Value, b: Value) -> Result<Value, String> {
314    if matches!(a, Value::Float(_)) || matches!(b, Value::Float(_)) {
315        return Ok(Value::Double(numeric_to_f64(&a)? + numeric_to_f64(&b)?));
316    }
317    match (&a, &b) {
318        (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x + y)),
319        (Value::Double(x), Value::Double(y)) => Ok(Value::Double(x + y)),
320        (Value::Int64(x), Value::Double(y)) => Ok(Value::Double(*x as f64 + y)),
321        (Value::Double(x), Value::Int64(y)) => Ok(Value::Double(x + *y as f64)),
322        (Value::String(x), Value::String(y)) => Ok(Value::String(format!("{}{}", x, y))),
323        (Value::UInt64(x), Value::UInt64(y)) => x
324            .checked_add(*y)
325            .map(Value::UInt64)
326            .ok_or_else(|| "UInt64 addition overflow".into()),
327        (Value::UInt64(x), Value::Int64(y)) if *y >= 0 => x
328            .checked_add(*y as u64)
329            .map(Value::UInt64)
330            .ok_or_else(|| "UInt64 addition overflow".into()),
331        (Value::Int64(x), Value::UInt64(y)) if *x >= 0 => (*x as u64)
332            .checked_add(*y)
333            .map(Value::UInt64)
334            .ok_or_else(|| "UInt64 addition overflow".into()),
335        _ => Err(format!("Cannot add {:?} and {:?}", a.logical_type(), b.logical_type())),
336    }
337}
338
339fn sub_values(a: Value, b: Value) -> Result<Value, String> {
340    if matches!(a, Value::Float(_)) || matches!(b, Value::Float(_)) {
341        return Ok(Value::Double(numeric_to_f64(&a)? - numeric_to_f64(&b)?));
342    }
343    match (&a, &b) {
344        (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x - y)),
345        (Value::Double(x), Value::Double(y)) => Ok(Value::Double(x - y)),
346        (Value::UInt64(x), Value::UInt64(y)) => x
347            .checked_sub(*y)
348            .map(Value::UInt64)
349            .ok_or_else(|| "UInt64 subtraction underflow".into()),
350        (Value::UInt64(x), Value::Int64(y)) if *y >= 0 => x
351            .checked_sub(*y as u64)
352            .map(Value::UInt64)
353            .ok_or_else(|| "UInt64 subtraction underflow".into()),
354        (Value::UInt64(x), Value::Int64(y)) if *y < 0 => x
355            .checked_add(y.unsigned_abs())
356            .map(Value::UInt64)
357            .ok_or_else(|| "UInt64 subtraction overflow".into()),
358        (Value::Int64(x), Value::UInt64(y)) if *x >= 0 => (*x as u64)
359            .checked_sub(*y)
360            .map(Value::UInt64)
361            .ok_or_else(|| "UInt64 subtraction underflow".into()),
362        _ => Err("Cannot subtract non-numeric".into()),
363    }
364}
365
366fn mul_values(a: Value, b: Value) -> Result<Value, String> {
367    if matches!(a, Value::Float(_)) || matches!(b, Value::Float(_)) {
368        return Ok(Value::Double(numeric_to_f64(&a)? * numeric_to_f64(&b)?));
369    }
370    match (&a, &b) {
371        (Value::Int64(x), Value::Int64(y)) => Ok(Value::Int64(x * y)),
372        (Value::Double(x), Value::Double(y)) => Ok(Value::Double(x * y)),
373        (Value::UInt64(x), Value::UInt64(y)) => x
374            .checked_mul(*y)
375            .map(Value::UInt64)
376            .ok_or_else(|| "UInt64 multiplication overflow".into()),
377        (Value::UInt64(x), Value::Int64(y)) if *y >= 0 => x
378            .checked_mul(*y as u64)
379            .map(Value::UInt64)
380            .ok_or_else(|| "UInt64 multiplication overflow".into()),
381        (Value::Int64(x), Value::UInt64(y)) if *x >= 0 => (*x as u64)
382            .checked_mul(*y)
383            .map(Value::UInt64)
384            .ok_or_else(|| "UInt64 multiplication overflow".into()),
385        _ => Err("Cannot multiply non-numeric".into()),
386    }
387}
388
389fn div_values(a: Value, b: Value) -> Result<Value, String> {
390    if matches!(a, Value::Float(_)) || matches!(b, Value::Float(_)) {
391        let y = numeric_to_f64(&b)?;
392        if y == 0.0 {
393            return Err("Division by zero".into());
394        }
395        return Ok(Value::Double(numeric_to_f64(&a)? / y));
396    }
397    match (&a, &b) {
398        (Value::Int64(x), Value::Int64(y)) => {
399            if *y == 0 {
400                return Err("Division by zero".into());
401            }
402            Ok(Value::Int64(x / y))
403        }
404        (Value::Double(x), Value::Double(y)) => {
405            if *y == 0.0 {
406                return Err("Division by zero".into());
407            }
408            Ok(Value::Double(x / y))
409        }
410        (Value::UInt64(x), Value::UInt64(y)) => {
411            if *y == 0 {
412                return Err("Division by zero".into());
413            }
414            Ok(Value::UInt64(x / y))
415        }
416        (Value::UInt64(x), Value::Int64(y)) if *y > 0 => Ok(Value::UInt64(x / *y as u64)),
417        (Value::Int64(x), Value::UInt64(y)) if *x >= 0 && *y > 0 => Ok(Value::UInt64(*x as u64 / y)),
418        _ => Err("Cannot divide non-numeric".into()),
419    }
420}
421
422fn mod_values(a: Value, b: Value) -> Result<Value, String> {
423    if matches!(a, Value::Float(_)) || matches!(b, Value::Float(_)) {
424        let y = numeric_to_f64(&b)?;
425        if y == 0.0 {
426            return Err("Modulo by zero".into());
427        }
428        return Ok(Value::Double(numeric_to_f64(&a)? % y));
429    }
430    match (&a, &b) {
431        (Value::Int64(x), Value::Int64(y)) => {
432            if *y == 0 {
433                return Err("Modulo by zero".into());
434            }
435            Ok(Value::Int64(x % y))
436        }
437        (Value::UInt64(x), Value::UInt64(y)) => {
438            if *y == 0 {
439                return Err("Modulo by zero".into());
440            }
441            Ok(Value::UInt64(x % y))
442        }
443        (Value::UInt64(x), Value::Int64(y)) if *y > 0 => Ok(Value::UInt64(x % *y as u64)),
444        (Value::Int64(x), Value::UInt64(y)) if *x >= 0 && *y > 0 => Ok(Value::UInt64(*x as u64 % y)),
445        _ => Err("Cannot modulo non-integer".into()),
446    }
447}