fix

Author: myrrc

vortex-buffer/src/bit/take.rs

@@ -80,7 +80,6 @@ where
             })
         } else {
             #[cfg(target_arch = "x86_64")]
-            #[target_feature(enable = "avx2")]
             if let Some(taken) = avx2::take(bits, indices, None) {
                 return taken;
             }
@@ -98,7 +97,6 @@ where
     };
 
     #[cfg(target_arch = "x86_64")]
-    #[target_feature(enable = "avx2")]
     if let Some(taken) = avx2::take(bits, indices, None) {
         return taken;
     }
@@ -111,7 +109,7 @@ where
             // SAFETY: we're already iterating over every index, a bound
             // is excessive
             let idx = unsafe { indices.get_unchecked(i).as_() };
-            // SAFETY: we're verified validity has same size as indices before
+            // SAFETY: we've verified validity has same size as indices before
             let mask_idx: usize = unsafe { get_bit_unchecked(ptr, offset + i) }.as_();
             bools[idx & mask_idx.wrapping_neg()]
         });
@@ -124,7 +122,7 @@ where
         // SAFETY: we're already iterating over every index, a bound
         // is excessive
         let idx = unsafe { indices.get_unchecked(i).as_() };
-        // SAFETY: we're verified validity has same size as indices before
+        // SAFETY: we've verified validity has same size as indices before
         let mask_idx: usize = unsafe { get_bit_unchecked(ptr, validity_offset + i) }.as_();
         let masked_idx = idx & mask_idx.wrapping_neg();
         get_bit(buf, bits_offset + masked_idx)
@@ -145,10 +143,7 @@ fn first_unset(bits: BitBufferView<'_>) -> Option<usize> {
 }
 
 #[cfg(target_arch = "x86_64")]
-#[target_feature(enable = "avx2")]
 mod avx2 {
-    use super::index_at;
-    use super::keep_mask;
     use crate::BitBuffer;
     use crate::BitBufferView;
     use crate::BufferMut;
@@ -186,43 +181,43 @@ mod avx2 {
             return None;
         }
 
+        if size_of::<I>() != size_of::<u16>() && size_of::<I>() != size_of::<u32>() {
+            return None;
+        }
+
+        // SAFETY: the feature detection above guarantees avx2 is present.
+        Some(unsafe { run(bits, indices, validity) })
+    }
+
+    #[target_feature(enable = "avx2")]
+    unsafe fn run<I: AsPrimitive<usize>>(
+        bits: BitBufferView<'_>,
+        indices: &[I],
+        validity: Option<BitBufferView<'_>>,
+    ) -> BitBuffer {
         if size_of::<I>() == size_of::<u16>() {
             // SAFETY: u16 or i16 can be reinterpreted as u16 since we guarantee
             // no negative indices
             let indices = unsafe { from_raw_parts(indices.as_ptr().cast::<u16>(), indices.len()) };
-            return Some(unsafe { take_u16(bits, indices, validity) });
-        } else if size_of::<I> == size_of::<u32>() {
+            let ptr = indices.as_ptr();
+            // SAFETY: avx2 is enabled for this function.
+            unsafe {
+                take_lanes(bits, indices, validity, |group| {
+                    _mm256_cvtepu16_epi32(_mm_loadu_si128(ptr.add(group * 8).cast()))
+                })
+            }
+        } else {
             // SAFETY: u32 or i32 can be reinterpreted as u32 since we guarantee
             // no negative indices
             let indices = unsafe { from_raw_parts(indices.as_ptr().cast::<u32>(), indices.len()) };
-            return Some(unsafe { take_u32(bits, indices, validity) });
+            let ptr = indices.as_ptr();
+            // SAFETY: avx2 is enabled for this function.
+            unsafe { take_lanes(bits, indices, validity, |group| _mm256_loadu_si256(ptr.add(group * 8).cast())) }
         }
-        None
-    }
-
-    unsafe fn take_u16(
-        bits: BitBufferView<'_>,
-        indices: &[u16],
-        validity: Option<BitBufferView<'_>>,
-    ) -> BitBuffer {
-        let ptr = indices.as_ptr();
-        take_lanes(bits, indices, validity, |group| {
-            _mm256_cvtepu16_epi32(_mm_loadu_si128(ptr.add(group * 8).cast()))
-        })
-    }
-
-    unsafe fn take_u32(
-        bits: BitBufferView<'_>,
-        indices: &[u32],
-        validity: Option<BitBufferView<'_>>,
-    ) -> BitBuffer {
-        let ptr = indices.as_ptr();
-        take_lanes(bits, indices, validity, |group| {
-            _mm256_loadu_si256(ptr.add(group * 8).cast())
-        })
     }
 
     #[expect(clippy::cast_possible_truncation)]
+    #[target_feature(enable = "avx2")]
     unsafe fn take_lanes<I>(
         bits: BitBufferView<'_>,
         indices: &[I],
@@ -271,10 +266,12 @@ mod avx2 {
             let start = full_groups * 8;
             let mut byte = 0u8;
             for bit in 0..tail {
+                let pos = start + bit;
                 let keep = validity.map_or(usize::MAX, |(valid_bytes, valid_offset)| {
-                    keep_mask(valid_bytes.as_ptr(), valid_offset, start + bit)
+                    (unsafe { get_bit_unchecked(valid_bytes.as_ptr(), valid_offset + pos) } as usize)
+                        .wrapping_neg()
                 });
-                let idx = index_at(indices, start + bit) & keep;
+                let idx = unsafe { indices.get_unchecked(pos) }.as_() & keep;
                 let value = unsafe { get_bit_unchecked(base, bits.offset() + idx) };
                 byte |= (value as u8) << bit;
             }