[clang][PAC] Upstream pointer auth constant cachingusers/ojhunt/ptrauth-constant-cache-upstreaming

4 files changed, 97 insertions, 16 deletions

diff --git a/clang/lib/CodeGen/CGPointerAuth.cpp b/clang/lib/CodeGen/CGPointerAuth.cpp
index 375f87a..ff569e9 100644
--- a/clang/lib/CodeGen/CGPointerAuth.cpp
+++ b/clang/lib/CodeGen/CGPointerAuth.cpp
@@ -21,6 +21,40 @@
 using namespace clang;
 using namespace CodeGen;
 
+struct CodeGenModule::PointerAuthCachesTy {
+  struct PointerAuthConstantEntry {
+    unsigned Key;
+    llvm::Constant *OtherDiscriminator;
+    llvm::GlobalVariable *Global;
+  };
+
+  using PointerAuthConstantEntries =
+    std::vector<PointerAuthConstantEntry>;
+  using ByConstantCacheTy =
+    llvm::ValueMap<llvm::Constant*, PointerAuthConstantEntries>;
+  using ByDeclCacheTy =
+    llvm::DenseMap<const Decl *, llvm::Constant*>;
+  using PtrAuthDiscriminatorHashCacheTy = llvm::DenseMap<GlobalDecl, uint16_t>;
+  PtrAuthDiscriminatorHashCacheTy PtrAuthDiscriminatorHashes;
+  using VTablePtrAuthInfoCacheTy =
+    llvm::DenseMap<const CXXRecordDecl *, std::optional<PointerAuthQualifier>>;
+  VTablePtrAuthInfoCacheTy VTablePtrAuthInfos;
+  ByConstantCacheTy ConstantSignedPointersByConstant;
+  ByDeclCacheTy ConstantSignedPointersByDecl;
+  ByDeclCacheTy SignedThunkPointers;
+};
+
+CodeGenModule::PointerAuthCachesTy&
+CodeGenModule::getPointerAuthCaches() const {
+  if (!PointerAuthCaches)
+    PointerAuthCaches = new PointerAuthCachesTy;
+  return *PointerAuthCaches;
+}
+
+void CodeGenModule::destroyPointerAuthCaches() {
+  delete PointerAuthCaches;
+}
+
 /// Given a pointer-authentication schema, return a concrete "other"
 /// discriminator for it.
 llvm::ConstantInt *CodeGenModule::getPointerAuthOtherDiscriminator(
@@ -59,7 +93,7 @@ uint16_t CodeGen::getPointerAuthDeclDiscriminator(CodeGenModule &CGM,
 /// Return the "other" decl-specific discriminator for the given decl.
 uint16_t
 CodeGenModule::getPointerAuthDeclDiscriminator(GlobalDecl Declaration) {
-  uint16_t &EntityHash = PtrAuthDiscriminatorHashes[Declaration];
+  uint16_t &EntityHash = getPointerAuthCaches().PtrAuthDiscriminatorHashes[Declaration];
 
   if (EntityHash == 0) {
     StringRef Name = getMangledName(Declaration);
@@ -477,16 +511,39 @@ CodeGen::getConstantSignedPointer(CodeGenModule &CGM, llvm::Constant *Pointer,
 /// If applicable, sign a given constant function pointer with the ABI rules for
 /// functionType.
 llvm::Constant *CodeGenModule::getFunctionPointer(llvm::Constant *Pointer,
-                                                  QualType FunctionType) {
+                                                  QualType FunctionType,
+                                                  GlobalDecl GD) {
   assert(FunctionType->isFunctionType() ||
          FunctionType->isFunctionReferenceType() ||
          FunctionType->isFunctionPointerType());
 
-  if (auto PointerAuth = getFunctionPointerAuthInfo(FunctionType))
-    return getConstantSignedPointer(
-        Pointer, PointerAuth.getKey(), /*StorageAddress=*/nullptr,
+  if (auto PointerAuth = getFunctionPointerAuthInfo(FunctionType)) {
+    // Check a cache that, for now, just has entries for functions signed
+    // with the standard function-pointer scheme.
+    // Cache function pointers based on their decl.  Anything without a decl is
+    // going to be a one-off that doesn't need to be cached anyway.
+    llvm::Constant **Entry = nullptr;
+    if (GD) {
+      const FunctionDecl *FD = cast<FunctionDecl>(GD.getDecl());
+      PointerAuthCachesTy::ByDeclCacheTy &Cache =
+        getPointerAuthCaches().ConstantSignedPointersByDecl;
+      Entry = &Cache[FD->getCanonicalDecl()];
+      if (*Entry)
+        return llvm::ConstantExpr::getBitCast(*Entry, Pointer->getType());
+    }
+
+    // If the cache misses, build a new constant.  It's not a *problem* to
+    // have more than one of these for a particular function, but it's nice
+    // to avoid it.
+    Pointer = getConstantSignedPointer(
+        Pointer, PointerAuth.getKey(), nullptr,
         cast_or_null<llvm::ConstantInt>(PointerAuth.getDiscriminator()));
 
+    // Store the result back into the cache, if any.
+    if (Entry)
+      *Entry = Pointer;
+  }
+
   return Pointer;
 }
 
@@ -522,12 +579,26 @@ CGPointerAuthInfo CodeGenModule::getMemberFunctionPointerAuthInfo(QualType FT) {
 }
 
 llvm::Constant *CodeGenModule::getMemberFunctionPointer(llvm::Constant *Pointer,
-                                                        QualType FT) {
-  if (CGPointerAuthInfo PointerAuth = getMemberFunctionPointerAuthInfo(FT))
-    return getConstantSignedPointer(
+                                                        QualType FT,
+                                                        const FunctionDecl *FD) {
+  if (CGPointerAuthInfo PointerAuth = getMemberFunctionPointerAuthInfo(FT)) {
+    llvm::Constant **Entry = nullptr;
+    if (FD) {
+      PointerAuthCachesTy::ByDeclCacheTy &Cache =
+        getPointerAuthCaches().SignedThunkPointers;
+      Entry = &Cache[FD->getCanonicalDecl()];
+      if (*Entry)
+        return llvm::ConstantExpr::getBitCast(*Entry, Pointer->getType());
+    }
+
+    Pointer = getConstantSignedPointer(
         Pointer, PointerAuth.getKey(), nullptr,
         cast_or_null<llvm::ConstantInt>(PointerAuth.getDiscriminator()));
 
+    if (Entry)
+      *Entry = Pointer;
+  }
+
   if (const auto *MFT = dyn_cast<MemberPointerType>(FT.getTypePtr())) {
     if (MFT->hasPointeeToToCFIUncheckedCalleeFunctionType())
       Pointer = llvm::NoCFIValue::get(cast<llvm::GlobalValue>(Pointer));
@@ -541,7 +612,7 @@ llvm::Constant *CodeGenModule::getMemberFunctionPointer(const FunctionDecl *FD,
   QualType FT = FD->getType();
   FT = getContext().getMemberPointerType(FT, /*Qualifier=*/std::nullopt,
                                          cast<CXXMethodDecl>(FD)->getParent());
-  return getMemberFunctionPointer(getRawFunctionPointer(FD, Ty), FT);
+  return getMemberFunctionPointer(getRawFunctionPointer(FD, Ty), FT, FD);
 }
 
 std::optional<PointerAuthQualifier>
@@ -615,6 +686,7 @@ CodeGenModule::getVTablePointerAuthentication(const CXXRecordDecl *Record) {
   if (!Record->getDefinition() || !Record->isPolymorphic())
     return std::nullopt;
 
+  auto &VTablePtrAuthInfos = getPointerAuthCaches().VTablePtrAuthInfos;
   auto Existing = VTablePtrAuthInfos.find(Record);
   std::optional<PointerAuthQualifier> Authentication;
   if (Existing != VTablePtrAuthInfos.end()) {
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 7064421..79954d5 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -555,7 +555,9 @@ CodeGenModule::CodeGenModule(ASTContext &C,
     checkDataLayoutConsistency(Context.getTargetInfo(), LLVMContext, LangOpts);
 }
 
-CodeGenModule::~CodeGenModule() {}
+CodeGenModule::~CodeGenModule() {
+  destroyPointerAuthCaches();
+}
 
 void CodeGenModule::createObjCRuntime() {
   // This is just isGNUFamily(), but we want to force implementors of
diff --git a/clang/lib/CodeGen/CodeGenModule.h b/clang/lib/CodeGen/CodeGenModule.h
index 705d9a3..6ab8e7f 100644
--- a/clang/lib/CodeGen/CodeGenModule.h
+++ b/clang/lib/CodeGen/CodeGenModule.h
@@ -669,10 +669,15 @@ private:
   std::pair<std::unique_ptr<CodeGenFunction>, const TopLevelStmtDecl *>
       GlobalTopLevelStmtBlockInFlight;
 
-  llvm::DenseMap<GlobalDecl, uint16_t> PtrAuthDiscriminatorHashes;
+  struct PointerAuthCachesTy;
+
+  // We use a raw pointer, and a manual destruction function here as we want
+  // to avoid fully specifying PointerAuthCachesTy outside of the pointer auth
+  // codegen.
+  mutable PointerAuthCachesTy* PointerAuthCaches = nullptr;
+  PointerAuthCachesTy& getPointerAuthCaches() const;
+  void destroyPointerAuthCaches();
 
-  llvm::DenseMap<const CXXRecordDecl *, std::optional<PointerAuthQualifier>>
-      VTablePtrAuthInfos;
   std::optional<PointerAuthQualifier>
   computeVTPointerAuthentication(const CXXRecordDecl *ThisClass);
 
@@ -1033,13 +1038,15 @@ public:
   /// to the given function.  This will apply a pointer signature if
   /// necessary.
   llvm::Constant *getFunctionPointer(llvm::Constant *Pointer,
-                                     QualType FunctionType);
+                                     QualType FunctionType,
+                                     GlobalDecl GD = GlobalDecl());
 
   llvm::Constant *getMemberFunctionPointer(const FunctionDecl *FD,
                                            llvm::Type *Ty = nullptr);
 
   llvm::Constant *getMemberFunctionPointer(llvm::Constant *Pointer,
-                                           QualType FT);
+                                           QualType FT,
+                                           const FunctionDecl *FD);
 
   CGPointerAuthInfo getFunctionPointerAuthInfo(QualType T);
 
diff --git a/clang/lib/CodeGen/ItaniumCXXABI.cpp b/clang/lib/CodeGen/ItaniumCXXABI.cpp
index 569fbe9..f6e164f 100644
--- a/clang/lib/CodeGen/ItaniumCXXABI.cpp
+++ b/clang/lib/CodeGen/ItaniumCXXABI.cpp
@@ -5186,7 +5186,7 @@ ItaniumCXXABI::getSignedVirtualMemberFunctionPointer(const CXXMethodDecl *MD) {
   llvm::Constant *thunk = getOrCreateVirtualFunctionPointerThunk(origMD);
   QualType funcType = CGM.getContext().getMemberPointerType(
       MD->getType(), /*Qualifier=*/std::nullopt, MD->getParent());
-  return CGM.getMemberFunctionPointer(thunk, funcType);
+  return CGM.getMemberFunctionPointer(thunk, funcType, MD);
 }
 
 void WebAssemblyCXXABI::emitBeginCatch(CodeGenFunction &CGF,