Skip to content

Commit 6b1b15f

Browse files
AdamGSgatesn
andauthored
fix: Fix shr overflow in FoR compare (#2176)
closes #2175 --------- Co-authored-by: Nicholas Gates <[email protected]>
1 parent 8b585e7 commit 6b1b15f

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

encodings/fastlanes/src/for/compute/compare.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ fn compare_constant<T>(
3535
operator: Operator,
3636
) -> VortexResult<Option<Array>>
3737
where
38-
T: NativePType + Shr<u32, Output = T> + WrappingSub,
38+
T: NativePType + WrappingSub + Shr<usize, Output = T>,
3939
T: TryFrom<PValue, Error = VortexError>,
4040
Scalar: From<Option<T>>,
4141
{
@@ -53,11 +53,11 @@ where
5353
if let Some(reference) = reference {
5454
rhs = rhs.wrapping_sub(&reference);
5555
}
56-
if lhs.shift() > 0 {
57-
// Since compare requires that both sides are of same dtype this will always succeed and not panic
58-
rhs = rhs >> (lhs.shift() as u32)
59-
}
60-
rhs
56+
57+
// Since compare requires that both sides are of same dtype, shifting by bit_width should never have an effect.
58+
let bit_width = <T as NativePType>::PTYPE.bit_width();
59+
assert!(bit_width >= lhs.shift() as usize);
60+
rhs >> (lhs.shift() as usize % bit_width)
6161
});
6262

6363
// Wrap up the RHS into a scalar and cast to the encoded DType (this will be the equivalent

0 commit comments

Comments
 (0)