From 15b4e280049a2b495635027ebfcb3b6930fbc12a Mon Sep 17 00:00:00 2001
From: Caleb Zulawski <caleb.zulawski@gmail.com>
Date: Mon, 21 Jun 2021 17:05:43 -0400
Subject: [PATCH] Add from_bitmask (#136)

* Add from_bitmask
* Add mips workaround
---
 crates/core_simd/src/masks/bitmask.rs    |  9 +++++++++
 crates/core_simd/src/masks/full_masks.rs | 24 ++++++++++++++++++++++++
 crates/core_simd/src/masks/mod.rs        |  5 +++++
 crates/core_simd/tests/masks.rs          |  6 ++++--
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/crates/core_simd/src/masks/bitmask.rs b/crates/core_simd/src/masks/bitmask.rs
index 31c7f6e7c289..fab136d2b241 100644
--- a/crates/core_simd/src/masks/bitmask.rs
+++ b/crates/core_simd/src/masks/bitmask.rs
@@ -100,6 +100,15 @@ impl<T: Mask, const LANES: usize> BitMask<T, LANES> {
         unsafe { core::mem::transmute_copy(&self.0) }
     }
 
+    #[inline]
+    pub fn from_bitmask<U: Mask>(bitmask: U::BitMask) -> Self {
+        assert_eq!(
+            core::mem::size_of::<T::BitMask>(),
+            core::mem::size_of::<U::BitMask>()
+        );
+        unsafe { core::mem::transmute_copy(&bitmask) }
+    }
+
     #[inline]
     pub fn any(self) -> bool {
         self != Self::splat(false)
diff --git a/crates/core_simd/src/masks/full_masks.rs b/crates/core_simd/src/masks/full_masks.rs
index c2bfa03dfc6e..7d98333ef607 100644
--- a/crates/core_simd/src/masks/full_masks.rs
+++ b/crates/core_simd/src/masks/full_masks.rs
@@ -126,6 +126,30 @@ macro_rules! define_mask {
                     bitmask
                 }
             }
+
+            #[inline]
+            pub fn from_bitmask<U: crate::Mask>(mut bitmask: U::BitMask) -> Self {
+                unsafe {
+                    // There is a bug where LLVM appears to implement this operation with the wrong
+                    // bit order.
+                    // TODO fix this in a better way
+                    if cfg!(any(target_arch = "mips", target_arch = "mips64")) {
+                        for x in bitmask.as_mut() {
+                            *x = x.reverse_bits();
+                        }
+                    }
+
+                    // TODO remove the transmute when rustc is more flexible
+                    assert_eq!(core::mem::size_of::<U::IntBitMask>(), core::mem::size_of::<U::BitMask>());
+                    let bitmask: U::IntBitMask = core::mem::transmute_copy(&bitmask);
+
+                    Self::from_int_unchecked(crate::intrinsics::simd_select_bitmask(
+                        bitmask,
+                        Self::splat(true).to_int(),
+                        Self::splat(false).to_int(),
+                    ))
+                }
+            }
         }
 
         impl<T: Mask, const LANES: usize> core::convert::From<$name<T, LANES>> for crate::$type<LANES>
diff --git a/crates/core_simd/src/masks/mod.rs b/crates/core_simd/src/masks/mod.rs
index 81a410de200e..1d6b2e45224a 100644
--- a/crates/core_simd/src/masks/mod.rs
+++ b/crates/core_simd/src/masks/mod.rs
@@ -178,6 +178,11 @@ macro_rules! define_opaque_mask {
             pub fn to_bitmask(self) -> <Self as Mask>::BitMask {
                 self.0.to_bitmask::<Self>()
             }
+
+            /// Convert a bitmask to a mask.
+            pub fn from_bitmask(bitmask: <Self as Mask>::BitMask) -> Self {
+                Self(<$inner_ty>::from_bitmask::<Self>(bitmask))
+            }
         }
 
         // vector/array conversion
diff --git a/crates/core_simd/tests/masks.rs b/crates/core_simd/tests/masks.rs
index 5c2e60dd7c82..32dea49729f0 100644
--- a/crates/core_simd/tests/masks.rs
+++ b/crates/core_simd/tests/masks.rs
@@ -67,13 +67,15 @@ macro_rules! test_mask_api {
             }
 
             #[test]
-            fn to_bitmask() {
+            fn roundtrip_bitmask_conversion() {
                 let values = [
                     true, false, false, true, false, false, true, false,
                     true, true, false, false, false, false, false, true,
                 ];
                 let mask = core_simd::$name::<16>::from_array(values);
-                assert_eq!(mask.to_bitmask(), [0b01001001, 0b10000011]);
+                let bitmask = mask.to_bitmask();
+                assert_eq!(bitmask, [0b01001001, 0b10000011]);
+                assert_eq!(core_simd::$name::<16>::from_bitmask(bitmask), mask);
             }
         }
     }