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