From 542a6c67aef7cf5fa234e5d8134b88cffb9a3c36 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Wed, 11 Sep 2024 11:18:35 +0200
Subject: [PATCH] miri: support vector index arguments in simd_shuffle

---
 src/tools/miri/src/intrinsics/simd.rs             | 15 ++++++---------
 .../miri/tests/pass/intrinsics/portable-simd.rs   |  8 ++++++++
 .../simd-intrinsic-generic-elements.rs            |  0
 3 files changed, 14 insertions(+), 9 deletions(-)
 rename src/tools/miri/tests/pass/{ => intrinsics}/simd-intrinsic-generic-elements.rs (100%)

diff --git a/src/tools/miri/src/intrinsics/simd.rs b/src/tools/miri/src/intrinsics/simd.rs
index a5afd029b6519..85d7459bb4595 100644
--- a/src/tools/miri/src/intrinsics/simd.rs
+++ b/src/tools/miri/src/intrinsics/simd.rs
@@ -666,22 +666,19 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 let (right, right_len) = this.operand_to_simd(right)?;
                 let (dest, dest_len) = this.mplace_to_simd(dest)?;
 
-                // `index` is an array, not a SIMD type
-                let ty::Array(_, index_len) = index.layout.ty.kind() else {
-                    span_bug!(
-                        this.cur_span(),
-                        "simd_shuffle index argument has non-array type {}",
-                        index.layout.ty
-                    )
+                // `index` is an array or a SIMD type
+                let (index, index_len) = match index.layout.ty.kind() {
+                    // FIXME: remove this once `index` must always be a SIMD vector.
+                    ty::Array(..) => (index.assert_mem_place(), index.len(this)?),
+                    _ => this.operand_to_simd(index)?,
                 };
-                let index_len = index_len.eval_target_usize(*this.tcx, this.param_env());
 
                 assert_eq!(left_len, right_len);
                 assert_eq!(index_len, dest_len);
 
                 for i in 0..dest_len {
                     let src_index: u64 = this
-                        .read_immediate(&this.project_index(index, i)?)?
+                        .read_immediate(&this.project_index(&index, i)?)?
                         .to_scalar()
                         .to_u32()?
                         .into();
diff --git a/src/tools/miri/tests/pass/intrinsics/portable-simd.rs b/src/tools/miri/tests/pass/intrinsics/portable-simd.rs
index c4ba11d0a436f..daf75fee8fe13 100644
--- a/src/tools/miri/tests/pass/intrinsics/portable-simd.rs
+++ b/src/tools/miri/tests/pass/intrinsics/portable-simd.rs
@@ -620,6 +620,10 @@ fn simd_intrinsics() {
         );
         assert_eq!(simd_shuffle_generic::<_, i32x4, { &[3, 1, 0, 2] }>(a, b), a,);
         assert_eq!(simd_shuffle::<_, _, i32x4>(a, b, const { [3u32, 1, 0, 2] }), a,);
+        assert_eq!(
+            simd_shuffle::<_, _, i32x4>(a, b, const { u32x4::from_array([3u32, 1, 0, 2]) }),
+            a,
+        );
         assert_eq!(
             simd_shuffle_generic::<_, i32x4, { &[7, 5, 4, 6] }>(a, b),
             i32x4::from_array([4, 2, 1, 10]),
@@ -628,6 +632,10 @@ fn simd_intrinsics() {
             simd_shuffle::<_, _, i32x4>(a, b, const { [7u32, 5, 4, 6] }),
             i32x4::from_array([4, 2, 1, 10]),
         );
+        assert_eq!(
+            simd_shuffle::<_, _, i32x4>(a, b, const { u32x4::from_array([7u32, 5, 4, 6]) }),
+            i32x4::from_array([4, 2, 1, 10]),
+        );
     }
 }
 
diff --git a/src/tools/miri/tests/pass/simd-intrinsic-generic-elements.rs b/src/tools/miri/tests/pass/intrinsics/simd-intrinsic-generic-elements.rs
similarity index 100%
rename from src/tools/miri/tests/pass/simd-intrinsic-generic-elements.rs
rename to src/tools/miri/tests/pass/intrinsics/simd-intrinsic-generic-elements.rs