[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.

Author: oskarwirga

llvm/lib/Transforms/Utils/FunctionComparator.cpp

@@ -518,6 +518,20 @@ 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.");

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()
+@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
+}