Skip to content

Commit d3b0c87

Browse files
committed
Add additional math functions to math module
1 parent bccb6f3 commit d3b0c87

File tree

1 file changed

+203
-0
lines changed

1 file changed

+203
-0
lines changed

vm/src/stdlib/math.rs

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,123 @@ macro_rules! make_math_func {
2222
};
2323
}
2424

25+
// Number theory functions:
26+
make_math_func!(math_fabs, abs);
27+
28+
fn math_isfinite(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
29+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
30+
let value = objfloat::get_value(value).is_finite();
31+
Ok(vm.ctx.new_bool(value))
32+
}
33+
34+
fn math_isinf(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
35+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
36+
let value = objfloat::get_value(value).is_infinite();
37+
Ok(vm.ctx.new_bool(value))
38+
}
39+
40+
fn math_isnan(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
41+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
42+
let value = objfloat::get_value(value).is_nan();
43+
Ok(vm.ctx.new_bool(value))
44+
}
45+
46+
// Power and logarithmic functions:
47+
make_math_func!(math_exp, exp);
48+
make_math_func!(math_expm1, exp_m1);
49+
50+
fn math_log(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
51+
arg_check!(
52+
vm,
53+
args,
54+
required = [(x, Some(vm.ctx.float_type()))],
55+
optional = [(base, Some(vm.ctx.float_type()))]
56+
);
57+
let x = objfloat::get_value(x);
58+
match base {
59+
None => Ok(vm.ctx.new_float(x.ln())),
60+
Some(base) => {
61+
let base = objfloat::get_value(base);
62+
Ok(vm.ctx.new_float(x.log(base)))
63+
}
64+
}
65+
}
66+
67+
fn math_log1p(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
68+
arg_check!(vm, args, required = [(x, Some(vm.ctx.float_type()))]);
69+
let x = objfloat::get_value(x);
70+
Ok(vm.ctx.new_float((x + 1.0).ln()))
71+
}
72+
73+
make_math_func!(math_log2, log2);
74+
make_math_func!(math_log10, log10);
75+
76+
fn math_pow(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
77+
arg_check!(
78+
vm,
79+
args,
80+
required = [
81+
(x, Some(vm.ctx.float_type())),
82+
(y, Some(vm.ctx.float_type()))
83+
]
84+
);
85+
let x = objfloat::get_value(x);
86+
let y = objfloat::get_value(y);
87+
Ok(vm.ctx.new_float(x.powf(y)))
88+
}
89+
90+
make_math_func!(math_sqrt, sqrt);
91+
2592
// Trigonometric functions:
2693
make_math_func!(math_acos, acos);
2794
make_math_func!(math_asin, asin);
2895
make_math_func!(math_atan, atan);
96+
97+
fn math_atan2(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
98+
arg_check!(
99+
vm,
100+
args,
101+
required = [
102+
(y, Some(vm.ctx.float_type())),
103+
(x, Some(vm.ctx.float_type()))
104+
]
105+
);
106+
let y = objfloat::get_value(y);
107+
let x = objfloat::get_value(x);
108+
Ok(vm.ctx.new_float(y.atan2(x)))
109+
}
110+
29111
make_math_func!(math_cos, cos);
112+
113+
fn math_hypot(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
114+
arg_check!(
115+
vm,
116+
args,
117+
required = [
118+
(x, Some(vm.ctx.float_type())),
119+
(y, Some(vm.ctx.float_type()))
120+
]
121+
);
122+
let x = objfloat::get_value(x);
123+
let y = objfloat::get_value(y);
124+
Ok(vm.ctx.new_float(x.hypot(y)))
125+
}
126+
30127
make_math_func!(math_sin, sin);
31128
make_math_func!(math_tan, tan);
32129

130+
fn math_degrees(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
131+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
132+
let x = objfloat::get_value(value);
133+
Ok(vm.ctx.new_float(x * (180.0 / std::f64::consts::PI)))
134+
}
135+
136+
fn math_radians(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
137+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
138+
let x = objfloat::get_value(value);
139+
Ok(vm.ctx.new_float(x * (std::f64::consts::PI / 180.0)))
140+
}
141+
33142
// Hyperbolic functions:
34143
make_math_func!(math_acosh, acosh);
35144
make_math_func!(math_asinh, asinh);
@@ -38,21 +147,115 @@ make_math_func!(math_cosh, cosh);
38147
make_math_func!(math_sinh, sinh);
39148
make_math_func!(math_tanh, tanh);
40149

150+
// Special functions:
151+
fn math_erf(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
152+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
153+
let x = objfloat::get_value(value);
154+
155+
if x.is_nan() {
156+
Ok(vm.ctx.new_float(x))
157+
} else {
158+
// TODO: implement algorithm
159+
unimplemented!("TODO");
160+
}
161+
}
162+
163+
fn math_erfc(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
164+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
165+
let x = objfloat::get_value(value);
166+
167+
if x.is_nan() {
168+
Ok(vm.ctx.new_float(x))
169+
} else {
170+
// TODO: implement algorithm
171+
unimplemented!("TODO");
172+
}
173+
}
174+
175+
fn math_gamma(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
176+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
177+
let x = objfloat::get_value(value);
178+
179+
if x.is_finite() {
180+
// TODO: implement algorithm
181+
unimplemented!("TODO");
182+
} else {
183+
if x.is_nan() || x.is_sign_positive() {
184+
Ok(vm.ctx.new_float(x))
185+
} else {
186+
Ok(vm.ctx.new_float(std::f64::NAN))
187+
}
188+
}
189+
}
190+
191+
fn math_lgamma(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
192+
arg_check!(vm, args, required = [(value, Some(vm.ctx.float_type()))]);
193+
let x = objfloat::get_value(value);
194+
195+
if x.is_finite() {
196+
// TODO: implement algorithm
197+
unimplemented!("TODO");
198+
} else {
199+
if x.is_nan() {
200+
Ok(vm.ctx.new_float(x))
201+
} else {
202+
Ok(vm.ctx.new_float(std::f64::INFINITY))
203+
}
204+
}
205+
}
206+
41207
pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
42208
let py_mod = ctx.new_module(&"math".to_string(), ctx.new_scope(None));
209+
210+
// Number theory functions:
211+
py_mod.set_item("fabs", ctx.new_rustfunc(math_fabs));
212+
py_mod.set_item("isfinite", ctx.new_rustfunc(math_isfinite));
213+
py_mod.set_item("isinf", ctx.new_rustfunc(math_isinf));
214+
py_mod.set_item("isnan", ctx.new_rustfunc(math_isnan));
215+
216+
// Power and logarithmic functions:
217+
py_mod.set_item("exp", ctx.new_rustfunc(math_exp));
218+
py_mod.set_item("expm1", ctx.new_rustfunc(math_expm1));
219+
py_mod.set_item("log", ctx.new_rustfunc(math_log));
220+
py_mod.set_item("log1p", ctx.new_rustfunc(math_log1p));
221+
py_mod.set_item("log2", ctx.new_rustfunc(math_log2));
222+
py_mod.set_item("log10", ctx.new_rustfunc(math_log10));
223+
py_mod.set_item("pow", ctx.new_rustfunc(math_pow));
224+
py_mod.set_item("sqrt", ctx.new_rustfunc(math_sqrt));
225+
226+
// Trigonometric functions:
43227
py_mod.set_item("acos", ctx.new_rustfunc(math_acos));
44228
py_mod.set_item("asin", ctx.new_rustfunc(math_asin));
45229
py_mod.set_item("atan", ctx.new_rustfunc(math_atan));
230+
py_mod.set_item("atan2", ctx.new_rustfunc(math_atan2));
46231
py_mod.set_item("cos", ctx.new_rustfunc(math_cos));
232+
py_mod.set_item("hypot", ctx.new_rustfunc(math_hypot));
47233
py_mod.set_item("sin", ctx.new_rustfunc(math_sin));
48234
py_mod.set_item("tan", ctx.new_rustfunc(math_tan));
49235

236+
py_mod.set_item("degrees", ctx.new_rustfunc(math_degrees));
237+
py_mod.set_item("radians", ctx.new_rustfunc(math_radians));
238+
239+
// Hyperbolic functions:
50240
py_mod.set_item("acosh", ctx.new_rustfunc(math_acosh));
51241
py_mod.set_item("asinh", ctx.new_rustfunc(math_asinh));
52242
py_mod.set_item("atanh", ctx.new_rustfunc(math_atanh));
53243
py_mod.set_item("cosh", ctx.new_rustfunc(math_cosh));
54244
py_mod.set_item("sinh", ctx.new_rustfunc(math_sinh));
55245
py_mod.set_item("tanh", ctx.new_rustfunc(math_tanh));
56246

247+
// Special functions:
248+
py_mod.set_item("erf", ctx.new_rustfunc(math_erf));
249+
py_mod.set_item("erfc", ctx.new_rustfunc(math_erfc));
250+
py_mod.set_item("gamma", ctx.new_rustfunc(math_gamma));
251+
py_mod.set_item("lgamma", ctx.new_rustfunc(math_lgamma));
252+
253+
// Constants:
254+
py_mod.set_item("pi", ctx.new_float(std::f64::consts::PI)); // 3.14159...
255+
py_mod.set_item("e", ctx.new_float(std::f64::consts::E)); // 2.71..
256+
py_mod.set_item("tau", ctx.new_float(2.0 * std::f64::consts::PI));
257+
py_mod.set_item("inf", ctx.new_float(std::f64::INFINITY));
258+
py_mod.set_item("nan", ctx.new_float(std::f64::NAN));
259+
57260
py_mod
58261
}

0 commit comments

Comments
 (0)