Skip to main content

mrubyedge_math/
lib.rs

1use std::rc::Rc;
2
3use mrubyedge::{
4    Error,
5    yamrb::{
6        helpers::mrb_define_singleton_cmethod,
7        value::{RObject, RValue},
8        vm::VM,
9    },
10};
11
12pub fn init_math(vm: &mut VM) {
13    if vm.get_const_by_name("Math").is_some() {
14        return;
15    }
16
17    let math_module = vm.define_module("Math", None);
18
19    // Define constants
20    math_module.consts.borrow_mut().insert(
21        "PI".to_string(),
22        RObject::float(std::f64::consts::PI).to_refcount_assigned(),
23    );
24    math_module.consts.borrow_mut().insert(
25        "E".to_string(),
26        RObject::float(std::f64::consts::E).to_refcount_assigned(),
27    );
28
29    // Get the module object to define singleton methods (module methods)
30    let math_module_obj = vm.get_const_by_name("Math").expect("Math module not found");
31
32    // Trigonometric functions
33    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "sin", Box::new(mrb_math_sin));
34    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "cos", Box::new(mrb_math_cos));
35    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "tan", Box::new(mrb_math_tan));
36
37    // Inverse trigonometric functions
38    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "asin", Box::new(mrb_math_asin));
39    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "acos", Box::new(mrb_math_acos));
40    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "atan", Box::new(mrb_math_atan));
41    mrb_define_singleton_cmethod(
42        vm,
43        math_module_obj.clone(),
44        "atan2",
45        Box::new(mrb_math_atan2),
46    );
47
48    // Hyperbolic functions
49    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "sinh", Box::new(mrb_math_sinh));
50    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "cosh", Box::new(mrb_math_cosh));
51    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "tanh", Box::new(mrb_math_tanh));
52
53    // Inverse hyperbolic functions
54    mrb_define_singleton_cmethod(
55        vm,
56        math_module_obj.clone(),
57        "asinh",
58        Box::new(mrb_math_asinh),
59    );
60    mrb_define_singleton_cmethod(
61        vm,
62        math_module_obj.clone(),
63        "acosh",
64        Box::new(mrb_math_acosh),
65    );
66    mrb_define_singleton_cmethod(
67        vm,
68        math_module_obj.clone(),
69        "atanh",
70        Box::new(mrb_math_atanh),
71    );
72
73    // Exponential and logarithmic functions
74    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "exp", Box::new(mrb_math_exp));
75    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "log", Box::new(mrb_math_log));
76    mrb_define_singleton_cmethod(
77        vm,
78        math_module_obj.clone(),
79        "log10",
80        Box::new(mrb_math_log10),
81    );
82    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "log2", Box::new(mrb_math_log2));
83
84    // Root functions
85    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "sqrt", Box::new(mrb_math_sqrt));
86    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "cbrt", Box::new(mrb_math_cbrt));
87
88    // Other mathematical functions
89    mrb_define_singleton_cmethod(
90        vm,
91        math_module_obj.clone(),
92        "hypot",
93        Box::new(mrb_math_hypot),
94    );
95    mrb_define_singleton_cmethod(
96        vm,
97        math_module_obj.clone(),
98        "ldexp",
99        Box::new(mrb_math_ldexp),
100    );
101    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "erf", Box::new(mrb_math_erf));
102    mrb_define_singleton_cmethod(vm, math_module_obj.clone(), "erfc", Box::new(mrb_math_erfc));
103}
104
105// Helper function to get a float from RObject
106fn get_float_arg(obj: &RObject) -> Result<f64, Error> {
107    match &obj.value {
108        RValue::Integer(i) => Ok(*i as f64),
109        RValue::Float(f) => Ok(*f),
110        _ => Err(Error::internal("expected Numeric for Math function")),
111    }
112}
113
114// Helper function to check argument count
115fn check_args_count(args: &[Rc<RObject>], expected: usize) -> Result<Vec<Rc<RObject>>, Error> {
116    if args.len() != expected {
117        return Err(Error::ArgumentError(format!(
118            "wrong number of arguments (given {}, expected {})",
119            args.len(),
120            expected
121        )));
122    }
123
124    Ok(args.to_vec())
125}
126
127// Trigonometric functions
128pub fn mrb_math_sin(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
129    let args = check_args_count(args, 1)?;
130    let x = get_float_arg(&args[0])?;
131    Ok(RObject::float(x.sin()).to_refcount_assigned())
132}
133
134pub fn mrb_math_cos(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
135    let args = check_args_count(args, 1)?;
136    let x = get_float_arg(&args[0])?;
137    Ok(RObject::float(x.cos()).to_refcount_assigned())
138}
139
140pub fn mrb_math_tan(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
141    let args = check_args_count(args, 1)?;
142    let x = get_float_arg(&args[0])?;
143    Ok(RObject::float(x.tan()).to_refcount_assigned())
144}
145
146// Inverse trigonometric functions
147pub fn mrb_math_asin(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
148    let args = check_args_count(args, 1)?;
149    let x = get_float_arg(&args[0])?;
150    Ok(RObject::float(x.asin()).to_refcount_assigned())
151}
152
153pub fn mrb_math_acos(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
154    let args = check_args_count(args, 1)?;
155    let x = get_float_arg(&args[0])?;
156    Ok(RObject::float(x.acos()).to_refcount_assigned())
157}
158
159pub fn mrb_math_atan(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
160    let args = check_args_count(args, 1)?;
161    let x = get_float_arg(&args[0])?;
162    Ok(RObject::float(x.atan()).to_refcount_assigned())
163}
164
165pub fn mrb_math_atan2(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
166    let args = check_args_count(args, 2)?;
167    let y = get_float_arg(&args[0])?;
168    let x = get_float_arg(&args[1])?;
169    Ok(RObject::float(y.atan2(x)).to_refcount_assigned())
170}
171
172// Hyperbolic functions
173pub fn mrb_math_sinh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
174    let args = check_args_count(args, 1)?;
175    let x = get_float_arg(&args[0])?;
176    Ok(RObject::float(x.sinh()).to_refcount_assigned())
177}
178
179pub fn mrb_math_cosh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
180    let args = check_args_count(args, 1)?;
181    let x = get_float_arg(&args[0])?;
182    Ok(RObject::float(x.cosh()).to_refcount_assigned())
183}
184
185pub fn mrb_math_tanh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
186    let args = check_args_count(args, 1)?;
187    let x = get_float_arg(&args[0])?;
188    Ok(RObject::float(x.tanh()).to_refcount_assigned())
189}
190
191// Inverse hyperbolic functions
192pub fn mrb_math_asinh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
193    let args = check_args_count(args, 1)?;
194    let x = get_float_arg(&args[0])?;
195    Ok(RObject::float(x.asinh()).to_refcount_assigned())
196}
197
198pub fn mrb_math_acosh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
199    let args = check_args_count(args, 1)?;
200    let x = get_float_arg(&args[0])?;
201    Ok(RObject::float(x.acosh()).to_refcount_assigned())
202}
203
204pub fn mrb_math_atanh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
205    let args = check_args_count(args, 1)?;
206    let x = get_float_arg(&args[0])?;
207    Ok(RObject::float(x.atanh()).to_refcount_assigned())
208}
209
210// Exponential and logarithmic functions
211pub fn mrb_math_exp(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
212    let args = check_args_count(args, 1)?;
213    let x = get_float_arg(&args[0])?;
214    Ok(RObject::float(x.exp()).to_refcount_assigned())
215}
216
217pub fn mrb_math_log(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
218    if args.len() == 1 {
219        let x = get_float_arg(&args[0])?;
220        Ok(RObject::float(x.ln()).to_refcount_assigned())
221    } else if args.len() == 2 {
222        let x = get_float_arg(&args[0])?;
223        let base = get_float_arg(&args[1])?;
224        Ok(RObject::float(x.log(base)).to_refcount_assigned())
225    } else {
226        Err(Error::ArgumentError(format!(
227            "wrong number of arguments (given {}, expected 1..2)",
228            args.len()
229        )))
230    }
231}
232
233pub fn mrb_math_log10(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
234    let args = check_args_count(args, 1)?;
235    let x = get_float_arg(&args[0])?;
236    Ok(RObject::float(x.log10()).to_refcount_assigned())
237}
238
239pub fn mrb_math_log2(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
240    let args = check_args_count(args, 1)?;
241    let x = get_float_arg(&args[0])?;
242    Ok(RObject::float(x.log2()).to_refcount_assigned())
243}
244
245// Root functions
246pub fn mrb_math_sqrt(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
247    let args = check_args_count(args, 1)?;
248    let x = get_float_arg(&args[0])?;
249    Ok(RObject::float(x.sqrt()).to_refcount_assigned())
250}
251
252pub fn mrb_math_cbrt(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
253    let args = check_args_count(args, 1)?;
254    let x = get_float_arg(&args[0])?;
255    Ok(RObject::float(x.cbrt()).to_refcount_assigned())
256}
257
258// Other mathematical functions
259pub fn mrb_math_hypot(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
260    let args = check_args_count(args, 2)?;
261    let x = get_float_arg(&args[0])?;
262    let y = get_float_arg(&args[1])?;
263    Ok(RObject::float(x.hypot(y)).to_refcount_assigned())
264}
265
266pub fn mrb_math_ldexp(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
267    let args = check_args_count(args, 2)?;
268    let fraction = get_float_arg(&args[0])?;
269    let exponent: i32 = args[1].as_ref().try_into()?;
270    Ok(RObject::float(fraction * 2f64.powi(exponent)).to_refcount_assigned())
271}
272
273pub fn mrb_math_erf(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
274    let args = check_args_count(args, 1)?;
275    let x = get_float_arg(&args[0])?;
276    let result = erf_approximation(x);
277    Ok(RObject::float(result).to_refcount_assigned())
278}
279
280pub fn mrb_math_erfc(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
281    let args = check_args_count(args, 1)?;
282    let x = get_float_arg(&args[0])?;
283    let result = 1.0 - erf_approximation(x);
284    Ok(RObject::float(result).to_refcount_assigned())
285}
286
287// Abramowitz and Stegun approximation of error function
288fn erf_approximation(x: f64) -> f64 {
289    let a1 = 0.254829592;
290    let a2 = -0.284496736;
291    let a3 = 1.421413741;
292    let a4 = -1.453152027;
293    let a5 = 1.061405429;
294    let p = 0.3275911;
295
296    let sign = if x < 0.0 { -1.0 } else { 1.0 };
297    let x = x.abs();
298
299    let t = 1.0 / (1.0 + p * x);
300    let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();
301
302    sign * y
303}