[llvm] [PtrAuth] Add ConstantPtrAuth comparator to FunctionComparator.cpp (PR #159480)

[Oskar Wirga via llvm-commits]

https://github.com/oskarwirga updated https://github.com/llvm/llvm-project/pull/159480

>From 10f5e779499abb474607e855631a74dacd386f6f Mon Sep 17 00:00:00 2001
From: Oskar Wirga <oskar.wirga at gmail.com>
Date: Thu, 11 Sep 2025 19:16:40 -0700
Subject: [PATCH] [PtrAuth] Add ConstantPtrAuth comparator to
 FunctionComparator.cpp

When building rustc std for arm64e, the optimizations result in devirtualization of indirect calls.

This can happen during function merging so I modified FunctionComparator.cpp as the ConstantPtrAuth value would go unchecked in the switch statement.
---
 .../Transforms/Utils/FunctionComparator.cpp   |  17 +++
 .../MergeFunc/ptrauth-const-compare.ll        | 101 ++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 llvm/test/Transforms/MergeFunc/ptrauth-const-compare.ll

diff --git a/llvm/lib/Transforms/Utils/FunctionComparator.cpp b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
index 6d4026e8209de..8aba7886c9051 100644
--- a/llvm/lib/Transforms/Utils/FunctionComparator.cpp
+++ b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
@@ -354,6 +354,23 @@ int FunctionComparator::cmpConstants(const Constant *L,
   if (!L->isNullValue() && R->isNullValue())
     return -1;
 
+  // Handle authenticated pointer constants produced by ConstantPtrAuth::get.
+  if (auto *PA1 = dyn_cast<ConstantPtrAuth>(L)) {
+    auto *PA2 = dyn_cast<ConstantPtrAuth>(R);
+    if (!PA2)
+      return cmpNumbers(L->getValueID(), R->getValueID());
+
+    if (int Res = cmpConstants(PA1->getPointer(), PA2->getPointer()))
+      return Res;
+    if (int Res = cmpConstants(PA1->getKey(), PA2->getKey()))
+      return Res;
+    if (int Res =
+            cmpConstants(PA1->getDiscriminator(), PA2->getDiscriminator()))
+      return Res;
+    return cmpConstants(PA1->getAddrDiscriminator(),
+                        PA2->getAddrDiscriminator());
+  }
+
   auto GlobalValueL = const_cast<GlobalValue *>(dyn_cast<GlobalValue>(L));
   auto GlobalValueR = const_cast<GlobalValue *>(dyn_cast<GlobalValue>(R));
   if (GlobalValueL && GlobalValueR) {
diff --git a/llvm/test/Transforms/MergeFunc/ptrauth-const-compare.ll b/llvm/test/Transforms/MergeFunc/ptrauth-const-compare.ll
new file mode 100644
index 0000000000000..b3877f21213f9
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/ptrauth-const-compare.ll
@@ -0,0 +1,101 @@
+; RUN: opt -passes=mergefunc -S %s | FileCheck %s
+; Ensure MergeFunc handles ConstantPtrAuth correctly and does not
+; merge when any ptrauth operand differs (ptr, key, int disc, addr disc).
+
+target triple = "arm64e-apple-ios14.0.0"
+
+declare void @baz()
+ at ADDR = external global i8
+
+declare void @sink0()
+declare void @sink1()
+declare void @sink2()
+declare void @sink3()
+
+; different base pointer (null vs @baz)
+
+define void @f_ptr_null() {
+; CHECK-LABEL: define void @f_ptr_null()
+; CHECK:       call void @sink0()
+; CHECK:       call void ptrauth (ptr null, i32 0)(
+entry:
+  call void @sink0()
+  call void ptrauth (ptr null, i32 0)(ptr null)
+  ret void
+}
+
+define void @g_ptr_baz() {
+; CHECK-LABEL: define void @g_ptr_baz()
+; CHECK:       call void @sink0()
+; CHECK:       call void ptrauth (ptr @baz, i32 0)(
+entry:
+  call void @sink0()
+  call void ptrauth (ptr @baz, i32 0)(ptr null)
+  ret void
+}
+
+; different key (i32 0 vs i32 1)
+
+define void @f_key0() {
+; CHECK-LABEL: define void @f_key0()
+; CHECK:       call void @sink1()
+; CHECK:       call void ptrauth (ptr @baz, i32 0)(
+entry:
+  call void @sink1()
+  call void ptrauth (ptr @baz, i32 0)(ptr null)
+  ret void
+}
+
+define void @g_key1() {
+; CHECK-LABEL: define void @g_key1()
+; CHECK:       call void @sink1()
+; CHECK:       call void ptrauth (ptr @baz, i32 1)(
+entry:
+  call void @sink1()
+  call void ptrauth (ptr @baz, i32 1)(ptr null)
+  ret void
+}
+
+; different integer disc (i64 0 vs i64 7)
+
+define void @f_disc0() {
+; CHECK-LABEL: define void @f_disc0()
+; CHECK:       call void @sink2()
+; CHECK:       call void ptrauth (ptr @baz, i32 0)(
+entry:
+  call void @sink2()
+  call void ptrauth (ptr @baz, i32 0)(ptr null)
+  ret void
+}
+
+define void @g_disc7() {
+; CHECK-LABEL: define void @g_disc7()
+; CHECK:       call void @sink2()
+; CHECK:       call void ptrauth (ptr @baz, i32 0, i64 7)(
+entry:
+  call void @sink2()
+  call void ptrauth (ptr @baz, i32 0, i64 7)(ptr null)
+  ret void
+}
+
+; different addr disc (ptr null vs @ADDR)
+
+define void @f_addr_null() {
+; CHECK-LABEL: define void @f_addr_null()
+; CHECK:       call void @sink3()
+; CHECK:       call void ptrauth (ptr @baz, i32 0)(
+entry:
+  call void @sink3()
+  call void ptrauth (ptr @baz, i32 0)(ptr null)
+  ret void
+}
+
+define void @g_addr_ADDR() {
+; CHECK-LABEL: define void @g_addr_ADDR()
+; CHECK:       call void @sink3()
+; CHECK:       call void ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR)(
+entry:
+  call void @sink3()
+  call void ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR)(ptr null)
+  ret void
+}