From beeb0441ef70047e0d5f7addd2593d7e9d5671f1 Mon Sep 17 00:00:00 2001
From: Devon Hollowood <devonhollowood@google.com>
Date: Sat, 22 Oct 2022 14:55:14 -0700
Subject: [PATCH] Allow log2() to work on large integers

---
 Cargo.lock            |  5 +++--
 Lib/test/test_math.py |  2 --
 stdlib/Cargo.toml     |  1 +
 stdlib/src/math.rs    | 37 ++++++++++++++++++++++++++++++-------
 4 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 40baabbaa..683e9e004 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1250,9 +1250,9 @@ dependencies = [
 
 [[package]]
 name = "num-rational"
-version = "0.4.0"
+version = "0.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d41702bd167c2df5520b384281bc111a4b5efcf7fbc4c9c222c815b07e0a6a6a"
+checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0"
 dependencies = [
  "autocfg",
  "num-bigint",
@@ -1936,6 +1936,7 @@ dependencies = [
  "num-bigint",
  "num-complex",
  "num-integer",
+ "num-rational",
  "num-traits",
  "num_enum",
  "once_cell",
diff --git a/Lib/test/test_math.py b/Lib/test/test_math.py
index 8e5c43512..be13a363b 100644
--- a/Lib/test/test_math.py
+++ b/Lib/test/test_math.py
@@ -1162,8 +1162,6 @@ class MathTests(unittest.TestCase):
         self.assertRaises(ValueError, math.log1p, -1)
         self.assertEqual(math.log1p(INF), INF)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @requires_IEEE_754
     def testLog2(self):
         self.assertRaises(TypeError, math.log2)
diff --git a/stdlib/Cargo.toml b/stdlib/Cargo.toml
index 34268fc58..7ef5cf34d 100644
--- a/stdlib/Cargo.toml
+++ b/stdlib/Cargo.toml
@@ -53,6 +53,7 @@ bzip2 = { version = "0.4", optional = true }
 num-complex = "0.4.0"
 num-bigint = "0.4.3"
 num-integer = "0.1.44"
+num-rational = "0.4.1"
 
 crossbeam-utils = "=0.8.9"
 itertools = "0.10.3"
diff --git a/stdlib/src/math.rs b/stdlib/src/math.rs
index 34795a6eb..c771cb438 100644
--- a/stdlib/src/math.rs
+++ b/stdlib/src/math.rs
@@ -8,7 +8,8 @@ mod math {
         identifier, PyObject, PyObjectRef, PyRef, PyResult, VirtualMachine,
     };
     use num_bigint::BigInt;
-    use num_traits::{One, Signed, Zero};
+    use num_rational::Ratio;
+    use num_traits::{One, Signed, ToPrimitive, Zero};
     use rustpython_common::float_ops;
     use std::cmp::Ordering;
 
@@ -155,13 +156,35 @@ mod math {
         }
     }
 
+    /// Generates the base-2 logarithm of a BigInt `x`
+    fn ilog2(x: &BigInt) -> f64 {
+        // log2(x) = log2(2^n * 2^-n * x) = n + log2(x/2^n)
+        // If we set 2^n to be the greatest power of 2 below x, then x/2^n is in [1, 2), and can
+        // thus be converted into a float.
+        let n = x.bits() as u32 - 1;
+        let frac = Ratio::new(x.clone(), BigInt::from(2).pow(n));
+        f64::from(n) + frac.to_f64().unwrap().log2()
+    }
+
     #[pyfunction]
-    fn log2(x: ArgIntoFloat, vm: &VirtualMachine) -> PyResult<f64> {
-        let x = *x;
-        if x.is_nan() || x > 0.0_f64 {
-            Ok(x.log2())
-        } else {
-            Err(vm.new_value_error("math domain error".to_owned()))
+    fn log2(x: PyObjectRef, vm: &VirtualMachine) -> PyResult<f64> {
+        match x.try_float(vm) {
+            Ok(x) => {
+                let x = x.to_f64();
+                if x.is_nan() || x > 0.0_f64 {
+                    Ok(x.log2())
+                } else {
+                    Err(vm.new_value_error("math domain error".to_owned()))
+                }
+            }
+            Err(float_err) => {
+                if let Ok(x) = x.try_int(vm) {
+                    Ok(ilog2(x.as_bigint()))
+                } else {
+                    // Return the float error, as it will be more intuitive to users
+                    Err(float_err)
+                }
+            }
         }
     }