extend macro rule

Signed-off-by: Ruihang Xia <[email protected]>
apache · Jun 2, 2024 · 8ab09e1 · 8ab09e1
1 parent 95da017
commit 8ab09e1
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 54 deletions.
diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml
@@ -37,6 +37,7 @@ bench = false
 bytes = { version = "1.4" }
 num = { version = "0.4", default-features = false, features = ["std"] }
 half = { version = "2.1", default-features = false }
+paste = { version = "1.0" }
 
 [dev-dependencies]
 criterion = { version = "0.5", default-features = false }

diff --git a/arrow-buffer/src/arith.rs b/arrow-buffer/src/arith.rs
@@ -15,7 +15,7 @@
 // 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) => {
@@ -34,6 +34,20 @@ macro_rules! derive_arith {
             }
         }
 
+        ::paste::paste! {
+            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/interval.rs b/arrow-buffer/src/interval.rs
@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::arith::derive_arith;
-use std::ops::{AddAssign, Neg, SubAssign};
+use std::ops::Neg;
 
 /// Value of an IntervalMonthDayNano array
 #[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
@@ -225,34 +225,6 @@ impl Neg for IntervalMonthDayNano {
     }
 }
 
-impl AddAssign for IntervalMonthDayNano {
-    #[cfg(debug_assertions)]
-    fn add_assign(&mut self, rhs: Self) {
-        *self = self
-            .checked_add(rhs)
-            .expect("IntervalMonthDayNano overflow");
-    }
-
-    #[cfg(not(debug_assertions))]
-    fn add_assign(&mut self, rhs: Self) {
-        *self = self.wrapping_add(rhs);
-    }
-}
-
-impl SubAssign for IntervalMonthDayNano {
-    #[cfg(debug_assertions)]
-    fn sub_assign(&mut self, rhs: Self) {
-        *self = self
-            .checked_sub(rhs)
-            .expect("IntervalMonthDayNano underflow");
-    }
-
-    #[cfg(not(debug_assertions))]
-    fn sub_assign(&mut self, rhs: Self) {
-        *self = self.wrapping_sub(rhs);
-    }
-}
-
 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);
@@ -445,30 +417,6 @@ impl Neg for IntervalDayTime {
     }
 }
 
-impl AddAssign for IntervalDayTime {
-    #[cfg(debug_assertions)]
-    fn add_assign(&mut self, rhs: Self) {
-        *self = self.checked_add(rhs).expect("IntervalDayTime overflow");
-    }
-
-    #[cfg(not(debug_assertions))]
-    fn add_assign(&mut self, rhs: Self) {
-        *self = self.wrapping_add(rhs);
-    }
-}
-
-impl SubAssign for IntervalDayTime {
-    #[cfg(debug_assertions)]
-    fn sub_assign(&mut self, rhs: Self) {
-        *self = self.checked_sub(rhs).expect("IntervalDayTime underflow");
-    }
-
-    #[cfg(not(debug_assertions))]
-    fn sub_assign(&mut self, rhs: Self) {
-        *self = self.wrapping_sub(rhs);
-    }
-}
-
 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);