[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 095255ef68116dc9baebd1ad8ba7d9246152488d 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   | 13 ++++
 .../MergeFunc/ptrauth-const-compare.ll        | 67 +++++++++++++++++++
 2 files changed, 80 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..581e6b89fa649 100644
--- a/llvm/lib/Transforms/Utils/FunctionComparator.cpp
+++ b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
@@ -354,6 +354,19 @@ 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;
+    return cmpConstants(PA1->getDiscriminator(), PA2->getDiscriminator());
+  }
+
   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..75ebb834c6219
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/ptrauth-const-compare.ll
@@ -0,0 +1,67 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs --version 2
+;; Check the case for ConstantPtrAuth to be compared properly in FunctionComparator.cpp
+; RUN: opt -S -passes=mergefunc < %s | FileCheck %s
+target triple = "arm64e-apple-ios14.0.0"
+
+define { ptr, i64 } @"foo"() {
+start:
+  %func = alloca [8 x i8], align 8
+  br i1 false, label %bb5, label %bb9
+
+bb9:                                              ; preds = %bb2, %start
+  %self = load i8, ptr null, align 1
+  br i1 false, label %bb2, label %bb5
+
+bb5:                                              ; preds = %bb9, %start
+  ret { ptr, i64 } zeroinitializer
+
+bb2:                                              ; preds = %bb9
+  call void ptrauth (ptr null, i32 0)(ptr null, i8 0)
+  br label %bb9
+}
+
+define { ptr, i64 } @"bar"() {
+start:
+  %func = alloca [8 x i8], align 8
+  br i1 false, label %bb5, label %bb9
+
+bb9:                                              ; preds = %bb2, %start
+  %self = load i8, ptr null, align 1
+  br i1 false, label %bb2, label %bb5
+
+bb5:                                              ; preds = %bb9, %start
+  ret { ptr, i64 } zeroinitializer
+
+bb2:                                              ; preds = %bb9
+  call void ptrauth (ptr @"baz", i32 0)(ptr null, i8 0)
+  br label %bb9
+}
+
+declare void @"baz"()
+; CHECK-LABEL: define { ptr, i64 } @foo() {
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[FUNC:%.*]] = alloca [8 x i8], align 8
+; CHECK-NEXT:    br i1 false, label [[BB5:%.*]], label [[BB9:%.*]]
+; CHECK:       bb9:
+; CHECK-NEXT:    [[SELF:%.*]] = load i8, ptr null, align 1
+; CHECK-NEXT:    br i1 false, label [[BB2:%.*]], label [[BB5]]
+; CHECK:       bb5:
+; CHECK-NEXT:    ret { ptr, i64 } zeroinitializer
+; CHECK:       bb2:
+; CHECK-NEXT:    call void ptrauth (ptr null, i32 0)(ptr null, i8 0)
+; CHECK-NEXT:    br label [[BB9]]
+;
+;
+; CHECK-LABEL: define { ptr, i64 } @bar() {
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[FUNC:%.*]] = alloca [8 x i8], align 8
+; CHECK-NEXT:    br i1 false, label [[BB5:%.*]], label [[BB9:%.*]]
+; CHECK:       bb9:
+; CHECK-NEXT:    [[SELF:%.*]] = load i8, ptr null, align 1
+; CHECK-NEXT:    br i1 false, label [[BB2:%.*]], label [[BB5]]
+; CHECK:       bb5:
+; CHECK-NEXT:    ret { ptr, i64 } zeroinitializer
+; CHECK:       bb2:
+; CHECK-NEXT:    call void ptrauth (ptr @baz, i32 0)(ptr null, i8 0)
+; CHECK-NEXT:    br label [[BB9]]
+;