Separating the cache for the existence of address discrimination

Codegen needs to be able to test this, so we need to be able to
rely on queries to the ASTContext.

Author: ojhunt

clang/include/clang/AST/ASTContext.h

@@ -620,61 +620,25 @@ class ASTContext : public RefCountedBase<ASTContext> {
   ParameterIndexTable ParamIndices;
 
 public:
-  struct TypeRelocationInfo {
-    static TypeRelocationInfo invalid() {
-      TypeRelocationInfo Result;
-      Result.IsRelocatable = false;
-      Result.IsReplaceable = false;
-      return Result;
+  struct CXXRecordDeclRelocationInfo {
+    static CXXRecordDeclRelocationInfo NonRelocatable() {
+      return {false, false};
     }
 
-    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;
+    unsigned IsRelocatable;
+    unsigned IsReplaceable;
   };
-  std::optional<TypeRelocationInfo>
+  std::optional<CXXRecordDeclRelocationInfo>
   getRelocationInfoForCXXRecord(const CXXRecordDecl *) const;
-  void setRelocationInfoForCXXRecord(const CXXRecordDecl *, TypeRelocationInfo);
+  void setRelocationInfoForCXXRecord(const CXXRecordDecl *,
+                                     const CXXRecordDeclRelocationInfo &);
+  bool containsAddressDiscriminatedPointerAuth(QualType T);
 
 private:
-  llvm::DenseMap<const CXXRecordDecl *, TypeRelocationInfo> RelocationInfo;
+  llvm::DenseMap<const CXXRecordDecl *, CXXRecordDeclRelocationInfo>
+      RelocatableClasses;
+  llvm::DenseMap<const RecordDecl *, bool>
+      RecordContainsAddressDiscriminatedPointerAuth;
 
   ImportDecl *FirstLocalImport = nullptr;
   ImportDecl *LastLocalImport = nullptr;

clang/include/clang/Sema/Sema.h

@@ -4343,10 +4343,10 @@ class Sema final : public SemaBase {
   void ActOnTagFinishDefinition(Scope *S, Decl *TagDecl,
                                 SourceRange BraceRange);
 
-  ASTContext::TypeRelocationInfo
-  GetCXX2CTypeRelocationInfo(const clang::CXXRecordDecl *D);
-
-  ASTContext::TypeRelocationInfo GetCXX2CTypeRelocationInfo(QualType T);
+  using CXXRecordDeclRelocationInfo = ASTContext::CXXRecordDeclRelocationInfo;
+  CXXRecordDeclRelocationInfo
+  CheckCXX2CRelocatableAndReplaceable(const clang::CXXRecordDecl *D);
+  CXXRecordDeclRelocationInfo CheckCXX2CRelocatableAndReplaceable(QualType T);
 
   void ActOnTagFinishSkippedDefinition(SkippedDefinitionContext Context);
 

clang/lib/AST/ASTContext.cpp

@@ -1687,22 +1687,56 @@ void ASTContext::getOverriddenMethods(
   Overridden.append(OverDecls.begin(), OverDecls.end());
 }
 
-std::optional<ASTContext::TypeRelocationInfo>
+std::optional<ASTContext::CXXRecordDeclRelocationInfo>
 ASTContext::getRelocationInfoForCXXRecord(const CXXRecordDecl *RD) const {
   assert(RD);
   CXXRecordDecl *D = RD->getDefinition();
-  auto it = RelocationInfo.find(D);
-  if (it != RelocationInfo.end())
+  auto it = RelocatableClasses.find(D);
+  if (it != RelocatableClasses.end())
     return it->getSecond();
   return std::nullopt;
 }
 
-void ASTContext::setRelocationInfoForCXXRecord(const CXXRecordDecl *RD,
-                                               TypeRelocationInfo Info) {
+void ASTContext::setRelocationInfoForCXXRecord(
+    const CXXRecordDecl *RD, const CXXRecordDeclRelocationInfo &Info) {
   assert(RD);
   CXXRecordDecl *D = RD->getDefinition();
-  assert(RelocationInfo.find(D) == RelocationInfo.end());
-  RelocationInfo.insert({D, Info});
+  assert(RelocatableClasses.find(D) == RelocatableClasses.end());
+  RelocatableClasses.insert({D, Info});
+}
+
+bool ASTContext::containsAddressDiscriminatedPointerAuth(QualType T) {
+  if (!LangOpts.PointerAuthCalls && !LangOpts.PointerAuthIntrinsics)
+    return false;
+
+  T = T.getCanonicalType();
+  if (T.hasAddressDiscriminatedPointerAuth())
+    return true;
+  const RecordDecl *RD = T->getAsRecordDecl();
+  if (!RD)
+    return false;
+
+  auto SaveReturn = [this, RD](bool Result) {
+    RecordContainsAddressDiscriminatedPointerAuth.insert({RD, Result});
+    return Result;
+  };
+  if (auto Existing = RecordContainsAddressDiscriminatedPointerAuth.find(RD);
+      Existing != RecordContainsAddressDiscriminatedPointerAuth.end())
+    return Existing->second;
+  if (const CXXRecordDecl *CXXRD = dyn_cast<CXXRecordDecl>(RD)) {
+    if (CXXRD->isPolymorphic() &&
+        hasAddressDiscriminatedVTableAuthentication(CXXRD))
+      return SaveReturn(true);
+    for (auto Base : CXXRD->bases()) {
+      if (containsAddressDiscriminatedPointerAuth(Base.getType()))
+        return SaveReturn(true);
+    }
+  }
+  for (auto *FieldDecl : RD->fields()) {
+    if (containsAddressDiscriminatedPointerAuth(FieldDecl->getType()))
+      return SaveReturn(true);
+  }
+  return SaveReturn(false);
 }
 
 void ASTContext::addedLocalImportDecl(ImportDecl *Import) {

clang/lib/Sema/SemaDeclCXX.cpp

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

clang/lib/Sema/SemaTypeTraits.cpp

@@ -226,13 +226,13 @@ static bool IsEligibleForReplacement(Sema &SemaRef, const CXXRecordDecl *D) {
   return !D->hasDeletedDestructor();
 }
 
-ASTContext::TypeRelocationInfo
-Sema::GetCXX2CTypeRelocationInfo(const CXXRecordDecl *D) {
+ASTContext::CXXRecordDeclRelocationInfo
+Sema::CheckCXX2CRelocatableAndReplaceable(const CXXRecordDecl *D) {
   if (auto ExistingInfo = Context.getRelocationInfoForCXXRecord(D))
     return *ExistingInfo;
 
   if (!getLangOpts().CPlusPlus || D->isInvalidDecl())
-    return ASTContext::TypeRelocationInfo::invalid();
+    return CXXRecordDeclRelocationInfo::NonRelocatable();
 
   assert(D->hasDefinition());
 
@@ -260,9 +260,8 @@ Sema::GetCXX2CTypeRelocationInfo(const CXXRecordDecl *D) {
     return *Is;
   };
 
-  ASTContext::TypeRelocationInfo Info;
-
-  Info.updateRelocatable([&] {
+  ASTContext::CXXRecordDeclRelocationInfo Info;
+  Info.IsRelocatable = [&] {
     if (D->isDependentType())
       return false;
 
@@ -280,9 +279,9 @@ Sema::GetCXX2CTypeRelocationInfo(const CXXRecordDecl *D) {
 
     // is default-movable.
     return IsDefaultMovable();
-  }());
+  }();
 
-  Info.updateReplaceable([&] {
+  Info.IsReplaceable = [&] {
     if (D->isDependentType())
       return false;
 
@@ -300,42 +299,45 @@ Sema::GetCXX2CTypeRelocationInfo(const CXXRecordDecl *D) {
 
     // is default-movable.
     return IsDefaultMovable();
-  }());
-
-  Info.setIsUnion(D->isUnion());
+  }();
 
   bool PtrauthMatters = LangOpts.PointerAuthIntrinsics ||
                         LangOpts.PointerAuthVTPtrAddressDiscrimination;
   if (PtrauthMatters) {
-    auto IsBottomRelocationInfo =
-        [](const ASTContext::TypeRelocationInfo &Info) {
-          return !Info.isReplaceable() && !Info.isRelocatable() &&
-                 Info.hasAddressDiscriminatedPointerAuth();
-        };
+    bool IsUnion = D->isUnion();
+    auto RecordPointerAuth = [&](bool HasAddressDiscrimination) {
+      if (HasAddressDiscrimination && IsUnion) {
+        Info.IsRelocatable = false;
+        Info.IsReplaceable = false;
+      }
+    };
+    auto IsBottomRelocationInfo = [](const CXXRecordDeclRelocationInfo &Info) {
+      return !Info.IsReplaceable && !Info.IsRelocatable;
+    };
 
     if (D->isPolymorphic())
-      Info.updateContainsAddressDiscriminatedValues(
-          Context.hasAddressDiscriminatedVTableAuthentication(D));
+      RecordPointerAuth(Context.hasAddressDiscriminatedVTableAuthentication(D));
     for (auto Base : D->bases()) {
       if (IsBottomRelocationInfo(Info))
         break;
-      bool BaseHasPtrauth = GetCXX2CTypeRelocationInfo(Base.getType())
-                                .hasAddressDiscriminatedPointerAuth();
-      Info.updateContainsAddressDiscriminatedValues(BaseHasPtrauth);
+      bool BaseHasPtrauth =
+          Context.containsAddressDiscriminatedPointerAuth(Base.getType());
+      RecordPointerAuth(BaseHasPtrauth);
     }
     for (auto *FieldDecl : D->fields()) {
       if (IsBottomRelocationInfo(Info))
         break;
-      bool FieldHasPtrauth = GetCXX2CTypeRelocationInfo(FieldDecl->getType())
-                                 .hasAddressDiscriminatedPointerAuth();
-      Info.updateContainsAddressDiscriminatedValues(FieldHasPtrauth);
+      bool FieldHasPtrauth =
+          Context.containsAddressDiscriminatedPointerAuth(FieldDecl->getType());
+      RecordPointerAuth(FieldHasPtrauth);
     }
   }
   Context.setRelocationInfoForCXXRecord(D, Info);
   return Info;
 }
 
-ASTContext::TypeRelocationInfo Sema::GetCXX2CTypeRelocationInfo(QualType T) {
+ASTContext::CXXRecordDeclRelocationInfo
+Sema::CheckCXX2CRelocatableAndReplaceable(QualType T) {
   T = T.getCanonicalType();
   enum class DirectRelocationInformation { Yes, No, Unknown };
   DirectRelocationInformation Relocatable =
@@ -356,10 +358,8 @@ ASTContext::TypeRelocationInfo Sema::GetCXX2CTypeRelocationInfo(QualType T) {
       Replaceable = DRI;
   };
   auto UpdateAddressDiscrimination = [&](DirectRelocationInformation DRI) {
-    if (ContainsAddressDiscriminatedValues ==
-            DirectRelocationInformation::Unknown ||
-        ContainsAddressDiscriminatedValues == DirectRelocationInformation::No)
-      ContainsAddressDiscriminatedValues = DRI;
+    if (ContainsAddressDiscriminatedValues == DirectRelocationInformation::Yes)
+      Replaceable = DirectRelocationInformation::No;
   };
 
   if (T->isVariableArrayType()) {
@@ -393,12 +393,9 @@ ASTContext::TypeRelocationInfo Sema::GetCXX2CTypeRelocationInfo(QualType T) {
   if (BaseElementType->isVectorType())
     UpdateRelocatable(DirectRelocationInformation::Yes);
 
-  auto CreateInfo = [=]() -> ASTContext::TypeRelocationInfo {
-    ASTContext::TypeRelocationInfo Result;
+  auto CreateInfo = [=]() -> CXXRecordDeclRelocationInfo {
     return {Relocatable == DirectRelocationInformation::Yes,
-            Replaceable == DirectRelocationInformation::Yes,
-            ContainsAddressDiscriminatedValues ==
-                DirectRelocationInformation::Yes};
+            Replaceable == DirectRelocationInformation::Yes};
   };
 
   if (BaseElementType->isIncompleteType())
@@ -408,20 +405,18 @@ ASTContext::TypeRelocationInfo Sema::GetCXX2CTypeRelocationInfo(QualType T) {
   if (!RD)
     return CreateInfo();
 
-  ASTContext::TypeRelocationInfo Info = GetCXX2CTypeRelocationInfo(RD);
-  Info.updateRelocatable(Relocatable != DirectRelocationInformation::No);
-  Info.updateReplaceable(Replaceable != DirectRelocationInformation::No);
-  Info.updateContainsAddressDiscriminatedValues(
-      ContainsAddressDiscriminatedValues == DirectRelocationInformation::Yes);
+  CXXRecordDeclRelocationInfo Info = CheckCXX2CRelocatableAndReplaceable(RD);
+  Info.IsRelocatable &= Relocatable != DirectRelocationInformation::No;
+  Info.IsReplaceable &= Replaceable != DirectRelocationInformation::No;
   return Info;
 }
 
 bool Sema::IsCXXTriviallyRelocatableType(QualType Type) {
-  return GetCXX2CTypeRelocationInfo(Type).isRelocatable();
+  return CheckCXX2CRelocatableAndReplaceable(Type).IsRelocatable;
 }
 
 bool Sema::IsCXXReplaceableType(QualType Type) {
-  return GetCXX2CTypeRelocationInfo(Type).isReplaceable();
+  return CheckCXX2CRelocatableAndReplaceable(Type).IsReplaceable;
 }
 
 /// Checks that type T is not a VLA.
@@ -722,7 +717,8 @@ static bool isTriviallyEqualityComparableType(Sema &S, QualType Type,
 }
 
 static bool IsTriviallyRelocatableType(Sema &SemaRef, QualType T) {
-  QualType BaseElementType = SemaRef.getASTContext().getBaseElementType(T);
+  ASTContext &Ctx = SemaRef.getASTContext();
+  QualType BaseElementType = Ctx.getBaseElementType(T);
 
   if (BaseElementType->isIncompleteType())
     return false;
@@ -732,13 +728,14 @@ static bool IsTriviallyRelocatableType(Sema &SemaRef, QualType T) {
   if (T.hasAddressDiscriminatedPointerAuth())
     return false;
 
+  if (Ctx.containsAddressDiscriminatedPointerAuth(BaseElementType))
+    return false;
+
   if (const auto *RD = BaseElementType->getAsCXXRecordDecl();
       RD && !RD->isPolymorphic()) {
-    ASTContext::TypeRelocationInfo Info =
-        SemaRef.GetCXX2CTypeRelocationInfo(RD);
-    if (Info.hasAddressDiscriminatedPointerAuth())
-      return false;
-    if (Info.isRelocatable())
+    ASTContext::CXXRecordDeclRelocationInfo Info =
+        SemaRef.CheckCXX2CRelocatableAndReplaceable(RD);
+    if (Info.IsRelocatable)
       return true;
   }