Skip to content

arrow-buffer: i256: Implement num_traits wrapping shift #9418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
theirix wants to merge 1 commit into
base: main
Choose a base branch
from
Open

arrow-buffer: i256: Implement num_traits wrapping shift #9418

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 52 additions & 2 deletions arrow-buffer/src/bigint/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ use crate::bigint::div::div_rem;
use num_bigint::BigInt;
use num_traits::{
Bounded, CheckedAdd, CheckedDiv, CheckedMul, CheckedNeg, CheckedRem, CheckedSub, FromPrimitive,
Num, One, Signed, ToPrimitive, WrappingAdd, WrappingMul, WrappingNeg, WrappingSub, Zero,
cast::AsPrimitive,
Num, One, Signed, ToPrimitive, WrappingAdd, WrappingMul, WrappingNeg, WrappingShl, WrappingShr,
WrappingSub, Zero, cast::AsPrimitive,
};
use std::cmp::Ordering;
use std::num::ParseIntError;
Expand Down Expand Up @@ -807,6 +807,46 @@ impl Shr<u8> for i256 {
}
}

impl WrappingShl for i256 {
#[inline]
fn wrapping_shl(&self, rhs: u32) -> i256 {
(*self).shl(rhs)
Copy link
Contributor

@Rafferty97 Rafferty97 Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The dereference isn't necessary, and I think the wrapping behaviour should be explicit rather than delegated to the shl implementation. i.e.:

impl WrappingShl for i256 {
    fn wrapping_shl(&self, rhs: u32) -> Self {
        // Truncating to a `u8` masks off the higher-order bits as required
        self.shl(rhs as u8)
    }
}

}
}

impl WrappingShr for i256 {
#[inline]
fn wrapping_shr(&self, rhs: u32) -> i256 {
Copy link
Contributor

@Rafferty97 Rafferty97 Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly, I'd suggest:

impl WrappingShr for i256 {
    fn wrapping_shr(&self, rhs: u32) -> Self {
        // Truncating to a `u8` masks off the higher-order bits as required
        self.shr(rhs as u8)
    }
}

(*self).shr(rhs)
}
}

// Define Shr<T> and Shl<T> for specified integer types
macro_rules! define_wrapping_shift {
Copy link
Contributor

@Rafferty97 Rafferty97 Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sold on the idea of implementing Shl/Shr with wrapping semantics, as this is inconsistent with the other built-in integer types, which panic instead.

Copy link
Contributor Author

@theirix theirix Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The behaviour for num_traits and core Shl is different

At the same time, WrappingShl requires Shl to be satisfied. I guess we can implement a panicking core Shl (to be consistent with core trait), and provide a non-panicking wrapping WrappingShl::wrapping_shl behaviour instead of delegating to Shl::shl as you suggested in https://github.com/apache/arrow-rs/pull/9418/changes#r2928017509

Copy link
Contributor

@Rafferty97 Rafferty97 Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, that's exactly what I'd suggest. Technically only Shl<usize> needs to be implemented, so you could just do so directly, but the macro is fine too.

Copy link
Contributor Author

@theirix theirix Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, thank you. I'll update the implementation and get back to you soon.

// Handle multiple types
($trait_name:ident, $method:ident, [$($t:ty),+]) => {
$(define_wrapping_shift!($trait_name, $method, $t);)+
};
// Handle single type
($trait_name:ident, $method:ident, $t:ty) => {
impl $trait_name<$t> for i256 {
type Output = i256;

#[inline]
fn $method(self, rhs: $t) -> Self::Output {
// Take modulo 256
#[allow(clippy::suspicious_arithmetic_impl)]
let shift: u8 = (rhs % (256 as $t)) as u8;
// Use existing Shl<u8> / Shr<u8> implementation
<Self as $trait_name<u8>>::$method(self, shift)
}
}
};
}

define_wrapping_shift!(Shl, shl, [u16, u32, u64, usize, i16, i32, i64, isize]);
define_wrapping_shift!(Shr, shr, [u16, u32, u64, usize, i16, i32, i64, isize]);

macro_rules! define_as_primitive {
($native_ty:ty) => {
impl AsPrimitive<i256> for $native_ty {
Expand Down Expand Up @@ -1190,9 +1230,19 @@ mod tests {
let (expected, _) = i256::from_bigint_with_overflow(bl.clone() << shift);
assert_eq!(actual.to_string(), expected.to_string());

let wrapping_actual = <i256 as WrappingShl>::wrapping_shl(&il, shift as u32);
assert_eq!(wrapping_actual.to_string(), expected.to_string());

let actual = il >> shift;
let (expected, _) = i256::from_bigint_with_overflow(bl.clone() >> shift);
assert_eq!(actual.to_string(), expected.to_string());

let wrapping_actual = <i256 as WrappingShr>::wrapping_shr(&il, shift as u32);
assert_eq!(wrapping_actual.to_string(), expected.to_string());

// Check wrapping of the shift argument
let wrapping_actual = <i256 as WrappingShr>::wrapping_shr(&il, 512 + shift as u32);
assert_eq!(wrapping_actual.to_string(), expected.to_string());
}
}

Expand Down
Loading