fix big-endian bitmasks smaller than a byte

Author: RalfJung

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
         };

crates/core_simd/tests/masks.rs

@@ -84,13 +84,20 @@ macro_rules! test_mask_api {
             #[test]
             fn roundtrip_bitmask_conversion_short() {
                 use core_simd::ToBitMask;
+
                 let values = [
                     false, false, false, true,
                 ];
                 let mask = core_simd::Mask::<$type, 4>::from_array(values);
                 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);
             }
         }
     }