[clang] Fix PointerAuth semantics of __builtin_is_cpp_trivially_relocatable

This reworks the way we compute relocatability and replaceability of
types. We do this by having a single interface the provides a full
`TypeRelocationInfo` for a QualType. This simplifies the reasoning
for all cases, and reduces duplication of code.

We still limit our caching of results to CXXRecordDecls but i think
it might be worth caching on canonical types in future, but for now
I didn't want to change the caching semantics any further.

TypeRelocationInfo no longer permits raw field access as we can then
ensure we maintain correct state over pointer auth and similar.

As a basic optimization we gate the additional required work on the
existence pointer auth options.

Author: ojhunt

clang/include/clang/AST/ASTContext.h

@@ -620,18 +620,61 @@ class ASTContext : public RefCountedBase<ASTContext> {
   ParameterIndexTable ParamIndices;
 
 public:
-  struct CXXRecordDeclRelocationInfo {
-    unsigned IsRelocatable;
-    unsigned IsReplaceable;
+  struct TypeRelocationInfo {
+    static TypeRelocationInfo invalid() {
+      TypeRelocationInfo Result;
+      Result.IsRelocatable = false;
+      Result.IsReplaceable = false;
+      return Result;
+    }
+
+    TypeRelocationInfo() {}
+    TypeRelocationInfo(bool IsRelocatable, bool IsReplaceable,
+                       bool HasAddressDiscriminatedValues)
+        : IsUnion(false), IsRelocatable(IsRelocatable),
+          IsReplaceable(IsReplaceable),
+          HasAddressDiscriminatedValues(HasAddressDiscriminatedValues) {}
+
+    void setIsUnion(bool Union) {
+      IsUnion = Union;
+      normalizeState();
+    }
+    void updateRelocatable(bool Relocatable) { IsRelocatable &= Relocatable; }
+    void updateReplaceable(bool Replaceable) { IsReplaceable &= Replaceable; }
+    void
+    updateContainsAddressDiscriminatedValues(bool AddressDiscriminatedValues) {
+      HasAddressDiscriminatedValues |= AddressDiscriminatedValues;
+      normalizeState();
+    }
+    void mergeWithSiblingMember(bool IsUnion, const TypeRelocationInfo &Other) {
+      updateRelocatable(Other.IsRelocatable);
+      updateReplaceable(Other.IsReplaceable);
+      normalizeState();
+    }
+    bool isRelocatable() const { return IsRelocatable; }
+    bool isReplaceable() const { return IsReplaceable; }
+    bool hasAddressDiscriminatedPointerAuth() const {
+      return HasAddressDiscriminatedValues;
+    }
+
+  private:
+    void normalizeState() {
+      if (!IsUnion || !HasAddressDiscriminatedValues)
+        return;
+      IsRelocatable = false;
+      IsReplaceable = false;
+    }
+    bool IsUnion = false;
+    bool IsRelocatable = true;
+    bool IsReplaceable = true;
+    bool HasAddressDiscriminatedValues = false;
   };
-  std::optional<CXXRecordDeclRelocationInfo>
+  std::optional<TypeRelocationInfo>
   getRelocationInfoForCXXRecord(const CXXRecordDecl *) const;
-  void setRelocationInfoForCXXRecord(const CXXRecordDecl *,
-                                     CXXRecordDeclRelocationInfo);
+  void setRelocationInfoForCXXRecord(const CXXRecordDecl *, TypeRelocationInfo);
 
 private:
-  llvm::DenseMap<const CXXRecordDecl *, CXXRecordDeclRelocationInfo>
-      RelocatableClasses;
+  llvm::DenseMap<const CXXRecordDecl *, TypeRelocationInfo> RelocationInfo;
 
   ImportDecl *FirstLocalImport = nullptr;
   ImportDecl *LastLocalImport = nullptr;
@@ -3668,6 +3711,7 @@ OPT_LIST(V)
   /// authentication policy for the specified record.
   const CXXRecordDecl *
   baseForVTableAuthentication(const CXXRecordDecl *ThisClass);
+  bool hasAddressDiscriminatedVTableAuthentication(const CXXRecordDecl *Class);
   bool useAbbreviatedThunkName(GlobalDecl VirtualMethodDecl,
                                StringRef MangledName);
 

clang/include/clang/Sema/Sema.h

@@ -4343,8 +4343,10 @@ class Sema final : public SemaBase {
   void ActOnTagFinishDefinition(Scope *S, Decl *TagDecl,
                                 SourceRange BraceRange);
 
-  ASTContext::CXXRecordDeclRelocationInfo
-  CheckCXX2CRelocatableAndReplaceable(const clang::CXXRecordDecl *D);
+  ASTContext::TypeRelocationInfo
+  GetCXX2CTypeRelocationInfo(const clang::CXXRecordDecl *D);
+
+  ASTContext::TypeRelocationInfo GetCXX2CTypeRelocationInfo(QualType T);
 
   void ActOnTagFinishSkippedDefinition(SkippedDefinitionContext Context);
 
@@ -8680,7 +8682,6 @@ class Sema final : public SemaBase {
   // FIXME: This is in Sema because it requires
   // overload resolution, can we move to ASTContext?
   bool IsCXXTriviallyRelocatableType(QualType T);
-  bool IsCXXTriviallyRelocatableType(const CXXRecordDecl &RD);
 
   //// Determines if a type is replaceable
   /// according to the C++26 rules.

clang/lib/AST/ASTContext.cpp

@@ -1687,22 +1687,22 @@ void ASTContext::getOverriddenMethods(
   Overridden.append(OverDecls.begin(), OverDecls.end());
 }
 
-std::optional<ASTContext::CXXRecordDeclRelocationInfo>
+std::optional<ASTContext::TypeRelocationInfo>
 ASTContext::getRelocationInfoForCXXRecord(const CXXRecordDecl *RD) const {
   assert(RD);
   CXXRecordDecl *D = RD->getDefinition();
-  auto it = RelocatableClasses.find(D);
-  if (it != RelocatableClasses.end())
+  auto it = RelocationInfo.find(D);
+  if (it != RelocationInfo.end())
     return it->getSecond();
   return std::nullopt;
 }
 
-void ASTContext::setRelocationInfoForCXXRecord(
-    const CXXRecordDecl *RD, CXXRecordDeclRelocationInfo Info) {
+void ASTContext::setRelocationInfoForCXXRecord(const CXXRecordDecl *RD,
+                                               TypeRelocationInfo Info) {
   assert(RD);
   CXXRecordDecl *D = RD->getDefinition();
-  assert(RelocatableClasses.find(D) == RelocatableClasses.end());
-  RelocatableClasses.insert({D, Info});
+  assert(RelocationInfo.find(D) == RelocationInfo.end());
+  RelocationInfo.insert({D, Info});
 }
 
 void ASTContext::addedLocalImportDecl(ImportDecl *Import) {
@@ -15121,6 +15121,21 @@ ASTContext::baseForVTableAuthentication(const CXXRecordDecl *ThisClass) {
   return PrimaryBase;
 }
 
+bool ASTContext::hasAddressDiscriminatedVTableAuthentication(
+    const CXXRecordDecl *Class) {
+  assert(Class->isPolymorphic());
+  const CXXRecordDecl *BaseType = baseForVTableAuthentication(Class);
+  using AuthAttr = VTablePointerAuthenticationAttr;
+  const auto *ExplicitAuth = BaseType->getAttr<AuthAttr>();
+  if (!ExplicitAuth)
+    return LangOpts.PointerAuthVTPtrAddressDiscrimination;
+  AuthAttr::AddressDiscriminationMode AddressDiscrimination =
+      ExplicitAuth->getAddressDiscrimination();
+  if (AddressDiscrimination == AuthAttr::DefaultAddressDiscrimination)
+    return LangOpts.PointerAuthVTPtrAddressDiscrimination;
+  return AddressDiscrimination == AuthAttr::AddressDiscrimination;
+}
+
 bool ASTContext::useAbbreviatedThunkName(GlobalDecl VirtualMethodDecl,
                                          StringRef MangledName) {
   auto *Method = cast<CXXMethodDecl>(VirtualMethodDecl.getDecl());

clang/lib/Sema/SemaDeclCXX.cpp

@@ -10615,7 +10615,7 @@ void Sema::checkIllFormedTrivialABIStruct(CXXRecordDecl &RD) {
       }
   }
 
-  if (IsCXXTriviallyRelocatableType(RD))
+  if (GetCXX2CTypeRelocationInfo(&RD).isRelocatable())
     return;
 
   // Ill-formed if the copy and move constructors are deleted.