miri: support vector index arguments in simd_shuffle

3 files changed, 14 insertions, 9 deletions

diff --git a/src/tools/miri/src/intrinsics/simd.rs b/src/tools/miri/src/intrinsics/simd.rs
index a5afd029b65..85d7459bb45 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 c4ba11d0a43..daf75fee8fe 100644
--- a/src/tools/miri/tests/pass/intrinsics/portable-simd.rs
+++ b/src/tools/miri/tests/pass/intrinsics/portable-simd.rs
@@ -621,6 +621,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
index 9cf0c2ddef3..9cf0c2ddef3 100644
--- a/src/tools/miri/tests/pass/simd-intrinsic-generic-elements.rs
+++ b/src/tools/miri/tests/pass/intrinsics/simd-intrinsic-generic-elements.rs