[lldb][AArch64] Simplify MTE memory region test

By checking for cpu and toolchain features ahead
of time we don't need the custom return codes.

Reviewed By: omjavaid

Differential Revision: https://reviews.llvm.org/D97684

GitOrigin-RevId: 6cdc2239dbabeb6fb8a9f933693f744a60d50a8c
diff --git a/packages/Python/lldbsuite/test/decorators.py b/packages/Python/lldbsuite/test/decorators.py
index 6d781de..d67572e 100644
--- a/packages/Python/lldbsuite/test/decorators.py
+++ b/packages/Python/lldbsuite/test/decorators.py
@@ -823,6 +823,39 @@
     """Skip this test if the environment is set up to run LLDB *itself* under ASAN."""
     return skipTestIfFn(is_running_under_asan)(func)
 
+def skipUnlessAArch64MTELinuxCompiler(func):
+    """Decorate the item to skip test unless MTE is supported by the test compiler."""
+
+    def is_toolchain_with_mte(self):
+        compiler_path = self.getCompiler()
+        compiler = os.path.basename(compiler_path)
+        f = tempfile.NamedTemporaryFile()
+        if lldbplatformutil.getPlatform() == 'windows':
+            return "MTE tests are not compatible with 'windows'"
+
+        cmd = "echo 'int main() {}' | %s -x c -o %s -" % (compiler_path, f.name)
+        if os.popen(cmd).close() is not None:
+            # Cannot compile at all, don't skip the test
+            # so that we report the broken compiler normally.
+            return None
+
+        # We need the Linux headers and ACLE MTE intrinsics
+        test_src = """
+            #include <asm/hwcap.h>
+            #include <arm_acle.h>
+            #ifndef HWCAP2_MTE
+            #error
+            #endif
+            int main() {
+                void* ptr = __arm_mte_create_random_tag((void*)(0), 0);
+            }"""
+        cmd = "echo '%s' | %s -march=armv8.5-a+memtag -x c -o %s -" % (test_src, compiler_path, f.name)
+        if os.popen(cmd).close() is not None:
+            return "Toolchain does not support MTE"
+        return None
+
+    return skipTestIfFn(is_toolchain_with_mte)(func)
+
 def _get_bool_config(key, fail_value = True):
     """
     Returns the current LLDB's build config value.
diff --git a/test/API/linux/aarch64/mte_memory_region/TestAArch64LinuxMTEMemoryRegion.py b/test/API/linux/aarch64/mte_memory_region/TestAArch64LinuxMTEMemoryRegion.py
index 3855a5a..29035c3 100644
--- a/test/API/linux/aarch64/mte_memory_region/TestAArch64LinuxMTEMemoryRegion.py
+++ b/test/API/linux/aarch64/mte_memory_region/TestAArch64LinuxMTEMemoryRegion.py
@@ -17,8 +17,9 @@
 
     NO_DEBUG_INFO_TESTCASE = True
 
-    @skipIf(archs=no_match(["aarch64"]))
+    @skipUnlessArch("aarch64")
     @skipUnlessPlatform(["linux"])
+    @skipUnlessAArch64MTELinuxCompiler
     def test_mte_regions(self):
         if not self.isAArch64MTE():
             self.skipTest('Target must support MTE.')
@@ -35,15 +36,6 @@
         self.runCmd("run", RUN_SUCCEEDED)
 
         if self.process().GetState() == lldb.eStateExited:
-            # 47 = non MTE toolchain
-            # 48 = non MTE target
-            exit_status = self.process().GetExitStatus()
-            if exit_status == 47:
-                self.skipTest("MTE must be available in toolchain")
-            elif exit_status == 48:
-                self.skipTest("target must have MTE enabled")
-
-            # Otherwise we have MTE but another problem occured
             self.fail("Test program failed to run.")
 
         self.expect("thread list", STOPPED_DUE_TO_BREAKPOINT,
diff --git a/test/API/linux/aarch64/mte_memory_region/main.c b/test/API/linux/aarch64/mte_memory_region/main.c
index 17c135d..5f07e5c 100644
--- a/test/API/linux/aarch64/mte_memory_region/main.c
+++ b/test/API/linux/aarch64/mte_memory_region/main.c
@@ -5,22 +5,9 @@
 #include <sys/prctl.h>
 #include <unistd.h>
 
-#define INCOMPATIBLE_TOOLCHAIN 47
-#define INCOMPATIBLE_TARGET 48
-
-// This is in a seperate non static function
-// so that we can always breakpoint the return 0 here.
-// Even if main never reaches it because HWCAP2_MTE
-// is not defined.
-// If it were in main then you would effectively have:
-// return TEST_INCOMPATIBLE;
-// return 0;
-// So the two returns would have the same breakpoint location
-// and we couldn't tell them apart.
-int setup_mte_page(void) {
-#ifdef HWCAP2_MTE
+int main(int argc, char const *argv[]) {
   if (!(getauxval(AT_HWCAP2) & HWCAP2_MTE))
-    return INCOMPATIBLE_TARGET;
+    return 1;
 
   int got = prctl(PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE, 0, 0, 0);
   if (got)
@@ -30,15 +17,6 @@
                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
   if (the_page == MAP_FAILED)
     return 1;
-#endif
 
   return 0; // Set break point at this line.
 }
-
-int main(int argc, char const *argv[]) {
-#ifdef HWCAP2_MTE
-  return setup_mte_page();
-#else
-  return INCOMPATIBLE_TOOLCHAIN;
-#endif
-}