Skip to main content

mruby_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 (excluding trailing nil)
111fn check_args_count(args: &[Rc<RObject>], expected: usize) -> Result<Vec<Rc<RObject>>, Error> {
112    let args = if !args.is_empty() && args[args.len() - 1].is_nil() {
113        &args[0..args.len() - 1]
114    } else {
115        args
116    };
117
118    if args.len() != expected {
119        return Err(Error::ArgumentError(format!(
120            "wrong number of arguments (given {}, expected {})",
121            args.len(),
122            expected
123        )));
124    }
125
126    Ok(args.to_vec())
127}
128
129// Trigonometric functions
130pub fn mrb_math_sin(_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.sin()).to_refcount_assigned())
134}
135
136pub fn mrb_math_cos(_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.cos()).to_refcount_assigned())
140}
141
142pub fn mrb_math_tan(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
143    let args = check_args_count(args, 1)?;
144    let x = get_float_arg(&args[0])?;
145    Ok(RObject::float(x.tan()).to_refcount_assigned())
146}
147
148// Inverse trigonometric functions
149pub fn mrb_math_asin(_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.asin()).to_refcount_assigned())
153}
154
155pub fn mrb_math_acos(_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.acos()).to_refcount_assigned())
159}
160
161pub fn mrb_math_atan(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
162    let args = check_args_count(args, 1)?;
163    let x = get_float_arg(&args[0])?;
164    Ok(RObject::float(x.atan()).to_refcount_assigned())
165}
166
167pub fn mrb_math_atan2(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
168    let args = check_args_count(args, 2)?;
169    let y = get_float_arg(&args[0])?;
170    let x = get_float_arg(&args[1])?;
171    Ok(RObject::float(y.atan2(x)).to_refcount_assigned())
172}
173
174// Hyperbolic functions
175pub fn mrb_math_sinh(_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.sinh()).to_refcount_assigned())
179}
180
181pub fn mrb_math_cosh(_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.cosh()).to_refcount_assigned())
185}
186
187pub fn mrb_math_tanh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
188    let args = check_args_count(args, 1)?;
189    let x = get_float_arg(&args[0])?;
190    Ok(RObject::float(x.tanh()).to_refcount_assigned())
191}
192
193// Inverse hyperbolic functions
194pub fn mrb_math_asinh(_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.asinh()).to_refcount_assigned())
198}
199
200pub fn mrb_math_acosh(_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.acosh()).to_refcount_assigned())
204}
205
206pub fn mrb_math_atanh(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
207    let args = check_args_count(args, 1)?;
208    let x = get_float_arg(&args[0])?;
209    Ok(RObject::float(x.atanh()).to_refcount_assigned())
210}
211
212// Exponential and logarithmic functions
213pub fn mrb_math_exp(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
214    let args = check_args_count(args, 1)?;
215    let x = get_float_arg(&args[0])?;
216    Ok(RObject::float(x.exp()).to_refcount_assigned())
217}
218
219pub fn mrb_math_log(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
220    let args_vec = if !args.is_empty() && args[args.len() - 1].is_nil() {
221        args[0..args.len() - 1].to_vec()
222    } else {
223        args.to_vec()
224    };
225
226    if args_vec.len() == 1 {
227        let x = get_float_arg(&args_vec[0])?;
228        Ok(RObject::float(x.ln()).to_refcount_assigned())
229    } else if args_vec.len() == 2 {
230        let x = get_float_arg(&args_vec[0])?;
231        let base = get_float_arg(&args_vec[1])?;
232        Ok(RObject::float(x.log(base)).to_refcount_assigned())
233    } else {
234        Err(Error::ArgumentError(format!(
235            "wrong number of arguments (given {}, expected 1..2)",
236            args_vec.len()
237        )))
238    }
239}
240
241pub fn mrb_math_log10(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
242    let args = check_args_count(args, 1)?;
243    let x = get_float_arg(&args[0])?;
244    Ok(RObject::float(x.log10()).to_refcount_assigned())
245}
246
247pub fn mrb_math_log2(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
248    let args = check_args_count(args, 1)?;
249    let x = get_float_arg(&args[0])?;
250    Ok(RObject::float(x.log2()).to_refcount_assigned())
251}
252
253// Root functions
254pub fn mrb_math_sqrt(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
255    let args = check_args_count(args, 1)?;
256    let x = get_float_arg(&args[0])?;
257    Ok(RObject::float(x.sqrt()).to_refcount_assigned())
258}
259
260pub fn mrb_math_cbrt(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
261    let args = check_args_count(args, 1)?;
262    let x = get_float_arg(&args[0])?;
263    Ok(RObject::float(x.cbrt()).to_refcount_assigned())
264}
265
266// Other mathematical functions
267pub fn mrb_math_hypot(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
268    let args = check_args_count(args, 2)?;
269    let x = get_float_arg(&args[0])?;
270    let y = get_float_arg(&args[1])?;
271    Ok(RObject::float(x.hypot(y)).to_refcount_assigned())
272}
273
274pub fn mrb_math_ldexp(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
275    let args = check_args_count(args, 2)?;
276    let fraction = get_float_arg(&args[0])?;
277    let exponent: i32 = args[1].as_ref().try_into()?;
278    Ok(RObject::float(fraction * 2f64.powi(exponent)).to_refcount_assigned())
279}
280
281pub fn mrb_math_erf(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
282    let args = check_args_count(args, 1)?;
283    let x = get_float_arg(&args[0])?;
284    let result = erf_approximation(x);
285    Ok(RObject::float(result).to_refcount_assigned())
286}
287
288pub fn mrb_math_erfc(_vm: &mut VM, args: &[Rc<RObject>]) -> Result<Rc<RObject>, Error> {
289    let args = check_args_count(args, 1)?;
290    let x = get_float_arg(&args[0])?;
291    let result = 1.0 - erf_approximation(x);
292    Ok(RObject::float(result).to_refcount_assigned())
293}
294
295// Abramowitz and Stegun approximation of error function
296fn erf_approximation(x: f64) -> f64 {
297    let a1 = 0.254829592;
298    let a2 = -0.284496736;
299    let a3 = 1.421413741;
300    let a4 = -1.453152027;
301    let a5 = 1.061405429;
302    let p = 0.3275911;
303
304    let sign = if x < 0.0 { -1.0 } else { 1.0 };
305    let x = x.abs();
306
307    let t = 1.0 / (1.0 + p * x);
308    let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();
309
310    sign * y
311}