[llvm] [CFI][annotation] Leave alone function pointers in function annotations (PR #80173)

[via llvm-commits]

https://github.com/yozhu created https://github.com/llvm/llvm-project/pull/80173

Function annotation, as part of llvm.metadata, is for the function itself and doesn't apply to its corresponding jump table entry, so with CFI we shouldn't replace function pointer in function annotation with pointer to its corresponding jump table entry.

>From c91a3a6f07488703d9f939a0d2477023589fc1fe Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at meta.com>
Date: Thu, 25 Jan 2024 20:37:10 -0800
Subject: [PATCH] [CFI][annotation] Leave alone function pointers in function
 annotations

Function annotation, as part of llvm.metadata, is for the function itself
and doesn't apply to its corresponding jump table entry, so with CFI we
shouldn't replace function pointer in function annotation with pointer to
its corresponding jump table entry.
---
 llvm/lib/Transforms/IPO/LowerTypeTests.cpp | 30 +++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 733f290b1bc93..f6630019eaec6 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -470,6 +470,9 @@ class LowerTypeTestsModule {
 
   Function *WeakInitializerFn = nullptr;
 
+  GlobalVariable *GlobalAnnotation;
+  std::set<void *> FunctionAnnotations;
+
   bool shouldExportConstantsAsAbsoluteSymbols();
   uint8_t *exportTypeId(StringRef TypeId, const TypeIdLowering &TIL);
   TypeIdLowering importTypeId(StringRef TypeId);
@@ -531,6 +534,10 @@ class LowerTypeTestsModule {
   /// replace each use, which is a direct function call.
   void replaceDirectCalls(Value *Old, Value *New);
 
+  bool isFunctionAnnotation(void *GV) {
+    return FunctionAnnotations.count(GV) != 0;
+  }
+
 public:
   LowerTypeTestsModule(Module &M, ModuleAnalysisManager &AM,
                        ModuleSummaryIndex *ExportSummary,
@@ -1377,8 +1384,11 @@ void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
   // (all?) targets. Switch to a runtime initializer.
   SmallSetVector<GlobalVariable *, 8> GlobalVarUsers;
   findGlobalVariableUsersOf(F, GlobalVarUsers);
-  for (auto *GV : GlobalVarUsers)
+  for (auto *GV : GlobalVarUsers) {
+    if (GV == GlobalAnnotation)
+      continue;
     moveInitializerToModuleConstructor(GV);
+  }
 
   // Can not RAUW F with an expression that uses F. Replace with a temporary
   // placeholder first.
@@ -1392,6 +1402,10 @@ void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
   // Don't use range based loop, because use list will be modified.
   while (!PlaceholderFn->use_empty()) {
     Use &U = *PlaceholderFn->use_begin();
+    if (isFunctionAnnotation(U.getUser())) {
+      U.set(F);
+      continue;
+    }
     auto *InsertPt = dyn_cast<Instruction>(U.getUser());
     assert(InsertPt && "Non-instruction users should have been eliminated");
     auto *PN = dyn_cast<PHINode>(InsertPt);
@@ -1837,6 +1851,16 @@ LowerTypeTestsModule::LowerTypeTestsModule(
   }
   OS = TargetTriple.getOS();
   ObjectFormat = TargetTriple.getObjectFormat();
+
+  // Function annotation describes or applies to function itself, and
+  // shouldn't be associated with jump table thunk generated for CFI.
+  GlobalAnnotation = M.getGlobalVariable("llvm.global.annotations");
+  auto *C = dyn_cast_or_null<Constant>(GlobalAnnotation);
+  if (C && C->getNumOperands() == 1) {
+    C = cast<Constant>(C->getOperand(0));
+    for (auto &Op : C->operands())
+      FunctionAnnotations.insert(Op.get());
+  }
 }
 
 bool LowerTypeTestsModule::runForTesting(Module &M, ModuleAnalysisManager &AM) {
@@ -1900,6 +1924,10 @@ void LowerTypeTestsModule::replaceCfiUses(Function *Old, Value *New,
     if (isDirectCall(U) && (Old->isDSOLocal() || !IsJumpTableCanonical))
       continue;
 
+    // Skip function annotation
+    if (isFunctionAnnotation(U.getUser()))
+      continue;
+
     // Must handle Constants specially, we cannot call replaceUsesOfWith on a
     // constant because they are uniqued.
     if (auto *C = dyn_cast<Constant>(U.getUser())) {