[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 72d8ad7dca5ddc289cf3d0497f8ffff8a9bfe22c 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   |  12 ++
 .../MergeFunc/ptrauth-const-compare.ll        | 130 ++++++++++++++++++
 2 files changed, 142 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 7cf19177db545..6928de881f797 100644
--- a/llvm/lib/Transforms/Utils/FunctionComparator.cpp
+++ b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
@@ -518,6 +518,18 @@ int FunctionComparator::cmpConstants(const Constant *L,
     const auto *REquiv = cast<DSOLocalEquivalent>(R);
     return cmpGlobalValues(LEquiv->getGlobalValue(), REquiv->getGlobalValue());
   }
+  case Value::ConstantPtrAuthVal: {
+    // Handle authenticated pointer constants produced by ConstantPtrAuth::get.
+    const ConstantPtrAuth *LPA = cast<ConstantPtrAuth>(L);
+    const ConstantPtrAuth *RPA = cast<ConstantPtrAuth>(R);
+    if (int Res = cmpConstants(LPA->getPointer(), RPA->getPointer()))
+      return Res;
+    if (int Res = cmpConstants(LPA->getKey(), RPA->getKey()))
+      return Res;
+    if (int Res = cmpConstants(LPA->getDiscriminator(), RPA->getDiscriminator()))
+      return Res;
+    return cmpConstants(LPA->getAddrDiscriminator(), RPA->getAddrDiscriminator());
+  }
   default: // Unknown constant, abort.
     LLVM_DEBUG(dbgs() << "Looking at valueID " << L->getValueID() << "\n");
     llvm_unreachable("Constant ValueID not recognized.");
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..931e64e3bc733
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/ptrauth-const-compare.ll
@@ -0,0 +1,130 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; 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 @callee(ptr)
+
+; different base pointer (null vs @baz)
+
+define i32 @f_ptr_null() {
+; CHECK-LABEL: define i32 @f_ptr_null() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr null, i32 0))
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  call void @callee(ptr ptrauth (ptr null, i32 0))
+  ret i32 1
+}
+
+define i32 @g_ptr_baz() {
+; CHECK-LABEL: define i32 @g_ptr_baz() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0))
+; CHECK-NEXT:    ret i32 2
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0))
+  ret i32 2
+}
+
+; different key (i32 0 vs i32 1)
+
+define i32 @f_key0() {
+; CHECK-LABEL: define i32 @f_key0() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0))
+; CHECK-NEXT:    ret i32 3
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0))
+  ret i32 3
+}
+
+define i32 @g_key1() {
+; CHECK-LABEL: define i32 @g_key1() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 1))
+; CHECK-NEXT:    ret i32 4
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 1))
+  ret i32 4
+}
+
+; different integer disc (i64 0 vs i64 7)
+
+define i32 @f_disc0() {
+; CHECK-LABEL: define i32 @f_disc0() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0))
+; CHECK-NEXT:    ret i32 5
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0))
+  ret i32 5
+}
+
+define i32 @g_disc7() {
+; CHECK-LABEL: define i32 @g_disc7() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 7))
+; CHECK-NEXT:    ret i32 6
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 7))
+  ret i32 6
+}
+
+; different addr disc (ptr null vs @ADDR)
+
+define i32 @f_addr_null() {
+; CHECK-LABEL: define i32 @f_addr_null() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0))
+; CHECK-NEXT:    ret i32 7
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0))
+  ret i32 7
+}
+
+define i32 @g_addr_ADDR() {
+; CHECK-LABEL: define i32 @g_addr_ADDR() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR))
+; CHECK-NEXT:    ret i32 8
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR))
+  ret i32 8
+}
+
+; positive test: identical ptrauth operands, should be merged
+
+define void @merge_ptrauth_a() {
+; CHECK-LABEL: define void @merge_ptrauth_a() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR))
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR))
+  ret void
+}
+
+define void @merge_ptrauth_b() {
+; CHECK-LABEL: define void @merge_ptrauth_b() {
+; CHECK-NEXT:    tail call void @merge_ptrauth_a()
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @callee(ptr ptrauth (ptr @baz, i32 0, i64 0, ptr @ADDR))
+  ret void
+}