diff --git a/arrow-buffer/src/arith.rs b/arrow-buffer/src/arith.rs index ca693c3607dc..a576b2677131 100644 --- a/arrow-buffer/src/arith.rs +++ b/arrow-buffer/src/arith.rs @@ -15,10 +15,10 @@ // specific language governing permissions and limitations // under the License. -/// Derives `std::ops::$op` for `$ty` calling `$wrapping` or `$checked` variants +/// Derives `std::ops::$t` for `$ty` calling `$wrapping` or `$checked` variants /// based on if debug_assertions enabled macro_rules! derive_arith { - ($ty:ty, $t:ident, $op:ident, $wrapping:ident, $checked:ident) => { + ($ty:ty, $t:ident, $t_assign:ident, $op:ident, $op_assign:ident, $wrapping:ident, $checked:ident) => { impl std::ops::$t for $ty { type Output = $ty; @@ -34,6 +34,20 @@ macro_rules! derive_arith { } } + impl std::ops::$t_assign for $ty { + #[cfg(debug_assertions)] + fn $op_assign(&mut self, rhs: Self) { + *self = self + .$checked(rhs) + .expect(concat!(stringify!($ty), " overflow")); + } + + #[cfg(not(debug_assertions))] + fn $op_assign(&mut self, rhs: Self) { + *self = self.$wrapping(rhs); + } + } + impl<'a> std::ops::$t<$ty> for &'a $ty { type Output = $ty; diff --git a/arrow-buffer/src/bigint/mod.rs b/arrow-buffer/src/bigint/mod.rs index bbe65b073aa6..c3296fed3753 100644 --- a/arrow-buffer/src/bigint/mod.rs +++ b/arrow-buffer/src/bigint/mod.rs @@ -639,11 +639,51 @@ fn mulx(a: u128, b: u128) -> (u128, u128) { (low, high) } -derive_arith!(i256, Add, add, wrapping_add, checked_add); -derive_arith!(i256, Sub, sub, wrapping_sub, checked_sub); -derive_arith!(i256, Mul, mul, wrapping_mul, checked_mul); -derive_arith!(i256, Div, div, wrapping_div, checked_div); -derive_arith!(i256, Rem, rem, wrapping_rem, checked_rem); +derive_arith!( + i256, + Add, + AddAssign, + add, + add_assign, + wrapping_add, + checked_add +); +derive_arith!( + i256, + Sub, + SubAssign, + sub, + sub_assign, + wrapping_sub, + checked_sub +); +derive_arith!( + i256, + Mul, + MulAssign, + mul, + mul_assign, + wrapping_mul, + checked_mul +); +derive_arith!( + i256, + Div, + DivAssign, + div, + div_assign, + wrapping_div, + checked_div +); +derive_arith!( + i256, + Rem, + RemAssign, + rem, + rem_assign, + wrapping_rem, + checked_rem +); impl Neg for i256 { type Output = i256; diff --git a/arrow-buffer/src/interval.rs b/arrow-buffer/src/interval.rs index 7e8043e9a724..bed3b2e31ada 100644 --- a/arrow-buffer/src/interval.rs +++ b/arrow-buffer/src/interval.rs @@ -225,11 +225,51 @@ impl Neg for IntervalMonthDayNano { } } -derive_arith!(IntervalMonthDayNano, Add, add, wrapping_add, checked_add); -derive_arith!(IntervalMonthDayNano, Sub, sub, wrapping_sub, checked_sub); -derive_arith!(IntervalMonthDayNano, Mul, mul, wrapping_mul, checked_mul); -derive_arith!(IntervalMonthDayNano, Div, div, wrapping_div, checked_div); -derive_arith!(IntervalMonthDayNano, Rem, rem, wrapping_rem, checked_rem); +derive_arith!( + IntervalMonthDayNano, + Add, + AddAssign, + add, + add_assign, + wrapping_add, + checked_add +); +derive_arith!( + IntervalMonthDayNano, + Sub, + SubAssign, + sub, + sub_assign, + wrapping_sub, + checked_sub +); +derive_arith!( + IntervalMonthDayNano, + Mul, + MulAssign, + mul, + mul_assign, + wrapping_mul, + checked_mul +); +derive_arith!( + IntervalMonthDayNano, + Div, + DivAssign, + div, + div_assign, + wrapping_div, + checked_div +); +derive_arith!( + IntervalMonthDayNano, + Rem, + RemAssign, + rem, + rem_assign, + wrapping_rem, + checked_rem +); /// Value of an IntervalDayTime array #[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)] @@ -417,8 +457,48 @@ impl Neg for IntervalDayTime { } } -derive_arith!(IntervalDayTime, Add, add, wrapping_add, checked_add); -derive_arith!(IntervalDayTime, Sub, sub, wrapping_sub, checked_sub); -derive_arith!(IntervalDayTime, Mul, mul, wrapping_mul, checked_mul); -derive_arith!(IntervalDayTime, Div, div, wrapping_div, checked_div); -derive_arith!(IntervalDayTime, Rem, rem, wrapping_rem, checked_rem); +derive_arith!( + IntervalDayTime, + Add, + AddAssign, + add, + add_assign, + wrapping_add, + checked_add +); +derive_arith!( + IntervalDayTime, + Sub, + SubAssign, + sub, + sub_assign, + wrapping_sub, + checked_sub +); +derive_arith!( + IntervalDayTime, + Mul, + MulAssign, + mul, + mul_assign, + wrapping_mul, + checked_mul +); +derive_arith!( + IntervalDayTime, + Div, + DivAssign, + div, + div_assign, + wrapping_div, + checked_div +); +derive_arith!( + IntervalDayTime, + Rem, + RemAssign, + rem, + rem_assign, + wrapping_rem, + checked_rem +);