From 3ba203345a2680d6cec3a0dae5447d432cb0bf2a Mon Sep 17 00:00:00 2001
From: rdambrosio <rdambrosio016@gmail.com>
Date: Fri, 30 Apr 2021 17:33:25 -0400
Subject: [PATCH] Feat: add shift_left and shift_right

---
 crates/core_simd/src/permute.rs   | 52 +++++++++++++++++++++++++++++++
 crates/core_simd/tests/permute.rs | 22 +++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/crates/core_simd/src/permute.rs b/crates/core_simd/src/permute.rs
index dd63c69c63d..96fff8ca947 100644
--- a/crates/core_simd/src/permute.rs
+++ b/crates/core_simd/src/permute.rs
@@ -122,6 +122,58 @@ macro_rules! impl_shuffle_lane {
                 }
                 (self.shuffle::<{ even() }>(other), self.shuffle::<{ odd() }>(other))
             }
+
+            /// Shift a vector `N` times to the left and fill the rest with zeroes.
+            ///
+            /// # Examples
+            ///
+            /// ```
+            /// # use core_simd::SimdU32;
+            /// let a = SimdU32::from_array([0, 1, 2, 3]);
+            /// let b = SimdU32::from_array([2, 3, 0, 0]);
+            /// assert_eq!(a.shift_left::<{2}>(), b);
+            /// ```
+            #[inline]
+            pub fn shift_left<const N: u32>(self) -> Self {
+                const fn idx() -> [u32; $n] {
+                    let mut base = [$n; $n];
+                    let mut i = N;
+                    while i < $n {
+                        base[i - 2] += N;
+                        i += 1;
+                    }
+                    base
+                }
+                let zeroed = Self::default();
+
+                self.shuffle::<{ idx() }>(zeroed)
+            }
+
+            /// Shift a vector `N` times to the right and fill the rest with zeroes.
+            ///
+            /// # Examples
+            ///
+            /// ```
+            /// # use core_simd::SimdU32;
+            /// let a = SimdU32::from_array([0, 1, 2, 3]);
+            /// let b = SimdU32::from_array([0, 0, 0, 1]);
+            /// assert_eq!(a.shift_right::<{2}>(), b);
+            /// ```
+            #[inline]
+            pub fn shift_right<const N: u32>(self) -> Self {
+                const fn idx() -> [u32; $n] {
+                    let mut base = [$n; $n];
+                    let mut i = N;
+                    while i < $n {
+                        base[i] += N;
+                        i += 1;
+                    }
+                    base
+                }
+                let zeroed = Self::default();
+
+                self.shuffle::<{ idx() }>(zeroed)
+            }
         }
     }
 }
diff --git a/crates/core_simd/tests/permute.rs b/crates/core_simd/tests/permute.rs
index 2be43c9cf3c..d7745b009b3 100644
--- a/crates/core_simd/tests/permute.rs
+++ b/crates/core_simd/tests/permute.rs
@@ -33,3 +33,25 @@ fn interleave() {
     assert_eq!(even, a);
     assert_eq!(odd, b);
 }
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn shift_left() {
+    let a = SimdU32::from_array([0, 1, 2, 3]);
+    let b = SimdU32::from_array([2, 3, 0, 0]);
+    assert_eq!(a.shift_left::<{ 2 }>(), b);
+    let b = SimdU32::from_array([0, 0, 0, 0]);
+    assert_eq!(a.shift_left::<{ 5 }>(), b);
+    assert_eq!(a.shift_left::<{ 0 }>(), a);
+}
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn shift_right() {
+    let a = SimdU32::from_array([0, 1, 2, 3]);
+    let b = SimdU32::from_array([0, 0, 0, 1]);
+    assert_eq!(a.shift_right::<{ 2 }>(), b);
+    let b = SimdU32::from_array([0, 0, 0, 0]);
+    assert_eq!(a.shift_right::<{ 5 }>(), b);
+    assert_eq!(a.shift_right::<{ 0 }>(), a);
+}