fix big-endian bitmasks smaller than a byte

2 files changed, 23 insertions, 4 deletions

diff --git a/crates/core_simd/src/masks/full_masks.rs b/crates/core_simd/src/masks/full_masks.rs
index 8bbdf637de8..efa688b128f 100644
--- a/crates/core_simd/src/masks/full_masks.rs
+++ b/crates/core_simd/src/masks/full_masks.rs
@@ -68,14 +68,26 @@ where
 
 // Used for bitmask bit order workaround
 pub(crate) trait ReverseBits {
-    fn reverse_bits(self) -> Self;
+    // Reverse the least significant `n` bits of `self`.
+    // (Remaining bits must be 0.)
+    fn reverse_bits(self, n: usize) -> Self;
 }
 
 macro_rules! impl_reverse_bits {
     { $($int:ty),* } => {
         $(
         impl ReverseBits for $int {
-            fn reverse_bits(self) -> Self { <$int>::reverse_bits(self) }
+            #[inline(always)]
+            fn reverse_bits(self, n: usize) -> Self {
+                let rev = <$int>::reverse_bits(self);
+                let bitsize = core::mem::size_of::<$int>() * 8;
+                if n < bitsize {
+                    // Shift things back to the right
+                    rev >> (bitsize - n)
+                } else {
+                    rev
+                }
+            }
         }
         )*
     }
@@ -137,7 +149,7 @@ where
 
         // LLVM assumes bit order should match endianness
         if cfg!(target_endian = "big") {
-            bitmask.reverse_bits()
+            bitmask.reverse_bits(LANES)
         } else {
             bitmask
         }
@@ -150,7 +162,7 @@ where
     {
         // LLVM assumes bit order should match endianness
         let bitmask = if cfg!(target_endian = "big") {
-            bitmask.reverse_bits()
+            bitmask.reverse_bits(LANES)
         } else {
             bitmask
         };
diff --git a/crates/core_simd/tests/masks.rs b/crates/core_simd/tests/masks.rs
index 1c587630a36..d10c6610f50 100644
--- a/crates/core_simd/tests/masks.rs
+++ b/crates/core_simd/tests/masks.rs
@@ -84,6 +84,7 @@ macro_rules! test_mask_api {
             #[test]
             fn roundtrip_bitmask_conversion_short() {
                 use core_simd::ToBitMask;
+
                 let values = [
                     false, false, false, true,
                 ];
@@ -91,6 +92,12 @@ macro_rules! test_mask_api {
                 let bitmask = mask.to_bitmask();
                 assert_eq!(bitmask, 0b1000);
                 assert_eq!(core_simd::Mask::<$type, 4>::from_bitmask(bitmask), mask);
+
+                let values = [true, false];
+                let mask = core_simd::Mask::<$type, 2>::from_array(values);
+                let bitmask = mask.to_bitmask();
+                assert_eq!(bitmask, 0b01);
+                assert_eq!(core_simd::Mask::<$type, 2>::from_bitmask(bitmask), mask);
             }
         }
     }