[arrow-rs] branch master updated: Add pow to i256 (#2955)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 63417b1f3 Add pow to i256 (#2955)
63417b1f3 is described below

commit 63417b1f36461d4abf7c0aedc6f49e740a92e687
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Thu Oct 27 14:46:02 2022 -0700

    Add pow to i256 (#2955)
---
 arrow-buffer/src/bigint.rs | 82 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/arrow-buffer/src/bigint.rs b/arrow-buffer/src/bigint.rs
index 3518b85e4..fe135b329 100644
--- a/arrow-buffer/src/bigint.rs
+++ b/arrow-buffer/src/bigint.rs
@@ -305,6 +305,55 @@ impl i256 {
         let (val, overflow) = Self::from_bigint_with_overflow(l % r);
         (!overflow).then_some(val)
     }
+
+    /// Performs checked exponentiation
+    #[inline]
+    pub fn checked_pow(self, mut exp: u32) -> Option<Self> {
+        if exp == 0 {
+            return Some(i256::from_i128(1));
+        }
+
+        let mut base = self;
+        let mut acc: Self = i256::from_i128(1);
+
+        while exp > 1 {
+            if (exp & 1) == 1 {
+                acc = acc.checked_mul(base)?;
+            }
+            exp /= 2;
+            base = base.checked_mul(base)?;
+        }
+        // since exp!=0, finally the exp must be 1.
+        // Deal with the final bit of the exponent separately, since
+        // squaring the base afterwards is not necessary and may cause a
+        // needless overflow.
+        acc.checked_mul(base)
+    }
+
+    /// Performs wrapping exponentiation
+    #[inline]
+    pub fn wrapping_pow(self, mut exp: u32) -> Self {
+        if exp == 0 {
+            return i256::from_i128(1);
+        }
+
+        let mut base = self;
+        let mut acc: Self = i256::from_i128(1);
+
+        while exp > 1 {
+            if (exp & 1) == 1 {
+                acc = acc.wrapping_mul(base);
+            }
+            exp /= 2;
+            base = base.wrapping_mul(base);
+        }
+
+        // since exp!=0, finally the exp must be 1.
+        // Deal with the final bit of the exponent separately, since
+        // squaring the base afterwards is not necessary and may cause a
+        // needless overflow.
+        acc.wrapping_mul(base)
+    }
 }
 
 /// Performs an unsigned multiplication of `a * b` returning a tuple of
@@ -455,6 +504,39 @@ mod tests {
                 expected
             ),
         }
+
+        // Exponentiation
+        for exp in vec![0, 1, 3, 8, 100].into_iter() {
+            let actual = il.wrapping_pow(exp);
+            let (expected, overflow) =
+                i256::from_bigint_with_overflow(bl.clone().pow(exp));
+            assert_eq!(actual.to_string(), expected.to_string());
+
+            let checked = il.checked_pow(exp);
+            match overflow {
+                true => assert!(
+                    checked.is_none(),
+                    "{} ^ {} = {} vs {} * {} = {}",
+                    il,
+                    exp,
+                    actual,
+                    bl,
+                    exp,
+                    expected
+                ),
+                false => assert_eq!(
+                    checked.unwrap(),
+                    actual,
+                    "{} ^ {} = {} vs {} * {} = {}",
+                    il,
+                    exp,
+                    actual,
+                    bl,
+                    exp,
+                    expected
+                ),
+            }
+        }
     }
 
     #[test]