feat: impl *Assign ops for types in arrow-buffer (#5832)

* naive impl Signed-off-by: Ruihang Xia <waynestxia@gmail.com> * extend macro rule Signed-off-by: Ruihang Xia <waynestxia@gmail.com> * remove quote! Signed-off-by: Ruihang Xia <waynestxia@gmail.com> --------- Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
apache · Jun 3, 2024 · 13a26d6 · 13a26d6
1 parent 065151b
commit 13a26d6
Show file tree

Hide file tree

Showing 3 changed files with 151 additions and 17 deletions.
diff --git 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
@@ -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
@@ -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
+);