From 3c6b0b709377e5b4a47b7a299bfcc1f636a136d1 Mon Sep 17 00:00:00 2001
From: sayantn <sayantn05@gmail.com>
Date: Sat, 21 Dec 2024 13:54:12 +0530
Subject: [PATCH] Make `assert_instr` stricter

---
 crates/stdarch-test/src/disassembly.rs | 6 +++---
 crates/stdarch-test/src/lib.rs         | 7 ++++++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/crates/stdarch-test/src/disassembly.rs b/crates/stdarch-test/src/disassembly.rs
index a7b46a73fc..6a3ca06cb6 100644
--- a/crates/stdarch-test/src/disassembly.rs
+++ b/crates/stdarch-test/src/disassembly.rs
@@ -140,7 +140,7 @@ fn parse(output: &str) -> HashSet<Function> {
                     .filter(|&x| !x.is_empty())
                     .skip(1)
                     .map(str::to_lowercase)
-                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
+                    .skip_while(|s| matches!(&**s, "lock" | "vex")) // skip x86-specific prefix
                     .collect::<Vec<String>>()
             } else {
                 // objdump with --no-show-raw-insn
@@ -150,8 +150,8 @@ fn parse(output: &str) -> HashSet<Function> {
                 instruction
                     .split_whitespace()
                     .skip(1)
-                    .skip_while(|s| *s == "lock" || *s == "{evex}") // skip x86-specific prefix
-                    .map(std::string::ToString::to_string)
+                    .skip_while(|s| matches!(*s, "lock" | "{evex}" | "{vex}")) // skip x86-specific prefix
+                    .map(ToString::to_string)
                     .collect::<Vec<String>>()
             };
 
diff --git a/crates/stdarch-test/src/lib.rs b/crates/stdarch-test/src/lib.rs
index 3248de8f3a..41af91e8a5 100644
--- a/crates/stdarch-test/src/lib.rs
+++ b/crates/stdarch-test/src/lib.rs
@@ -84,7 +84,12 @@ pub fn assert(shim_addr: usize, fnname: &str, expected: &str) {
     // 2. It is a mark, indicating that the instruction will be
     // compiled into other instructions - mainly because of llvm
     // optimization.
-    let found = expected == "nop" || instrs.iter().any(|s| s.contains(expected));
+    let expected = if expected == "unknown" {
+        "<unknown>" // Workaround for rust-lang/stdarch#1674, todo: remove when the issue is fixed
+    } else {
+        expected
+    };
+    let found = expected == "nop" || instrs.iter().any(|s| s.starts_with(expected));
 
     // Look for subroutine call instructions in the disassembly to detect whether
     // inlining failed: all intrinsics are `#[inline(always)]`, so calling one