From d3dafda7c2de29cd6183413eb5d47f3e31265eea Mon Sep 17 00:00:00 2001
From: TheAnyKey <ich@simonroth.de>
Date: Fri, 8 May 2020 16:31:36 +0000
Subject: [PATCH] Implement LCM and GCD from Py3.9; merged from
 TheAnyKey/p39_math

---
 vm/src/stdlib/math.rs | 55 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 51 insertions(+), 4 deletions(-)

diff --git a/vm/src/stdlib/math.rs b/vm/src/stdlib/math.rs
index 6873e4352..1ddbdf289 100644
--- a/vm/src/stdlib/math.rs
+++ b/vm/src/stdlib/math.rs
@@ -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;
-    a.as_bigint().gcd(b.as_bigint())
+    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;
+    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),