Implement LCM and GCD from Py3.9; merged from TheAnyKey/p39_math
@@ -9,9 +9,9 @@ use statrs::function::gamma::{gamma, ln_gamma};
use num_bigint::BigInt;
use num_traits::{One, Zero};
-use crate::function::OptionalArg;
+use crate::function::{OptionalArg, PyFuncArgs};
use crate::obj::objfloat::{self, IntoPyFloat, PyFloatRef};
-use crate::obj::objint::{self, PyIntRef};
+use crate::obj::objint::{self, PyInt, PyIntRef};
use crate::obj::objtype;
use crate::pyobject::{Either, PyObjectRef, PyResult, TypeProtocol};
use crate::vm::VirtualMachine;
@@ -272,9 +272,55 @@ fn math_ldexp(
Ok(value * (2_f64).powf(objint::try_float(i.as_bigint(), vm)?))
}
-fn math_gcd(a: PyIntRef, b: PyIntRef) -> BigInt {
+fn math_perf_arb_len_int_op<F>(
+ args: PyFuncArgs,
+ vm: &VirtualMachine,
+ op: F,
+ default: BigInt,
+) -> PyResult<BigInt>
+where
+ F: Fn(&BigInt, &PyInt) -> BigInt,
+{
+ if !args.kwargs.is_empty() {
+ Err(vm.new_type_error("Takes no keyword arguments".to_owned()))
+ } else if args.args.is_empty() {
+ Ok(default)
+ } else if args.args.len() == 1 {
+ let a: PyObjectRef = args.args[0].clone();
+ if let Some(aa) = a.payload_if_subclass::<PyInt>(vm) {
+ let res = op(aa.as_bigint(), aa);
+ Ok(res)
+ } else {
+ Err(vm.new_type_error("Only integer arguments are supported".to_owned()))
+ }
+ } else {
+ let a = args.args[0].clone();
+ if let Some(aa) = a.payload_if_subclass::<PyInt>(vm) {
+ let mut res = aa.as_bigint().clone();
+ for b in args.args[1..].iter() {
+ if let Some(bb) = b.payload_if_subclass::<PyInt>(vm) {
+ res = op(&res, bb);
+ } else {
+ return Err(
+ vm.new_type_error("Only integer arguments are supported".to_owned())
+ );
+ }
+ }
+ Ok(res)
+ } else {
+ Err(vm.new_type_error("Only integer arguments are supported".to_owned()))
+ }
+ }
+}
+
+fn math_gcd(args: PyFuncArgs, vm: &VirtualMachine) -> PyResult<BigInt> {
+ use num_integer::Integer;
+ math_perf_arb_len_int_op(args, vm, |x, y| x.gcd(y.as_bigint()), BigInt::zero())
+}
+
+fn math_lcm(args: PyFuncArgs, vm: &VirtualMachine) -> PyResult<BigInt> {
use num_integer::Integer;
- a.as_bigint().gcd(b.as_bigint())
+ math_perf_arb_len_int_op(args, vm, |x, y| x.lcm(y.as_bigint()), BigInt::one())
}
fn math_factorial(value: PyIntRef, vm: &VirtualMachine) -> PyResult<BigInt> {
@@ -436,6 +482,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
// Gcd function
"gcd" => ctx.new_function(math_gcd),
+ "lcm" => ctx.new_function(math_lcm),
// Factorial function
"factorial" => ctx.new_function(math_factorial),
