Move implementation of SSE4.1 `ptest*` into a helper function

2 files changed, 41 insertions, 21 deletions

diff --git a/src/tools/miri/src/shims/x86/mod.rs b/src/tools/miri/src/shims/x86/mod.rs
index c4c361c3d63..6d361f5d2a5 100644
--- a/src/tools/miri/src/shims/x86/mod.rs
+++ b/src/tools/miri/src/shims/x86/mod.rs
@@ -615,3 +615,31 @@ fn horizontal_bin_op<'tcx>(
 
     Ok(())
 }
+
+/// Folds SIMD vectors `lhs` and `rhs` into a value of type `T` using `f`.
+fn bin_op_folded<'tcx, T>(
+    this: &crate::MiriInterpCx<'_, 'tcx>,
+    lhs: &OpTy<'tcx, Provenance>,
+    rhs: &OpTy<'tcx, Provenance>,
+    init: T,
+    mut f: impl FnMut(T, ImmTy<'tcx, Provenance>, ImmTy<'tcx, Provenance>) -> InterpResult<'tcx, T>,
+) -> InterpResult<'tcx, T> {
+    assert_eq!(lhs.layout, rhs.layout);
+
+    let (lhs, lhs_len) = this.operand_to_simd(lhs)?;
+    let (rhs, rhs_len) = this.operand_to_simd(rhs)?;
+
+    assert_eq!(lhs_len, rhs_len);
+
+    let mut acc = init;
+    for i in 0..lhs_len {
+        let lhs = this.project_index(&lhs, i)?;
+        let rhs = this.project_index(&rhs, i)?;
+
+        let lhs = this.read_immediate(&lhs)?;
+        let rhs = this.read_immediate(&rhs)?;
+        acc = f(acc, lhs, rhs)?;
+    }
+
+    Ok(acc)
+}
diff --git a/src/tools/miri/src/shims/x86/sse41.rs b/src/tools/miri/src/shims/x86/sse41.rs
index 10437919085..2683105b00b 100644
--- a/src/tools/miri/src/shims/x86/sse41.rs
+++ b/src/tools/miri/src/shims/x86/sse41.rs
@@ -2,7 +2,7 @@ use rustc_middle::mir;
 use rustc_span::Symbol;
 use rustc_target::spec::abi::Abi;
 
-use super::{round_all, round_first};
+use super::{bin_op_folded, round_all, round_first};
 use crate::*;
 use shims::foreign_items::EmulateForeignItemResult;
 
@@ -257,26 +257,18 @@ pub(super) trait EvalContextExt<'mir, 'tcx: 'mir>:
             "ptestz" | "ptestc" | "ptestnzc" => {
                 let [op, mask] = this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
 
-                let (op, op_len) = this.operand_to_simd(op)?;
-                let (mask, mask_len) = this.operand_to_simd(mask)?;
-
-                assert_eq!(op_len, mask_len);
-
-                let f = match unprefixed_name {
-                    "ptestz" => |op, mask| op & mask == 0,
-                    "ptestc" => |op, mask| op & mask == mask,
-                    "ptestnzc" => |op, mask| op & mask != 0 && op & mask != mask,
-                    _ => unreachable!(),
-                };
-
-                let mut all_zero = true;
-                for i in 0..op_len {
-                    let op = this.read_scalar(&this.project_index(&op, i)?)?.to_u64()?;
-                    let mask = this.read_scalar(&this.project_index(&mask, i)?)?.to_u64()?;
-                    all_zero &= f(op, mask);
-                }
-
-                this.write_scalar(Scalar::from_i32(all_zero.into()), dest)?;
+                let res = bin_op_folded(this, op, mask, true, |acc, op, mask| {
+                    let op = op.to_scalar().to_uint(op.layout.size)?;
+                    let mask = mask.to_scalar().to_uint(mask.layout.size)?;
+                    Ok(match unprefixed_name {
+                        "ptestz" => acc && (op & mask) == 0,
+                        "ptestc" => acc && (op & mask) == mask,
+                        "ptestnzc" => acc && (op & mask) != 0 && (op & mask) != mask,
+                        _ => unreachable!(),
+                    })
+                })?;
+
+                this.write_scalar(Scalar::from_i32(res.into()), dest)?;
             }
             _ => return Ok(EmulateForeignItemResult::NotSupported),
         }