From 6e8c86447798dc53890c60a167f6856752c1f8fd Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 16:10:26 -0700
Subject: [PATCH 01/48] bench

---
 arrow-buffer/Cargo.toml          |  6 +++-
 arrow-buffer/benches/bit_mask.rs | 58 ++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 arrow-buffer/benches/bit_mask.rs

diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml
index 8bc33b1874e4..68bfe8ddf732 100644
--- a/arrow-buffer/Cargo.toml
+++ b/arrow-buffer/Cargo.toml
@@ -44,10 +44,14 @@ rand = { version = "0.8", default-features = false, features = ["std", "std_rng"
 
 [build-dependencies]
 
+[[bench]]
+name = "bit_mask"
+harness = false
+
 [[bench]]
 name = "i256"
 harness = false
 
 [[bench]]
 name = "offset"
-harness = false
\ No newline at end of file
+harness = false
diff --git a/arrow-buffer/benches/bit_mask.rs b/arrow-buffer/benches/bit_mask.rs
new file mode 100644
index 000000000000..e941bbda3972
--- /dev/null
+++ b/arrow-buffer/benches/bit_mask.rs
@@ -0,0 +1,58 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_buffer::bit_mask::set_bits;
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("bit_mask");
+
+    for offset_write in (0..8).step_by(5) {
+        for offset_read in (0..8).step_by(5) {
+            for len in [1, 17, 65] {
+                for datum in [0u8, 0xADu8] {
+                    let x = (offset_write, offset_read, len, datum);
+                    group.bench_with_input(
+                        BenchmarkId::new(
+                            "set_bits",
+                            format!(
+                                "offset_write_{}_offset_read_{}_len_{}_datum_{}",
+                                x.0, x.1, x.2, x.3
+                            ),
+                        ),
+                        &x,
+                        |b, &x| {
+                            b.iter(|| {
+                                set_bits(
+                                    black_box(&mut [0u8; 9]),
+                                    black_box(&[x.3; 9]),
+                                    black_box(x.0),
+                                    black_box(x.1),
+                                    black_box(x.2),
+                                )
+                            });
+                        },
+                    );
+                }
+            }
+        }
+    }
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

From a81ba56b8e64e1ef909b1917afab32207b29308a Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 16:10:35 -0700
Subject: [PATCH 02/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 114 +++++++++++++++++++++++-------
 1 file changed, 88 insertions(+), 26 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 8f81cb7d0469..ba117f66178d 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -17,9 +17,6 @@
 
 //! Utils for working with packed bit masks
 
-use crate::bit_chunk_iterator::BitChunks;
-use crate::bit_util::{ceil, get_bit, set_bit};
-
 /// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
 /// bits in `data` in the range `[offset_read..offset_read+len]`
 /// returns the number of `0` bits `data[offset_read..offset_read+len]`
@@ -32,33 +29,98 @@ pub fn set_bits(
 ) -> usize {
     let mut null_count = 0;
 
-    let mut bits_to_align = offset_write % 8;
-    if bits_to_align > 0 {
-        bits_to_align = std::cmp::min(len, 8 - bits_to_align);
+    let mut acc = 0;
+    while len > acc {
+        let (n, l) = set_upto_64bits(
+            write_data,
+            data,
+            offset_write + acc,
+            offset_read + acc,
+            len - acc,
+        );
+        null_count += n;
+        acc += l;
     }
-    let mut write_byte_index = ceil(offset_write + bits_to_align, 8);
-
-    // Set full bytes provided by bit chunk iterator (which iterates in 64 bits at a time)
-    let chunks = BitChunks::new(data, offset_read + bits_to_align, len - bits_to_align);
-    chunks.iter().for_each(|chunk| {
-        null_count += chunk.count_zeros();
-        write_data[write_byte_index..write_byte_index + 8].copy_from_slice(&chunk.to_le_bytes());
-        write_byte_index += 8;
-    });
-
-    // Set individual bits both to align write_data to a byte offset and the remainder bits not covered by the bit chunk iterator
-    let remainder_offset = len - chunks.remainder_len();
-    (0..bits_to_align)
-        .chain(remainder_offset..len)
-        .for_each(|i| {
-            if get_bit(data, offset_read + i) {
-                set_bit(write_data, offset_write + i);
+
+    null_count
+}
+
+#[inline]
+fn set_upto_64bits(
+    write_data: &mut [u8],
+    data: &[u8],
+    offset_write: usize,
+    offset_read: usize,
+    len: usize,
+) -> (usize, usize) {
+    let read_byte = offset_read / 8;
+    let read_shift = offset_read % 8;
+    let write_byte = offset_write / 8;
+    let write_shift = offset_write % 8;
+
+    if len >= 64 {
+        let len = 64 - std::cmp::max(read_shift, write_shift);
+        let chunk = read_bytes_to_u64(data, read_byte, 8);
+        let chunk = chunk >> read_shift;
+        if chunk == 0 {
+            (len, len)
+        } else {
+            let chunk = chunk << write_shift;
+            let null_count = len - chunk.count_ones() as usize;
+            let c = chunk.to_le_bytes();
+            let ptr = write_data.as_ptr() as *mut u8;
+            unsafe {
+                *ptr.add(write_byte) |= c[0];
+                *ptr.add(write_byte + 1) = c[1];
+                *ptr.add(write_byte + 2) = c[2];
+                *ptr.add(write_byte + 3) = c[3];
+                *ptr.add(write_byte + 4) = c[4];
+                *ptr.add(write_byte + 5) = c[5];
+                *ptr.add(write_byte + 6) = c[6];
+                *ptr.add(write_byte + 7) |= c[7];
+            }
+            (null_count, len)
+        }
+    } else {
+        let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
+        let bytes = (len + read_shift).div_ceil(8);
+        let chunk = read_bytes_to_u64(data, read_byte, bytes);
+        let mask = u64::MAX >> (64 - len);
+        let chunk = (chunk >> read_shift) & mask;
+        if chunk == 0 {
+            (len, len)
+        } else {
+            if len == 1 {
+                let ptr = write_data.as_ptr() as *mut u8;
+                unsafe {
+                    *ptr.add(write_byte) |= 1 << write_shift;
+                }
+                (0, 1)
             } else {
-                null_count += 1;
+                let ptr = write_data.as_ptr() as *mut u8;
+                let chunk = chunk << write_shift;
+                let null_count = len - chunk.count_ones() as usize;
+                let bytes = (len + write_shift).div_ceil(8);
+                let c = chunk.to_le_bytes();
+                for i in 0..bytes {
+                    unsafe {
+                        *ptr.add(write_byte + i) |= c[i];
+                    }
+                }
+                (null_count, len)
             }
-        });
+        }
+    }
+}
 
-    null_count as usize
+#[inline]
+fn read_bytes_to_u64(data: &[u8], i: usize, bytes: usize) -> u64 {
+    let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
+    unsafe {
+        let src = data.as_ptr().add(i);
+        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, bytes);
+        tmp.assume_init()
+    }
 }
 
 #[cfg(test)]

From 57634f2ea93ef5a0a41e7dcb3e240b249d01ba56 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 17:00:37 -0700
Subject: [PATCH 03/48] clippy

---
 arrow-buffer/src/util/bit_mask.rs | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index ba117f66178d..1960b354dd87 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -17,6 +17,8 @@
 
 //! Utils for working with packed bit masks
 
+use crate::bit_util::ceil;
+
 /// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
 /// bits in `data` in the range `[offset_read..offset_read+len]`
 /// returns the number of `0` bits `data[offset_read..offset_read+len]`
@@ -83,7 +85,7 @@ fn set_upto_64bits(
         }
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
-        let bytes = (len + read_shift).div_ceil(8);
+        let bytes = ceil(len + read_shift, 8);
         let chunk = read_bytes_to_u64(data, read_byte, bytes);
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;
@@ -100,11 +102,10 @@ fn set_upto_64bits(
                 let ptr = write_data.as_ptr() as *mut u8;
                 let chunk = chunk << write_shift;
                 let null_count = len - chunk.count_ones() as usize;
-                let bytes = (len + write_shift).div_ceil(8);
-                let c = chunk.to_le_bytes();
-                for i in 0..bytes {
+                let bytes = ceil(len + write_shift, 8);
+                for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
                     unsafe {
-                        *ptr.add(write_byte + i) |= c[i];
+                        *ptr.add(write_byte + i) |= c;
                     }
                 }
                 (null_count, len)

From 32ff203d113beccfd6bda91bd708fef7263e6379 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 17:49:11 -0700
Subject: [PATCH 04/48] clippyj

---
 arrow-buffer/src/util/bit_mask.rs | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 1960b354dd87..935d96ee34a3 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -91,25 +91,23 @@ fn set_upto_64bits(
         let chunk = (chunk >> read_shift) & mask;
         if chunk == 0 {
             (len, len)
+        } else if len == 1 {
+            let ptr = write_data.as_ptr() as *mut u8;
+            unsafe {
+                *ptr.add(write_byte) |= 1 << write_shift;
+            }
+            (0, 1)
         } else {
-            if len == 1 {
-                let ptr = write_data.as_ptr() as *mut u8;
+            let ptr = write_data.as_ptr() as *mut u8;
+            let chunk = chunk << write_shift;
+            let null_count = len - chunk.count_ones() as usize;
+            let bytes = ceil(len + write_shift, 8);
+            for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
                 unsafe {
-                    *ptr.add(write_byte) |= 1 << write_shift;
+                    *ptr.add(write_byte + i) |= c;
                 }
-                (0, 1)
-            } else {
-                let ptr = write_data.as_ptr() as *mut u8;
-                let chunk = chunk << write_shift;
-                let null_count = len - chunk.count_ones() as usize;
-                let bytes = ceil(len + write_shift, 8);
-                for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
-                    unsafe {
-                        *ptr.add(write_byte + i) |= c;
-                    }
-                }
-                (null_count, len)
             }
+            (null_count, len)
         }
     }
 }

From f94f31250e933dbbe10da7c0967188a6e00d5cce Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 22:38:30 -0700
Subject: [PATCH 05/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 935d96ee34a3..b0691486407e 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -62,7 +62,7 @@ fn set_upto_64bits(
 
     if len >= 64 {
         let len = 64 - std::cmp::max(read_shift, write_shift);
-        let chunk = read_bytes_to_u64(data, read_byte, 8);
+        let chunk = unsafe { read_bytes_to_u64(data, read_byte, 8) };
         let chunk = chunk >> read_shift;
         if chunk == 0 {
             (len, len)
@@ -86,16 +86,14 @@ fn set_upto_64bits(
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
-        let chunk = read_bytes_to_u64(data, read_byte, bytes);
+        let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;
         if chunk == 0 {
             (len, len)
         } else if len == 1 {
             let ptr = write_data.as_ptr() as *mut u8;
-            unsafe {
-                *ptr.add(write_byte) |= 1 << write_shift;
-            }
+            unsafe { *ptr.add(write_byte) |= 1 << write_shift };
             (0, 1)
         } else {
             let ptr = write_data.as_ptr() as *mut u8;
@@ -103,23 +101,22 @@ fn set_upto_64bits(
             let null_count = len - chunk.count_ones() as usize;
             let bytes = ceil(len + write_shift, 8);
             for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
-                unsafe {
-                    *ptr.add(write_byte + i) |= c;
-                }
+                unsafe { *ptr.add(write_byte + i) |= c };
             }
             (null_count, len)
         }
     }
 }
 
+/// # Safety
+/// The caller must ensure all arguments are within the valid range.
+/// The caller must be aware `8 - count` bytes in the returned value might be uninitialized.
 #[inline]
-fn read_bytes_to_u64(data: &[u8], i: usize, bytes: usize) -> u64 {
+unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
-    unsafe {
-        let src = data.as_ptr().add(i);
-        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, bytes);
-        tmp.assume_init()
-    }
+    let src = data.as_ptr().add(offset);
+    std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
+    tmp.assume_init()
 }
 
 #[cfg(test)]

From e9cd77add154679a5bec7b177eccc05c20bc65af Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 23:28:00 -0700
Subject: [PATCH 06/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 48 ++++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index b0691486407e..7a811e7ec307 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -66,23 +66,39 @@ fn set_upto_64bits(
         let chunk = chunk >> read_shift;
         if chunk == 0 {
             (len, len)
+        } else if write_shift == 0 {
+            let null_count = len - chunk.count_ones() as usize;
+            unsafe {
+                let ptr = write_data.as_ptr().add(write_byte) as *mut u64;
+                ptr.write_unaligned(chunk);
+            }
+            (null_count, len)
         } else {
             let chunk = chunk << write_shift;
             let null_count = len - chunk.count_ones() as usize;
             let c = chunk.to_le_bytes();
-            let ptr = write_data.as_ptr() as *mut u8;
             unsafe {
-                *ptr.add(write_byte) |= c[0];
-                *ptr.add(write_byte + 1) = c[1];
-                *ptr.add(write_byte + 2) = c[2];
-                *ptr.add(write_byte + 3) = c[3];
-                *ptr.add(write_byte + 4) = c[4];
-                *ptr.add(write_byte + 5) = c[5];
-                *ptr.add(write_byte + 6) = c[6];
-                *ptr.add(write_byte + 7) |= c[7];
+                let ptr = write_data.as_mut_ptr().add(write_byte);
+                *ptr |= c[0];
+                *ptr.add(1) = c[1];
+                *ptr.add(2) = c[2];
+                *ptr.add(3) = c[3];
+                *ptr.add(4) = c[4];
+                *ptr.add(5) = c[5];
+                *ptr.add(6) = c[6];
+                *ptr.add(7) |= c[7];
             }
             (null_count, len)
         }
+    } else if len == 1 {
+        let c = unsafe { *data.as_ptr().add(read_byte) } & 1;
+        if c == 0 {
+            (1, 1)
+        } else {
+            let ptr = write_data.as_mut_ptr();
+            unsafe { *ptr.add(write_byte) |= 1 << write_shift };
+            (0, 1)
+        }
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
@@ -91,17 +107,13 @@ fn set_upto_64bits(
         let chunk = (chunk >> read_shift) & mask;
         if chunk == 0 {
             (len, len)
-        } else if len == 1 {
-            let ptr = write_data.as_ptr() as *mut u8;
-            unsafe { *ptr.add(write_byte) |= 1 << write_shift };
-            (0, 1)
         } else {
-            let ptr = write_data.as_ptr() as *mut u8;
             let chunk = chunk << write_shift;
             let null_count = len - chunk.count_ones() as usize;
             let bytes = ceil(len + write_shift, 8);
+            let ptr = unsafe { write_data.as_mut_ptr().add(write_byte) };
             for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
-                unsafe { *ptr.add(write_byte + i) |= c };
+                unsafe { *ptr.add(i) |= c };
             }
             (null_count, len)
         }
@@ -115,8 +127,10 @@ fn set_upto_64bits(
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);
-    std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
-    tmp.assume_init()
+    unsafe {
+        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
+        tmp.assume_init()
+    }
 }
 
 #[cfg(test)]

From 842c2b1d1e97a826ad38e7bb858835c686f42014 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 22 Aug 2024 23:42:33 -0700
Subject: [PATCH 07/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 7a811e7ec307..a38e6ec16b39 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -69,7 +69,7 @@ fn set_upto_64bits(
         } else if write_shift == 0 {
             let null_count = len - chunk.count_ones() as usize;
             unsafe {
-                let ptr = write_data.as_ptr().add(write_byte) as *mut u64;
+                let ptr = write_data.as_mut_ptr().add(write_byte) as *mut u64;
                 ptr.write_unaligned(chunk);
             }
             (null_count, len)
@@ -91,7 +91,7 @@ fn set_upto_64bits(
             (null_count, len)
         }
     } else if len == 1 {
-        let c = unsafe { *data.as_ptr().add(read_byte) } & 1;
+        let c = (unsafe { *data.as_ptr().add(read_byte) } >> read_shift) & 1;
         if c == 0 {
             (1, 1)
         } else {

From 06de184b28baf5f3b338f379b2d09443ae24b5d6 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 23 Aug 2024 00:11:50 -0700
Subject: [PATCH 08/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index a38e6ec16b39..a16517ab1d82 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -76,17 +76,10 @@ fn set_upto_64bits(
         } else {
             let chunk = chunk << write_shift;
             let null_count = len - chunk.count_ones() as usize;
-            let c = chunk.to_le_bytes();
             unsafe {
                 let ptr = write_data.as_mut_ptr().add(write_byte);
-                *ptr |= c[0];
-                *ptr.add(1) = c[1];
-                *ptr.add(2) = c[2];
-                *ptr.add(3) = c[3];
-                *ptr.add(4) = c[4];
-                *ptr.add(5) = c[5];
-                *ptr.add(6) = c[6];
-                *ptr.add(7) |= c[7];
+                let chunk = chunk | (*ptr) as u64;
+                (ptr as *mut u64).write_unaligned(chunk);
             }
             (null_count, len)
         }
@@ -122,11 +115,12 @@ fn set_upto_64bits(
 
 /// # Safety
 /// The caller must ensure all arguments are within the valid range.
-/// The caller must be aware `8 - count` bytes in the returned value might be uninitialized.
+/// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);
+    // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
     unsafe {
         std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
         tmp.assume_init()

From 03b0db86243c2ba724f809b6ca2f6d73d2554500 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 23 Aug 2024 06:11:08 -0700
Subject: [PATCH 09/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index a16517ab1d82..7b09953a7be0 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -73,6 +73,21 @@ fn set_upto_64bits(
                 ptr.write_unaligned(chunk);
             }
             (null_count, len)
+        } else if write_shift > read_shift {
+            let len = 64 - read_shift;
+            let null_count = len - chunk.count_ones() as usize;
+            unsafe {
+                let chunk = chunk << write_shift;
+                let ptr = write_data.as_mut_ptr().add(write_byte);
+                let chunk = chunk | (*ptr) as u64;
+                (ptr as *mut u64).write_unaligned(chunk);
+            }
+            unsafe {
+                let chunk = chunk >> (64 - write_shift);
+                let ptr = write_data.as_mut_ptr().add(write_byte + 8);
+                *ptr |= chunk as u8;
+            }
+            (null_count, len)
         } else {
             let chunk = chunk << write_shift;
             let null_count = len - chunk.count_ones() as usize;

From 7faa5f3d4db5a4d484ff653365e84632004715da Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 23 Aug 2024 13:11:21 -0700
Subject: [PATCH 10/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 68 +++++++++++++++++--------------
 1 file changed, 37 insertions(+), 31 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 7b09953a7be0..8e3f8b08f9e5 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -61,42 +61,48 @@ fn set_upto_64bits(
     let write_shift = offset_write % 8;
 
     if len >= 64 {
-        let len = 64 - std::cmp::max(read_shift, write_shift);
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, 8) };
-        let chunk = chunk >> read_shift;
-        if chunk == 0 {
-            (len, len)
-        } else if write_shift == 0 {
-            let null_count = len - chunk.count_ones() as usize;
-            unsafe {
-                let ptr = write_data.as_mut_ptr().add(write_byte) as *mut u64;
-                ptr.write_unaligned(chunk);
-            }
-            (null_count, len)
-        } else if write_shift > read_shift {
-            let len = 64 - read_shift;
-            let null_count = len - chunk.count_ones() as usize;
-            unsafe {
+        if read_shift == 0 {
+            if write_shift == 0 {
+                let len = 64;
+                let null_count = chunk.count_zeros() as usize;
+                unsafe {
+                    let ptr = write_data.as_mut_ptr().add(write_byte) as *mut u64;
+                    ptr.write_unaligned(chunk);
+                }
+                (null_count, len)
+            } else {
+                let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
-                let ptr = write_data.as_mut_ptr().add(write_byte);
-                let chunk = chunk | (*ptr) as u64;
-                (ptr as *mut u64).write_unaligned(chunk);
+                let null_count = len - chunk.count_ones() as usize;
+                unsafe {
+                    let ptr = write_data.as_mut_ptr().add(write_byte);
+                    let chunk = chunk | (*ptr) as u64;
+                    (ptr as *mut u64).write_unaligned(chunk);
+                }
+                (null_count, len)
             }
-            unsafe {
-                let chunk = chunk >> (64 - write_shift);
-                let ptr = write_data.as_mut_ptr().add(write_byte + 8);
-                *ptr |= chunk as u8;
-            }
-            (null_count, len)
         } else {
-            let chunk = chunk << write_shift;
-            let null_count = len - chunk.count_ones() as usize;
-            unsafe {
-                let ptr = write_data.as_mut_ptr().add(write_byte);
-                let chunk = chunk | (*ptr) as u64;
-                (ptr as *mut u64).write_unaligned(chunk);
+            if write_shift == 0 {
+                let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
+                let chunk = (chunk >> read_shift) | 0x00FFFFFFFFFFFFFF; // 56 bits mask
+                let null_count = len - chunk.count_ones() as usize;
+                unsafe {
+                    let ptr = write_data.as_mut_ptr().add(write_byte) as *mut u64;
+                    ptr.write_unaligned(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 {
+                    let ptr = write_data.as_mut_ptr().add(write_byte);
+                    let chunk = chunk | (*ptr) as u64;
+                    (ptr as *mut u64).write_unaligned(chunk);
+                }
+                (null_count, len)
             }
-            (null_count, len)
         }
     } else if len == 1 {
         let c = (unsafe { *data.as_ptr().add(read_byte) } >> read_shift) & 1;

From e3d812d8dfd911092751cf40b82cf0035c1abe5a Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 23 Aug 2024 13:45:34 -0700
Subject: [PATCH 11/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 71 +++++++++++++++----------------
 1 file changed, 35 insertions(+), 36 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 8e3f8b08f9e5..ef0ef9ad0a95 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -61,25 +61,19 @@ fn set_upto_64bits(
     let write_shift = offset_write % 8;
 
     if len >= 64 {
+        // SAFETY: chunk gets masked when necessary, so it is safe
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, 8) };
         if read_shift == 0 {
             if write_shift == 0 {
                 let len = 64;
                 let null_count = chunk.count_zeros() as usize;
-                unsafe {
-                    let ptr = write_data.as_mut_ptr().add(write_byte) as *mut u64;
-                    ptr.write_unaligned(chunk);
-                }
+                write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
             } else {
                 let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
                 let null_count = len - chunk.count_ones() as usize;
-                unsafe {
-                    let ptr = write_data.as_mut_ptr().add(write_byte);
-                    let chunk = chunk | (*ptr) as u64;
-                    (ptr as *mut u64).write_unaligned(chunk);
-                }
+                or_write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
             }
         } else {
@@ -87,50 +81,36 @@ fn set_upto_64bits(
                 let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
                 let chunk = (chunk >> read_shift) | 0x00FFFFFFFFFFFFFF; // 56 bits mask
                 let null_count = len - chunk.count_ones() as usize;
-                unsafe {
-                    let ptr = write_data.as_mut_ptr().add(write_byte) as *mut u64;
-                    ptr.write_unaligned(chunk);
-                }
+                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 {
-                    let ptr = write_data.as_mut_ptr().add(write_byte);
-                    let chunk = chunk | (*ptr) as u64;
-                    (ptr as *mut u64).write_unaligned(chunk);
-                }
+                or_write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
             }
         }
     } else if len == 1 {
         let c = (unsafe { *data.as_ptr().add(read_byte) } >> read_shift) & 1;
-        if c == 0 {
-            (1, 1)
-        } else {
-            let ptr = write_data.as_mut_ptr();
-            unsafe { *ptr.add(write_byte) |= 1 << write_shift };
-            (0, 1)
-        }
+        let ptr = write_data.as_mut_ptr();
+        unsafe { *ptr.add(write_byte) |= c << write_shift };
+        ((c ^ 1) as usize, 1)
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
+        // SAFETY: chunk gets masked, so it is safe
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;
-        if chunk == 0 {
-            (len, len)
-        } else {
-            let chunk = chunk << write_shift;
-            let null_count = len - chunk.count_ones() as usize;
-            let bytes = ceil(len + write_shift, 8);
-            let ptr = unsafe { write_data.as_mut_ptr().add(write_byte) };
-            for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
-                unsafe { *ptr.add(i) |= c };
-            }
-            (null_count, len)
+        let chunk = chunk << write_shift;
+        let null_count = len - chunk.count_ones() as usize;
+        let bytes = ceil(len + write_shift, 8);
+        let ptr = unsafe { write_data.as_mut_ptr().add(write_byte) };
+        for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
+            unsafe { *ptr.add(i) |= c };
         }
+        (null_count, len)
     }
 }
 
@@ -148,6 +128,25 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     }
 }
 
+#[inline]
+fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
+    // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range
+    unsafe {
+        let ptr = data.as_mut_ptr().add(offset) as *mut u64;
+        ptr.write_unaligned(chunk);
+    }
+}
+
+#[inline]
+fn or_write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
+    // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range
+    unsafe {
+        let ptr = data.as_mut_ptr().add(offset);
+        let chunk = chunk | (*ptr) as u64;
+        (ptr as *mut u64).write_unaligned(chunk);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

From 13dec63aeaa98c48fd0b4a5cba949b549e048532 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 23 Aug 2024 14:05:06 -0700
Subject: [PATCH 12/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index ef0ef9ad0a95..ad28d3f42c69 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -76,20 +76,18 @@ fn set_upto_64bits(
                 or_write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
             }
+        } else if write_shift == 0 {
+            let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
+            let chunk = (chunk >> read_shift) | 0x00FFFFFFFFFFFFFF; // 56 bits mask
+            let null_count = len - chunk.count_ones() as usize;
+            write_u64_bytes(write_data, write_byte, chunk);
+            (null_count, len)
         } else {
-            if write_shift == 0 {
-                let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
-                let chunk = (chunk >> read_shift) | 0x00FFFFFFFFFFFFFF; // 56 bits mask
-                let null_count = len - chunk.count_ones() as usize;
-                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;
-                or_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;
+            or_write_u64_bytes(write_data, write_byte, chunk);
+            (null_count, len)
         }
     } else if len == 1 {
         let c = (unsafe { *data.as_ptr().add(read_byte) } >> read_shift) & 1;

From 68cdaf24b1b3be772fb48217634e97105875bc09 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 23 Aug 2024 14:15:41 -0700
Subject: [PATCH 13/48] fix: Optimize set_bits

---
 arrow-buffer/src/util/bit_mask.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index ad28d3f42c69..f1fa232fb7c3 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -78,7 +78,7 @@ fn set_upto_64bits(
             }
         } else if write_shift == 0 {
             let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
-            let chunk = (chunk >> read_shift) | 0x00FFFFFFFFFFFFFF; // 56 bits mask
+            let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
             let null_count = len - chunk.count_ones() as usize;
             write_u64_bytes(write_data, write_byte, chunk);
             (null_count, len)

From 1e9de38de8812673a0786f72de09b171121b2c4f Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 00:23:25 -0700
Subject: [PATCH 14/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index f1fa232fb7c3..fd9be7bf5545 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -116,6 +116,7 @@ fn set_upto_64bits(
 /// The caller must ensure all arguments are within the valid range.
 /// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
+#[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);

From f1e1bbd2ee8cb364f6f46c50abc41c77b58d3595 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 00:25:29 -0700
Subject: [PATCH 15/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index fd9be7bf5545..0b178e2df95a 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -116,13 +116,13 @@ fn set_upto_64bits(
 /// The caller must ensure all arguments are within the valid range.
 /// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
-#[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);
     // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
     unsafe {
         std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
+        #[cfg(not(miri))]
         tmp.assume_init()
     }
 }

From f2946638dc0b93eda8a4a86b1ceedb3802bf1de4 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 00:27:02 -0700
Subject: [PATCH 16/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 0b178e2df95a..cc2f52f43eb0 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -98,6 +98,7 @@ fn set_upto_64bits(
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
         // SAFETY: chunk gets masked, so it is safe
+        #[cfg(not(miri))]
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;
@@ -122,7 +123,6 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
     unsafe {
         std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
-        #[cfg(not(miri))]
         tmp.assume_init()
     }
 }

From 39719c44b094ce27f339aa3a7d92ac8bbf51bb45 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 00:37:43 -0700
Subject: [PATCH 17/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index cc2f52f43eb0..a9d9442fa8d5 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -98,7 +98,6 @@ fn set_upto_64bits(
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
         // SAFETY: chunk gets masked, so it is safe
-        #[cfg(not(miri))]
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;
@@ -123,7 +122,11 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
     unsafe {
         std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
-        tmp.assume_init()
+        let mut res = 0;
+        #[cfg(not(miri))]
+        res = tmp.assume_init();
+        #[cfg(miri)]
+        res
     }
 }
 

From 9fbb87d2c166298de832e7bd14daba56a6f8b52f Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 00:57:06 -0700
Subject: [PATCH 18/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index a9d9442fa8d5..f39ce7c37a20 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -17,6 +17,7 @@
 
 //! Utils for working with packed bit masks
 
+#![cfg(not(miri))]
 use crate::bit_util::ceil;
 
 /// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
@@ -122,11 +123,7 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
     unsafe {
         std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
-        let mut res = 0;
-        #[cfg(not(miri))]
-        res = tmp.assume_init();
-        #[cfg(miri)]
-        res
+        tmp.assume_init()
     }
 }
 

From 74b9d80ab4171345edd7339234e05c80b858b70f Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 01:11:41 -0700
Subject: [PATCH 19/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index f39ce7c37a20..cc79c77f1747 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -17,7 +17,6 @@
 
 //! Utils for working with packed bit masks
 
-#![cfg(not(miri))]
 use crate::bit_util::ceil;
 
 /// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
@@ -33,6 +32,7 @@ pub fn set_bits(
     let mut null_count = 0;
 
     let mut acc = 0;
+    #[cfg(not(miri))]
     while len > acc {
         let (n, l) = set_upto_64bits(
             write_data,

From 25c309ec049176122738808f81839c4d611a68a6 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 01:32:46 -0700
Subject: [PATCH 20/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index cc79c77f1747..d786aa74da88 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -32,7 +32,6 @@ pub fn set_bits(
     let mut null_count = 0;
 
     let mut acc = 0;
-    #[cfg(not(miri))]
     while len > acc {
         let (n, l) = set_upto_64bits(
             write_data,
@@ -117,6 +116,7 @@ fn set_upto_64bits(
 /// The caller must ensure all arguments are within the valid range.
 /// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
+#[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);
@@ -127,6 +127,17 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     }
 }
 
+#[cfg(miri)]
+unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
+    let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
+    let src = data.as_ptr().add(offset);
+    // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
+    unsafe {
+        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
+        tmp.assume_init_read()
+    }
+}
+
 #[inline]
 fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
     // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range

From 79053304f225c4ef7f4ba13731f4c43b313fb19c Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 01:40:33 -0700
Subject: [PATCH 21/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index d786aa74da88..1fe75cfbd8c3 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -129,13 +129,8 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
 
 #[cfg(miri)]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
-    let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);
-    // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
-    unsafe {
-        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
-        tmp.assume_init_read()
-    }
+    u64::from_le_bytes(src.to_le_bytes())
 }
 
 #[inline]

From 6dd977111e2121318299623cf5fb98253c294f5b Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 02:01:31 -0700
Subject: [PATCH 22/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 1fe75cfbd8c3..d19538198623 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -129,8 +129,9 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
 
 #[cfg(miri)]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
-    let src = data.as_ptr().add(offset);
-    u64::from_le_bytes(src.to_le_bytes())
+    let mut arr = [0u8; 8];
+    arr[offset..(offset + count)].copy_from_slice(&data[offset..(offset + count)]);
+    u64::from_le_bytes(arr)
 }
 
 #[inline]

From 0e956cc9b3a1a04cf7637cd10a25fbcb508c63b4 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 02:13:01 -0700
Subject: [PATCH 23/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index d19538198623..299c8dedf678 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -131,7 +131,7 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut arr = [0u8; 8];
     arr[offset..(offset + count)].copy_from_slice(&data[offset..(offset + count)]);
-    u64::from_le_bytes(arr)
+    u64::from_be_bytes(arr)
 }
 
 #[inline]

From 272ecbbcb4490a0b551fa7980882a44a473c1bb5 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Sat, 24 Aug 2024 02:24:30 -0700
Subject: [PATCH 24/48] miri

---
 arrow-buffer/src/util/bit_mask.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 299c8dedf678..4e959956c104 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -130,8 +130,8 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
 #[cfg(miri)]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     let mut arr = [0u8; 8];
-    arr[offset..(offset + count)].copy_from_slice(&data[offset..(offset + count)]);
-    u64::from_be_bytes(arr)
+    arr[0..count].copy_from_slice(&data[offset..(offset + count)]);
+    u64::from_le_bytes(arr)
 }
 
 #[inline]

From 08ebf20fe8549537d98632133376b6e526f6be79 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Tue, 3 Sep 2024 14:42:16 -0700
Subject: [PATCH 25/48] address review comments

---
 arrow-buffer/Cargo.toml           |  6 +---
 arrow-buffer/benches/bit_mask.rs  | 58 -------------------------------
 arrow-buffer/src/util/bit_mask.rs | 11 +++---
 3 files changed, 7 insertions(+), 68 deletions(-)
 delete mode 100644 arrow-buffer/benches/bit_mask.rs

diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml
index 68bfe8ddf732..8bc33b1874e4 100644
--- a/arrow-buffer/Cargo.toml
+++ b/arrow-buffer/Cargo.toml
@@ -44,14 +44,10 @@ rand = { version = "0.8", default-features = false, features = ["std", "std_rng"
 
 [build-dependencies]
 
-[[bench]]
-name = "bit_mask"
-harness = false
-
 [[bench]]
 name = "i256"
 harness = false
 
 [[bench]]
 name = "offset"
-harness = false
+harness = false
\ No newline at end of file
diff --git a/arrow-buffer/benches/bit_mask.rs b/arrow-buffer/benches/bit_mask.rs
deleted file mode 100644
index e941bbda3972..000000000000
--- a/arrow-buffer/benches/bit_mask.rs
+++ /dev/null
@@ -1,58 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements.  See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership.  The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License.  You may obtain a copy of the License at
-//
-//   http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied.  See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use arrow_buffer::bit_mask::set_bits;
-use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
-
-fn criterion_benchmark(c: &mut Criterion) {
-    let mut group = c.benchmark_group("bit_mask");
-
-    for offset_write in (0..8).step_by(5) {
-        for offset_read in (0..8).step_by(5) {
-            for len in [1, 17, 65] {
-                for datum in [0u8, 0xADu8] {
-                    let x = (offset_write, offset_read, len, datum);
-                    group.bench_with_input(
-                        BenchmarkId::new(
-                            "set_bits",
-                            format!(
-                                "offset_write_{}_offset_read_{}_len_{}_datum_{}",
-                                x.0, x.1, x.2, x.3
-                            ),
-                        ),
-                        &x,
-                        |b, &x| {
-                            b.iter(|| {
-                                set_bits(
-                                    black_box(&mut [0u8; 9]),
-                                    black_box(&[x.3; 9]),
-                                    black_box(x.0),
-                                    black_box(x.1),
-                                    black_box(x.2),
-                                )
-                            });
-                        },
-                    );
-                }
-            }
-        }
-    }
-    group.finish();
-}
-
-criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);
diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 4e959956c104..36b8f74c538d 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -33,7 +33,7 @@ pub fn set_bits(
 
     let mut acc = 0;
     while len > acc {
-        let (n, l) = set_upto_64bits(
+        let (n, len_set) = set_upto_64bits(
             write_data,
             data,
             offset_write + acc,
@@ -41,7 +41,7 @@ pub fn set_bits(
             len - acc,
         );
         null_count += n;
-        acc += l;
+        acc += len_set;
     }
 
     null_count
@@ -64,19 +64,19 @@ fn set_upto_64bits(
         // SAFETY: chunk gets masked when necessary, so it is safe
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, 8) };
         if read_shift == 0 {
-            if write_shift == 0 {
+            if write_shift == 0 { // no shifting necessary
                 let len = 64;
                 let null_count = chunk.count_zeros() as usize;
                 write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
-            } else {
+            } else { // only write shifting necessary
                 let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
                 let null_count = len - chunk.count_ones() as usize;
                 or_write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
             }
-        } else if write_shift == 0 {
+        } else if write_shift == 0 { // only read shifting necessary
             let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
             let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
             let null_count = len - chunk.count_ones() as usize;
@@ -118,6 +118,7 @@ fn set_upto_64bits(
 #[inline]
 #[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
+    debug_assert!(count <= 8);
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
     let src = data.as_ptr().add(offset);
     // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.

From d751a7feb97475a6ea3498b663b9ba22cfbd0ad3 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Tue, 3 Sep 2024 14:58:26 -0700
Subject: [PATCH 26/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 46 +++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 36b8f74c538d..cc827eb5e939 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -64,19 +64,22 @@ fn set_upto_64bits(
         // SAFETY: chunk gets masked when necessary, so it is safe
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, 8) };
         if read_shift == 0 {
-            if write_shift == 0 { // no shifting necessary
+            if write_shift == 0 {
+                // no shifting necessary
                 let len = 64;
                 let null_count = chunk.count_zeros() as usize;
                 write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
-            } else { // only write shifting necessary
+            } else {
+                // only write shifting necessary
                 let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
                 let null_count = len - chunk.count_ones() as usize;
                 or_write_u64_bytes(write_data, write_byte, chunk);
                 (null_count, len)
             }
-        } else if write_shift == 0 { // only read shifting necessary
+        } else if write_shift == 0 {
+            // only read shifting necessary
             let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
             let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
             let null_count = len - chunk.count_ones() as usize;
@@ -278,4 +281,41 @@ mod tests {
         assert_eq!(destination, expected_data);
         assert_eq!(result, expected_null_count);
     }
+
+    #[test]
+    fn test_set_upto_64bits() {
+        // len >= 64
+        let write_data: &mut [u8] = &mut [
+            0b00000000, 0b00000000, 0b00000000, 0b00000000, 0b00000000, 0b00000000, 0b00000000,
+            0b00000000,
+        ];
+        let data: &[u8] = &[
+            0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001,
+            0b00000001,
+        ];
+        let offset_write = 1;
+        let offset_read = 0;
+        let len = 64;
+        let (n, len_set) = set_upto_64bits(write_data, data, offset_write, offset_read, len);
+        assert_eq!(n, 55);
+        assert_eq!(len_set, 63);
+        assert_eq!(
+            write_data,
+            &[
+                0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010,
+                0b00000010
+            ]
+        );
+
+        // len = 1
+        let write_data: &mut [u8] = &mut [0b00000000];
+        let data: &[u8] = &[0b00000001];
+        let offset_write = 1;
+        let offset_read = 0;
+        let len = 1;
+        let (n, len_set) = set_upto_64bits(write_data, data, offset_write, offset_read, len);
+        assert_eq!(n, 0);
+        assert_eq!(len_set, 1);
+        assert_eq!(write_data, &[0b00000010]);
+    }
 }

From ef2864fe15d2c856c05eae70693d68eb2ae00fa8 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Wed, 4 Sep 2024 15:31:48 -0700
Subject: [PATCH 27/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index cc827eb5e939..44871fb5004d 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -100,10 +100,12 @@ fn set_upto_64bits(
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
-        // SAFETY: chunk gets masked, so it is safe
-        let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
-        let chunk = (chunk >> read_shift) & mask;
+        // SAFETY: chunk gets masked, so it is safe
+        let chunk = unsafe {
+            let chunk = read_bytes_to_u64(data, read_byte, bytes);
+            (chunk >> read_shift) & mask
+        };
         let chunk = chunk << write_shift;
         let null_count = len - chunk.count_ones() as usize;
         let bytes = ceil(len + write_shift, 8);
@@ -119,7 +121,6 @@ fn set_upto_64bits(
 /// The caller must ensure all arguments are within the valid range.
 /// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
-#[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     debug_assert!(count <= 8);
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
@@ -131,13 +132,6 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     }
 }
 
-#[cfg(miri)]
-unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
-    let mut arr = [0u8; 8];
-    arr[0..count].copy_from_slice(&data[offset..(offset + count)]);
-    u64::from_le_bytes(arr)
-}
-
 #[inline]
 fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
     // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range

From e69cf9a4dfef5795f77db81ed6a36278081c5b16 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Wed, 4 Sep 2024 15:36:00 -0700
Subject: [PATCH 28/48] Revert "address review comments"

This reverts commit ef2864fe15d2c856c05eae70693d68eb2ae00fa8.
---
 arrow-buffer/src/util/bit_mask.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 44871fb5004d..cc827eb5e939 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -100,12 +100,10 @@ fn set_upto_64bits(
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
-        let mask = u64::MAX >> (64 - len);
         // SAFETY: chunk gets masked, so it is safe
-        let chunk = unsafe {
-            let chunk = read_bytes_to_u64(data, read_byte, bytes);
-            (chunk >> read_shift) & mask
-        };
+        let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
+        let mask = u64::MAX >> (64 - len);
+        let chunk = (chunk >> read_shift) & mask;
         let chunk = chunk << write_shift;
         let null_count = len - chunk.count_ones() as usize;
         let bytes = ceil(len + write_shift, 8);
@@ -121,6 +119,7 @@ fn set_upto_64bits(
 /// The caller must ensure all arguments are within the valid range.
 /// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
+#[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     debug_assert!(count <= 8);
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
@@ -132,6 +131,13 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     }
 }
 
+#[cfg(miri)]
+unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
+    let mut arr = [0u8; 8];
+    arr[0..count].copy_from_slice(&data[offset..(offset + count)]);
+    u64::from_le_bytes(arr)
+}
+
 #[inline]
 fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
     // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range

From b5f8bcacbea446188bd478b585e177ee83ab9659 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 5 Sep 2024 15:24:22 -0700
Subject: [PATCH 29/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index cc827eb5e939..0dc450b30be5 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -101,7 +101,12 @@ fn set_upto_64bits(
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
         // SAFETY: chunk gets masked, so it is safe
-        let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
+        let chunk = unsafe {
+            let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
+            let src = data.as_ptr().add(read_byte);
+            std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, bytes);
+            tmp.assume_init()
+        };
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;
         let chunk = chunk << write_shift;
@@ -119,7 +124,6 @@ fn set_upto_64bits(
 /// The caller must ensure all arguments are within the valid range.
 /// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
-#[cfg(not(miri))]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     debug_assert!(count <= 8);
     let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
@@ -131,13 +135,6 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     }
 }
 
-#[cfg(miri)]
-unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
-    let mut arr = [0u8; 8];
-    arr[0..count].copy_from_slice(&data[offset..(offset + count)]);
-    u64::from_le_bytes(arr)
-}
-
 #[inline]
 fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
     // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range

From 9c15417e8c4561054db1bdbdc0d073a0f663acea Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 5 Sep 2024 15:32:44 -0700
Subject: [PATCH 30/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 0dc450b30be5..30ec4df2ee49 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -100,15 +100,15 @@ fn set_upto_64bits(
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
+        let mask = u64::MAX >> (64 - len);
         // SAFETY: chunk gets masked, so it is safe
         let chunk = unsafe {
             let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
             let src = data.as_ptr().add(read_byte);
             std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, bytes);
-            tmp.assume_init()
+            let chunk = tmp.assume_init();
+            (chunk >> read_shift) & mask
         };
-        let mask = u64::MAX >> (64 - len);
-        let chunk = (chunk >> read_shift) & mask;
         let chunk = chunk << write_shift;
         let null_count = len - chunk.count_ones() as usize;
         let bytes = ceil(len + write_shift, 8);

From 533381a99d42b2a8df26f97df88698a082769e6a Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 5 Sep 2024 16:15:50 -0700
Subject: [PATCH 31/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 30ec4df2ee49..3e3e86131bde 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -100,15 +100,10 @@ fn set_upto_64bits(
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
-        let mask = u64::MAX >> (64 - len);
         // SAFETY: chunk gets masked, so it is safe
-        let chunk = unsafe {
-            let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
-            let src = data.as_ptr().add(read_byte);
-            std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, bytes);
-            let chunk = tmp.assume_init();
-            (chunk >> read_shift) & mask
-        };
+        let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
+        let mask = u64::MAX >> (64 - len);
+        let chunk = (chunk >> read_shift) & mask;
         let chunk = chunk << write_shift;
         let null_count = len - chunk.count_ones() as usize;
         let bytes = ceil(len + write_shift, 8);
@@ -122,16 +117,14 @@ fn set_upto_64bits(
 
 /// # Safety
 /// The caller must ensure all arguments are within the valid range.
-/// The caller must be aware `8 - count` bytes in the returned value are uninitialized.
 #[inline]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     debug_assert!(count <= 8);
-    let mut tmp = std::mem::MaybeUninit::<u64>::uninit();
+    let tmp = 0u64;
     let src = data.as_ptr().add(offset);
-    // SAFETY: the caller must not use the uninitialized `8 - count` bytes in the returned value.
     unsafe {
-        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
-        tmp.assume_init()
+        std::ptr::copy_nonoverlapping(src, &tmp as *const u64 as *mut u8, count);
+        tmp
     }
 }
 

From dca9ab842041123deb83d7a21dfb876a83072ca5 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 5 Sep 2024 16:21:53 -0700
Subject: [PATCH 32/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 3e3e86131bde..bb3b6496dc5c 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -120,11 +120,11 @@ fn set_upto_64bits(
 #[inline]
 unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     debug_assert!(count <= 8);
-    let tmp = 0u64;
+    let mut tmp = std::mem::MaybeUninit::<u64>::new(0);
     let src = data.as_ptr().add(offset);
     unsafe {
-        std::ptr::copy_nonoverlapping(src, &tmp as *const u64 as *mut u8, count);
-        tmp
+        std::ptr::copy_nonoverlapping(src, tmp.as_mut_ptr() as *mut u8, count);
+        tmp.assume_init()
     }
 }
 

From 7f3c3fb5483924c393a87fb8338bdce3dd16aff4 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 5 Sep 2024 16:36:15 -0700
Subject: [PATCH 33/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index bb3b6496dc5c..db45f0dd790b 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -61,8 +61,7 @@ fn set_upto_64bits(
     let write_shift = offset_write % 8;
 
     if len >= 64 {
-        // SAFETY: chunk gets masked when necessary, so it is safe
-        let chunk = unsafe { read_bytes_to_u64(data, read_byte, 8) };
+        let chunk = unsafe { (data.as_ptr().add(read_byte) as *const u64).read_unaligned() };
         if read_shift == 0 {
             if write_shift == 0 {
                 // no shifting necessary
@@ -100,7 +99,6 @@ fn set_upto_64bits(
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
-        // SAFETY: chunk gets masked, so it is safe
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
         let chunk = (chunk >> read_shift) & mask;

From 6ccedd2f08ca7e364a886fe2a7740aaf273af8ea Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 10:40:41 -0700
Subject: [PATCH 34/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 32 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index db45f0dd790b..9898d2337c40 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -67,14 +67,14 @@ fn set_upto_64bits(
                 // no shifting necessary
                 let len = 64;
                 let null_count = chunk.count_zeros() as usize;
-                write_u64_bytes(write_data, write_byte, chunk);
+                unsafe { write_u64_bytes(write_data, write_byte, chunk) };
                 (null_count, len)
             } else {
                 // only write shifting necessary
                 let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
                 let null_count = len - chunk.count_ones() as usize;
-                or_write_u64_bytes(write_data, write_byte, chunk);
+                unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
                 (null_count, len)
             }
         } else if write_shift == 0 {
@@ -82,13 +82,13 @@ fn set_upto_64bits(
             let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
             let chunk = (chunk >> read_shift) & 0x00FFFFFFFFFFFFFF; // 56 bits mask
             let null_count = len - chunk.count_ones() as usize;
-            write_u64_bytes(write_data, write_byte, chunk);
+            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;
-            or_write_u64_bytes(write_data, write_byte, chunk);
+            unsafe { or_write_u64_bytes(write_data, write_byte, chunk) };
             (null_count, len)
         }
     } else if len == 1 {
@@ -126,23 +126,21 @@ unsafe fn read_bytes_to_u64(data: &[u8], offset: usize, count: usize) -> u64 {
     }
 }
 
+/// # Safety
+/// The caller must ensure `data` has `offset..(offset + 8)` range
 #[inline]
-fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
-    // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range
-    unsafe {
-        let ptr = data.as_mut_ptr().add(offset) as *mut u64;
-        ptr.write_unaligned(chunk);
-    }
+unsafe fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
+    let ptr = data.as_mut_ptr().add(offset) as *mut u64;
+    ptr.write_unaligned(chunk);
 }
 
+/// # Safety
+/// The caller must ensure `data` has `offset..(offset + 8)` range
 #[inline]
-fn or_write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
-    // SAFETY: the caller must ensure `data` has `offset..(offset + 8)` range
-    unsafe {
-        let ptr = data.as_mut_ptr().add(offset);
-        let chunk = chunk | (*ptr) as u64;
-        (ptr as *mut u64).write_unaligned(chunk);
-    }
+unsafe fn or_write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
+    let ptr = data.as_mut_ptr().add(offset);
+    let chunk = chunk | (*ptr) as u64;
+    (ptr as *mut u64).write_unaligned(chunk);
 }
 
 #[cfg(test)]

From ff2f3caa932c20442dcefe72d38a5e8877dfd7c9 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 10:47:51 -0700
Subject: [PATCH 35/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 9898d2337c40..8655f750aad5 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -92,10 +92,10 @@ fn set_upto_64bits(
             (null_count, len)
         }
     } else if len == 1 {
-        let c = (unsafe { *data.as_ptr().add(read_byte) } >> read_shift) & 1;
+        let byte_chunk = (unsafe { data.get_unchecked(read_byte) } >> read_shift) & 1;
         let ptr = write_data.as_mut_ptr();
-        unsafe { *ptr.add(write_byte) |= c << write_shift };
-        ((c ^ 1) as usize, 1)
+        unsafe { *ptr.add(write_byte) |= byte_chunk << write_shift };
+        ((byte_chunk ^ 1) as usize, 1)
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);

From fb46cb00fb1ba515003e46ce8e0ef9a985d8ec89 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:04:42 -0700
Subject: [PATCH 36/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 8655f750aad5..9b9b2be0b6c8 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -134,6 +134,9 @@ unsafe fn write_u64_bytes(data: &mut [u8], offset: usize, chunk: u64) {
     ptr.write_unaligned(chunk);
 }
 
+/// Similar to `write_u64_bytes`, but this method ORs the offset addressed `data` and `chunk`
+/// instead of overwriting
+///
 /// # Safety
 /// The caller must ensure `data` has `offset..(offset + 8)` range
 #[inline]

From 3fd5e3e6e77ebd5ac4b14f70ec0a58045532e993 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:15:05 -0700
Subject: [PATCH 37/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 9b9b2be0b6c8..cc9b32189b6c 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -79,7 +79,7 @@ fn set_upto_64bits(
             }
         } else if write_shift == 0 {
             // only read shifting necessary
-            let len = 64 - 8; // 56 bits so that write_shift == 0 for the next iteration
+            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) };

From be3076edbaf37347b135ec352e0ada2731495315 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:31:15 -0700
Subject: [PATCH 38/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index cc9b32189b6c..8e073a849857 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -55,6 +55,8 @@ fn set_upto_64bits(
     offset_read: usize,
     len: usize,
 ) -> (usize, usize) {
+    debug_assert!(offset_read + len <= data.len());
+    debug_assert!(offset_write + len <= write_data.len());
     let read_byte = offset_read / 8;
     let read_shift = offset_read % 8;
     let write_byte = offset_write / 8;

From 58868c1229b6303b42fa609e0249589a1be7bff1 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:39:43 -0700
Subject: [PATCH 39/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 8e073a849857..e376d42731b2 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -55,8 +55,8 @@ fn set_upto_64bits(
     offset_read: usize,
     len: usize,
 ) -> (usize, usize) {
-    debug_assert!(offset_read + len <= data.len());
-    debug_assert!(offset_write + len <= write_data.len());
+    debug_assert!(offset_read * 8 + len <= data.len() * 8);
+    debug_assert!(offset_write * 8 + len <= write_data.len() * 8);
     let read_byte = offset_read / 8;
     let read_shift = offset_read % 8;
     let write_byte = offset_write / 8;

From a15db144effdfdae7dad4d93c8fb6eb93216dab0 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:45:18 -0700
Subject: [PATCH 40/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 35 +++++++++++++++++--------------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index e376d42731b2..5edfc74ddbdd 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -64,14 +64,15 @@ 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;
@@ -79,19 +80,21 @@ fn set_upto_64bits(
                 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;

From fefafa793f5537b7daf083e39810dad5656725fa Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:46:55 -0700
Subject: [PATCH 41/48] Revert "address review comments"

This reverts commit a15db144effdfdae7dad4d93c8fb6eb93216dab0.
---
 arrow-buffer/src/util/bit_mask.rs | 35 ++++++++++++++-----------------
 1 file changed, 16 insertions(+), 19 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 5edfc74ddbdd..e376d42731b2 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -64,15 +64,14 @@ fn set_upto_64bits(
 
     if len >= 64 {
         let chunk = unsafe { (data.as_ptr().add(read_byte) as *const u64).read_unaligned() };
-        match (read_byte, write_byte) {
-            (0, 0) => {
+        if read_shift == 0 {
+            if write_shift == 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)
-            }
-            (0, _) => {
+            } else {
                 // only write shifting necessary
                 let len = 64 - write_shift;
                 let chunk = chunk << write_shift;
@@ -80,21 +79,19 @@ fn set_upto_64bits(
                 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 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)
         }
     } else if len == 1 {
         let byte_chunk = (unsafe { data.get_unchecked(read_byte) } >> read_shift) & 1;

From d8c3f08ca04ee24eabf9e45ace5d88b63bb6e54b Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 11:54:50 -0700
Subject: [PATCH 42/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index e376d42731b2..d48c8794262d 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -55,8 +55,8 @@ fn set_upto_64bits(
     offset_read: usize,
     len: usize,
 ) -> (usize, usize) {
-    debug_assert!(offset_read * 8 + len <= data.len() * 8);
-    debug_assert!(offset_write * 8 + len <= write_data.len() * 8);
+    debug_assert!(offset_read + len <= data.len() * 8);
+    debug_assert!(offset_write + len <= write_data.len() * 8);
     let read_byte = offset_read / 8;
     let read_shift = offset_read % 8;
     let write_byte = offset_write / 8;

From cc5ec2b2e3600ee4930818a2cd0e8aee56cd27a1 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 6 Sep 2024 12:03:18 -0700
Subject: [PATCH 43/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index d48c8794262d..3d8dc2b6ab0b 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -276,10 +276,7 @@ mod tests {
     #[test]
     fn test_set_upto_64bits() {
         // len >= 64
-        let write_data: &mut [u8] = &mut [
-            0b00000000, 0b00000000, 0b00000000, 0b00000000, 0b00000000, 0b00000000, 0b00000000,
-            0b00000000,
-        ];
+        let write_data: &mut [u8] = &mut [0; 9];
         let data: &[u8] = &[
             0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001,
             0b00000001,
@@ -294,7 +291,7 @@ mod tests {
             write_data,
             &[
                 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010, 0b00000010,
-                0b00000010
+                0b00000010, 0b00000000
             ]
         );
 

From 4c39dc8eac61f06e6f46843617ef5d05a0ea9b72 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Mon, 9 Sep 2024 15:00:44 -0700
Subject: [PATCH 44/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 3d8dc2b6ab0b..d2a7463668af 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -29,17 +29,23 @@ pub fn set_bits(
     offset_read: usize,
     len: usize,
 ) -> usize {
+    assert!(offset_write + len <= write_data.len() * 8);
+    assert!(offset_read + len <= data.len() * 8);
     let mut null_count = 0;
-
     let mut acc = 0;
     while len > acc {
-        let (n, len_set) = set_upto_64bits(
-            write_data,
-            data,
-            offset_write + acc,
-            offset_read + acc,
-            len - acc,
-        );
+        // SAFETY: the arguments to `set_upto_64bits` are within the valid range because
+        // (offset_write + acc) + (len - acc) == offset_write + len <= write_data.len() * 8
+        // (offset_read + acc) + (len - acc) == offset_read + len <= data.len() * 8
+        let (n, len_set) = unsafe {
+            set_upto_64bits(
+                write_data,
+                data,
+                offset_write + acc,
+                offset_read + acc,
+                len - acc,
+            )
+        };
         null_count += n;
         acc += len_set;
     }
@@ -47,16 +53,16 @@ pub fn set_bits(
     null_count
 }
 
+/// # Safety
+/// The caller must ensure all arguments are within the valid range.
 #[inline]
-fn set_upto_64bits(
+unsafe fn set_upto_64bits(
     write_data: &mut [u8],
     data: &[u8],
     offset_write: usize,
     offset_read: usize,
     len: usize,
 ) -> (usize, usize) {
-    debug_assert!(offset_read + len <= data.len() * 8);
-    debug_assert!(offset_write + len <= write_data.len() * 8);
     let read_byte = offset_read / 8;
     let read_shift = offset_read % 8;
     let write_byte = offset_write / 8;

From f4789bee1d15d75e96e68dd8d39485535ca341eb Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Mon, 9 Sep 2024 15:19:01 -0700
Subject: [PATCH 45/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index d2a7463668af..c0f8c7c38071 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -290,7 +290,8 @@ mod tests {
         let offset_write = 1;
         let offset_read = 0;
         let len = 64;
-        let (n, len_set) = set_upto_64bits(write_data, data, offset_write, offset_read, len);
+        let (n, len_set) =
+            unsafe { set_upto_64bits(write_data, data, offset_write, offset_read, len) };
         assert_eq!(n, 55);
         assert_eq!(len_set, 63);
         assert_eq!(
@@ -307,7 +308,8 @@ mod tests {
         let offset_write = 1;
         let offset_read = 0;
         let len = 1;
-        let (n, len_set) = set_upto_64bits(write_data, data, offset_write, offset_read, len);
+        let (n, len_set) =
+            unsafe { set_upto_64bits(write_data, data, offset_write, offset_read, len) };
         assert_eq!(n, 0);
         assert_eq!(len_set, 1);
         assert_eq!(write_data, &[0b00000010]);

From 59fd805ee9e0a4c267a3a626b2a2d72eb2197276 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Tue, 10 Sep 2024 14:53:54 -0700
Subject: [PATCH 46/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index c0f8c7c38071..cc0ceb7a6042 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -101,21 +101,20 @@ unsafe fn set_upto_64bits(
         }
     } else if len == 1 {
         let byte_chunk = (unsafe { data.get_unchecked(read_byte) } >> read_shift) & 1;
-        let ptr = write_data.as_mut_ptr();
-        unsafe { *ptr.add(write_byte) |= byte_chunk << write_shift };
+        unsafe { *write_data.get_unchecked_mut(write_byte) |= byte_chunk << write_shift };
         ((byte_chunk ^ 1) as usize, 1)
     } else {
         let len = std::cmp::min(len, 64 - std::cmp::max(read_shift, write_shift));
         let bytes = ceil(len + read_shift, 8);
+        // SAFETY: the args of `read_bytes_to_u64` are valid as read_byte + bytes <= data.len()
         let chunk = unsafe { read_bytes_to_u64(data, read_byte, bytes) };
         let mask = u64::MAX >> (64 - len);
-        let chunk = (chunk >> read_shift) & mask;
-        let chunk = chunk << write_shift;
+        let chunk = (chunk >> read_shift) & mask; // masking to read `len` bits only
+        let chunk = chunk << write_shift; // shifting back to align with `write_data`
         let null_count = len - chunk.count_ones() as usize;
         let bytes = ceil(len + write_shift, 8);
-        let ptr = unsafe { write_data.as_mut_ptr().add(write_byte) };
         for (i, c) in chunk.to_le_bytes().iter().enumerate().take(bytes) {
-            unsafe { *ptr.add(i) |= c };
+            unsafe { *write_data.get_unchecked_mut(write_byte + i) |= c };
         }
         (null_count, len)
     }

From 7d8107613d5d272d62f326ff8b0b4a42232012f1 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 13 Sep 2024 14:24:25 -0700
Subject: [PATCH 47/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index cc0ceb7a6042..893c363163b1 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -22,6 +22,7 @@ use crate::bit_util::ceil;
 /// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
 /// bits in `data` in the range `[offset_read..offset_read+len]`
 /// returns the number of `0` bits `data[offset_read..offset_read+len]`
+/// `offset_write`, `offset_read`, and `len` are in terms of bits
 pub fn set_bits(
     write_data: &mut [u8],
     data: &[u8],
@@ -53,6 +54,9 @@ pub fn set_bits(
     null_count
 }
 
+/// Similar to `set_bits` but sets only upto 64 bits, actual number of bits set may vary.
+/// Returns a pair of the number of `0` bits and the number of bits set
+///
 /// # Safety
 /// The caller must ensure all arguments are within the valid range.
 #[inline]

From f185a195f8b2e2261bd0ddbd5439189088b7b309 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 13 Sep 2024 16:13:54 -0700
Subject: [PATCH 48/48] address review comments

---
 arrow-buffer/src/util/bit_mask.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs
index 893c363163b1..2074f0fab988 100644
--- a/arrow-buffer/src/util/bit_mask.rs
+++ b/arrow-buffer/src/util/bit_mask.rs
@@ -288,11 +288,11 @@ mod tests {
         let write_data: &mut [u8] = &mut [0; 9];
         let data: &[u8] = &[
             0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001, 0b00000001,
-            0b00000001,
+            0b00000001, 0b00000001,
         ];
         let offset_write = 1;
         let offset_read = 0;
-        let len = 64;
+        let len = 65;
         let (n, len_set) =
             unsafe { set_upto_64bits(write_data, data, offset_write, offset_read, len) };
         assert_eq!(n, 55);