Speed up i256 division and remainder operations

arrow-buffer/benches/i256.rs

@@ -17,8 +17,23 @@
 
 use arrow_buffer::i256;
 use criterion::*;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
 use std::str::FromStr;
 
+/// Returns fixed seedable RNG
+fn seedable_rng() -> StdRng {
+    StdRng::seed_from_u64(42)
+}
+
+fn create_i256_vec(size: usize) -> Vec<i256> {
+    let mut rng = seedable_rng();
+
+    (0..size)
+        .map(|_| i256::from_i128(rng.gen::<i128>()))
+        .collect()
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     let numbers = vec![
         i256::ZERO,
@@ -38,6 +53,28 @@ fn criterion_benchmark(c: &mut Criterion) {
             b.iter(|| i256::from_str(&t).unwrap());
         });
     }
+
+    c.bench_function("i256_div", |b| {
+        b.iter(|| {
+            for number_a in create_i256_vec(10) {
+                for number_b in create_i256_vec(5) {
+                    number_a.checked_div(number_b);
+                    number_a.wrapping_div(number_b);
+                }
+            }
+        });
+    });
+
+    c.bench_function("i256_rem", |b| {
+        b.iter(|| {
+            for number_a in create_i256_vec(10) {
+                for number_b in create_i256_vec(5) {
+                    number_a.checked_rem(number_b);
+                    number_a.wrapping_rem(number_b);
+                }
+            }
+        });
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

arrow-buffer/src/bigint.rs

@@ -22,22 +22,33 @@ use std::num::ParseIntError;
 use std::ops::{BitAnd, BitOr, BitXor, Neg, Shl, Shr};
 use std::str::FromStr;
 
-/// An opaque error similar to [`std::num::ParseIntError`]
+/// [`i256`] operations return this error type.
 #[derive(Debug)]
-pub struct ParseI256Error {}
+pub enum I256Error {
+    /// An opaque error similar to [`std::num::ParseIntError`]
+    ParseError,
+    /// Division by zero
+    DivideByZero,
+    /// Division overflow
+    DivideOverflow,
+}
 
-impl From<ParseIntError> for ParseI256Error {
+impl From<ParseIntError> for I256Error {
     fn from(_: ParseIntError) -> Self {
-        Self {}
+        I256Error::ParseError
     }
 }
 
-impl std::fmt::Display for ParseI256Error {
+impl std::fmt::Display for I256Error {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "Failed to parse as i256")
+        match self {
+            I256Error::ParseError => write!(f, "Failed to parse as i256"),
+            I256Error::DivideByZero => write!(f, "Division by zero"),
+            I256Error::DivideOverflow => write!(f, "Division overflow"),
+        }
     }
 }
-impl std::error::Error for ParseI256Error {}
+impl std::error::Error for I256Error {}
 
 /// A signed 256-bit integer
 #[allow(non_camel_case_types)]
@@ -60,7 +71,7 @@ impl std::fmt::Display for i256 {
 }
 
 impl FromStr for i256 {
-    type Err = ParseI256Error;
+    type Err = I256Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         // i128 can store up to 38 decimal digits
@@ -82,15 +93,15 @@ impl FromStr for i256 {
 
         if !s.as_bytes()[0].is_ascii_digit() {
             // Ensures no duplicate sign
-            return Err(ParseI256Error {});
+            return Err(I256Error::ParseError);
         }
 
         parse_impl(s, negative)
     }
 }
 
 /// Parse `s` with any sign and leading 0s removed
-fn parse_impl(s: &str, negative: bool) -> Result<i256, ParseI256Error> {
+fn parse_impl(s: &str, negative: bool) -> Result<i256, I256Error> {
     if s.len() <= 38 {
         let low = i128::from_str(s)?;
         return Ok(match negative {
@@ -102,7 +113,7 @@ fn parse_impl(s: &str, negative: bool) -> Result<i256, ParseI256Error> {
     let split = s.len() - 38;
     if !s.as_bytes()[split].is_ascii_digit() {
         // Ensures not splitting codepoint and no sign
-        return Err(ParseI256Error {});
+        return Err(I256Error::ParseError);
     }
     let (hs, ls) = s.split_at(split);
 
@@ -117,7 +128,7 @@ fn parse_impl(s: &str, negative: bool) -> Result<i256, ParseI256Error> {
 
     high.checked_mul(i256::from_i128(10_i128.pow(38)))
         .and_then(|high| high.checked_add(low))
-        .ok_or(ParseI256Error {})
+        .ok_or(I256Error::ParseError)
 }
 
 impl PartialOrd for i256 {
@@ -396,42 +407,101 @@ impl i256 {
             .then_some(Self { low, high })
     }
 
+    /// Return the least number of bits needed to represent the number
+    #[inline]
+    fn bits_required(&self) -> usize {
+        let le_bytes = self.to_le_bytes();
+        let arr: [u128; 2] = [
+            u128::from_le_bytes(le_bytes[0..16].try_into().unwrap()),
+            u128::from_le_bytes(le_bytes[16..32].try_into().unwrap()),
+        ];
+
+        let iter = arr.iter().rev().take(2 - 1);
+        if self.is_negative() {
+            let ctr = iter.take_while(|&&b| b == ::core::u128::MAX).count();
+            (128 * (2 - ctr)) + 1 - (!arr[2 - ctr - 1]).leading_zeros() as usize
+        } else {
+            let ctr = iter.take_while(|&&b| b == ::core::u128::MIN).count();
+            (128 * (2 - ctr)) + 1 - arr[2 - ctr - 1].leading_zeros() as usize
+        }
+    }
+
+    /// Division operation, returns (quotient, remainder).
+    /// This basically implements [Long division]: `<https://en.wikipedia.org/wiki/Division_algorithm>`
+    #[inline]
+    fn div_rem(self, other: Self) -> Result<(Self, Self), I256Error> {
+        if other == Self::ZERO {
+            return Err(I256Error::DivideByZero);
+        }
+        if other == Self::MINUS_ONE && self == Self::MIN {
+            return Err(I256Error::DivideOverflow);
+        }
+
+        if self == Self::MIN || other == Self::MIN {
+            let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
+            let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
+            let d = i256::from_bigint_with_overflow(&l / &r).0;
+            let r = i256::from_bigint_with_overflow(&l % &r).0;
+            return Ok((d, r));
+        }
+
+        let mut me = self.checked_abs().unwrap();
+        let mut you = other.checked_abs().unwrap();
+        let mut ret = [0u128; 2];
+        if me < you {
+            return Ok((Self::from_parts(ret[0], ret[1] as i128), self));
+        }
+
+        let shift = me.bits_required() - you.bits_required();
+        you = you.shl(shift as u8);
+        for i in (0..=shift).rev() {
+            if me >= you {
+                ret[i / 128] |= 1 << (i % 128);
+                me = me.checked_sub(you).unwrap();
+            }
+            you = you.shr(1);
+        }
+
+        Ok((
+            if self.is_negative() == other.is_negative() {
+                Self::from_parts(ret[0], ret[1] as i128)
+            } else {
+                -Self::from_parts(ret[0], ret[1] as i128)
+            },
+            if self.is_negative() { -me } else { me },
+        ))
+    }
+
     /// Performs wrapping division
     #[inline]
     pub fn wrapping_div(self, other: Self) -> Self {
-        let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
-        let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
-        Self::from_bigint_with_overflow(l / r).0
+        match self.div_rem(other) {
+            Ok((v, _)) => v,
+            Err(I256Error::DivideByZero) => panic!("attempt to divide by zero"),
+            Err(_) => Self::MIN,
+        }
     }
 
     /// Performs checked division
     #[inline]
     pub fn checked_div(self, other: Self) -> Option<Self> {
-        let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
-        let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
-        let (val, overflow) = Self::from_bigint_with_overflow(l / r);
-        (!overflow).then_some(val)
+        self.div_rem(other).map(|(v, _)| v).ok()
     }
 
     /// Performs wrapping remainder
     #[inline]
     pub fn wrapping_rem(self, other: Self) -> Self {
-        let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
-        let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
-        Self::from_bigint_with_overflow(l % r).0
+        match self.div_rem(other) {
+            Ok((_, v)) => v,
+            Err(I256Error::DivideByZero) => panic!("attempt to divide by zero"),
+            Err(_) => Self::ZERO,
+        }
     }
 
     /// Performs checked remainder
     #[inline]
     pub fn checked_rem(self, other: Self) -> Option<Self> {
-        if other == Self::ZERO {
-            return None;
-        }
-
-        let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
-        let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
-        let (val, overflow) = Self::from_bigint_with_overflow(l % r);
-        (!overflow).then_some(val)
+        self.div_rem(other).map(|(_, v)| v).ok()
     }
 
     /// Performs checked exponentiation
@@ -853,6 +923,43 @@ mod tests {
             ),
         }
 
+        // Division
+        if ir != i256::ZERO {
+            let actual = il.wrapping_div(ir);
+            let expected = bl.clone() / br.clone();
+            let checked = il.checked_div(ir);
+
+            if ir == i256::MINUS_ONE && il == i256::MIN {
+                // BigInt produces an integer over i256::MAX
+                assert_eq!(actual, i256::MIN);
+                assert!(checked.is_none());
+            } else {
+                assert_eq!(actual.to_string(), expected.to_string());
+                assert_eq!(checked.unwrap().to_string(), expected.to_string());
+            }
+        } else {
+            // `wrapping_div` panics on division by zero
+            assert!(il.checked_div(ir).is_none());
+        }
+
+        // Remainder
+        if ir != i256::ZERO {
+            let actual = il.wrapping_rem(ir);
+            let expected = bl.clone() % br.clone();
+            let checked = il.checked_rem(ir);
+
+            assert_eq!(actual.to_string(), expected.to_string());
+
+            if ir == i256::MINUS_ONE && il == i256::MIN {
+                assert!(checked.is_none());
+            } else {
+                assert_eq!(checked.unwrap().to_string(), expected.to_string());
+            }
+        } else {
+            // `wrapping_rem` panics on division by zero
+            assert!(il.checked_rem(ir).is_none());
+        }
+
         // Exponentiation
         for exp in vec![0, 1, 2, 3, 8, 100].into_iter() {
             let actual = il.wrapping_pow(exp);