address review comments

apache · Sep 6, 2024 · a15db14 · a15db14
1 parent 58868c1
commit a15db14
Showing 1 changed file with 19 additions and 16 deletions.
diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
@@ -64,34 +64,37 @@ fn set_upto_64bits(
 
     if len >= 64 {
         let chunk = unsafe { (data.as_ptr().add(read_byte) as *const u64).read_unaligned() };
-        if read_shift == 0 {
-            if write_shift == 0 {
+        match (read_byte, write_byte) {
+            (0, 0) => {
                 // no shifting necessary
                 let len = 64;
                 let null_count = chunk.count_zeros() as usize;
                 unsafe { write_u64_bytes(write_data, write_byte, chunk) };
                 (null_count, len)
-            } else {
+            }
+            (0, _) => {
                 // only write shifting necessary
                 let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
                 let null_count = len - chunk.count_ones() as usize;
                 unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
                 (null_count, len)
             }
-        } else if write_shift == 0 {
-            // only read shifting necessary
-            let len = 64 - 8; // 56 bits so the next set_upto_64bits call will see write_shift == 0
-            let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
-            let null_count = len - chunk.count_ones() as usize;
-            unsafe { write_u64_bytes(write_data, write_byte, chunk) };
-            (null_count, len)
-        } else {
-            let len = 64 - std::cmp::max(read_shift, write_shift);
-            let chunk = (chunk >> read_shift) << write_shift;
-            let null_count = len - chunk.count_ones() as usize;
-            unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
-            (null_count, len)
+            (_, 0) => {
+                // only read shifting necessary
+                let len = 64 - 8; // 56 bits so the next set_upto_64bits call gets write_shift == 0
+                let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
+                let null_count = len - chunk.count_ones() as usize;
+                unsafe { write_u64_bytes(write_data, write_byte, chunk) };
+                (null_count, len)
+            }
+            (_, _) => {
+                let len = 64 - std::cmp::max(read_shift, write_shift);
+                let chunk = (chunk >> read_shift) << write_shift;
+                let null_count = len - chunk.count_ones() as usize;
+                unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
+                (null_count, len)
+            }
         }
     } else if len == 1 {
         let byte_chunk = (unsafe { data.get_unchecked(read_byte) } >> read_shift) & 1;