Split partial specializations from the lookup table

Author: ChuanqiXu9

clang/include/clang/AST/DeclTemplate.h

@@ -806,13 +806,6 @@ class RedeclarableTemplateDecl : public TemplateDecl,
     /// was explicitly specialized.
     llvm::PointerIntPair<RedeclarableTemplateDecl *, 1, bool>
         InstantiatedFromMember;
-
-    /// If non-null, points to an array of specializations (including
-    /// partial specializations) known only by their external declaration IDs.
-    ///
-    /// The first value in the array is the number of specializations/partial
-    /// specializations that follow.
-    LazySpecializationInfo *LazySpecializations = nullptr;
   };
 
   /// Pointer to the common data shared by all declarations of this

clang/include/clang/Serialization/ASTBitCodes.h

@@ -737,8 +737,10 @@ enum ASTRecordTypes {
 
   /// Record code for updated specialization
   UPDATE_SPECIALIZATION = 73,
-  
+
   CXX_ADDED_TEMPLATE_SPECIALIZATION = 74,
+
+  CXX_ADDED_TEMPLATE_PARTIAL_SPECIALIZATION = 75,
 };
 
 /// Record types used within a source manager block.
@@ -1505,11 +1507,14 @@ enum DeclCode {
   /// A HLSLBufferDecl record.
   DECL_HLSL_BUFFER,
 
+  /// An ImplicitConceptSpecializationDecl record.
+  DECL_IMPLICIT_CONCEPT_SPECIALIZATION,
+
   // A decls specilization record.
   DECL_SPECIALIZATIONS,
 
-  /// An ImplicitConceptSpecializationDecl record.
-  DECL_IMPLICIT_CONCEPT_SPECIALIZATION,
+  // A decls specilization record.
+  DECL_PARTIAL_SPECIALIZATIONS,
 
   DECL_LAST = DECL_IMPLICIT_CONCEPT_SPECIALIZATION
 };

clang/include/clang/Serialization/ASTReader.h

@@ -635,10 +635,19 @@ class ASTReader
   llvm::DenseMap<const DeclContext *,
                  serialization::reader::DeclContextLookupTable> Lookups;
 
+  using SpecLookupTableTy =
+      llvm::DenseMap<const Decl *,
+                     serialization::reader::LazySpecializationInfoLookupTable>;
   /// Map from decls to specialized decls.
-  llvm::DenseMap<const Decl *,
-                 serialization::reader::LazySpecializationInfoLookupTable>
-      SpecializationsLookups;
+  SpecLookupTableTy SpecializationsLookups;
+  /// Split partial specialization from specialization to speed up lookups.
+  SpecLookupTableTy PartialSpecializationsLookups;
+
+  bool LoadExternalSpecializationsImpl(SpecLookupTableTy &SpecLookups,
+                                       const Decl *D);
+  bool LoadExternalSpecializationsImpl(SpecLookupTableTy &SpecLookups,
+                                       const Decl *D,
+                                       ArrayRef<TemplateArgument> TemplateArgs);
 
   // Updates for visible decls can occur for other contexts than just the
   // TU, and when we read those update records, the actual context may not
@@ -655,8 +664,10 @@ class ASTReader
   llvm::DenseMap<GlobalDeclID, DeclContextVisibleUpdates> PendingVisibleUpdates;
 
   using SpecializationsUpdate = SmallVector<UpdateData, 1>;
-  llvm::DenseMap<GlobalDeclID, SpecializationsUpdate>
-      PendingSpecializationsUpdates;
+  using SpecializationsUpdateMap =
+      llvm::DenseMap<GlobalDeclID, SpecializationsUpdate>;
+  SpecializationsUpdateMap PendingSpecializationsUpdates;
+  SpecializationsUpdateMap PendingPartialSpecializationsUpdates;
 
   /// The set of C++ or Objective-C classes that have forward
   /// declarations that have not yet been linked to their definitions.
@@ -691,9 +702,9 @@ class ASTReader
                                      uint64_t Offset, GlobalDeclID ID);
 
   bool ReadSpecializations(ModuleFile &M, llvm::BitstreamCursor &Cursor,
-                           uint64_t Offset, Decl *D);
+                           uint64_t Offset, Decl *D, bool IsPartial);
   void AddSpecializations(const Decl *D, const unsigned char *Data,
-                          ModuleFile &M);
+                          ModuleFile &M, bool IsPartial);
 
   /// A vector containing identifiers that have already been
   /// loaded.
@@ -1439,7 +1450,10 @@ class ASTReader
   /// Get the loaded specializations lookup tables for \p D,
   /// if any.
   serialization::reader::LazySpecializationInfoLookupTable *
-  getLoadedSpecializationsLookupTables(const Decl *D);
+  getLoadedSpecializationsLookupTables(const Decl *D, bool IsPartial);
+
+  /// If we have any unloaded specialization for \p D
+  bool haveUnloadedSpecializations(const Decl *D) const;
 
 private:
   struct ImportedModule {

clang/include/clang/Serialization/ASTWriter.h

@@ -428,6 +428,7 @@ class ASTWriter : public ASTDeserializationListener,
   using SpecializationUpdateMap =
       llvm::MapVector<const NamedDecl *, SmallVector<const Decl *>>;
   SpecializationUpdateMap SpecializationsUpdates;
+  SpecializationUpdateMap PartialSpecializationsUpdates;
 
   using FirstLatestDeclMap = llvm::DenseMap<Decl *, Decl *>;
 
@@ -580,9 +581,10 @@ class ASTWriter : public ASTDeserializationListener,
 
   void GenerateSpecializationInfoLookupTable(
       const NamedDecl *D, llvm::SmallVectorImpl<const Decl *> &Specializations,
-      llvm::SmallVectorImpl<char> &LookupTable);
+      llvm::SmallVectorImpl<char> &LookupTable, bool IsPartial);
   uint64_t WriteSpecializationInfoLookupTable(
-      const NamedDecl *D, llvm::SmallVectorImpl<const Decl *> &Specializations);
+      const NamedDecl *D, llvm::SmallVectorImpl<const Decl *> &Specializations,
+      bool IsPartial);
   void GenerateNameLookupTable(ASTContext &Context, const DeclContext *DC,
                                llvm::SmallVectorImpl<char> &LookupTable);
   uint64_t WriteDeclContextLexicalBlock(ASTContext &Context,
@@ -598,7 +600,7 @@ class ASTWriter : public ASTDeserializationListener,
   void WriteDeclAndTypes(ASTContext &Context);
   void PrepareWritingSpecialDecls(Sema &SemaRef);
   void WriteSpecialDeclRecords(Sema &SemaRef);
-  void WriteSpecializationsUpdates();
+  void WriteSpecializationsUpdates(bool IsPartial);
   void WriteDeclUpdatesBlocks(ASTContext &Context,
                               RecordDataImpl &OffsetsRecord);
   void WriteDeclContextVisibleUpdate(ASTContext &Context,
@@ -629,6 +631,7 @@ class ASTWriter : public ASTDeserializationListener,
   unsigned DeclObjCIvarAbbrev = 0;
   unsigned DeclCXXMethodAbbrev = 0;
   unsigned DeclSpecializationsAbbrev = 0;
+  unsigned DeclPartialSpecializationsAbbrev = 0;
 
   unsigned DeclDependentNonTemplateCXXMethodAbbrev = 0;
   unsigned DeclTemplateCXXMethodAbbrev = 0;

clang/lib/Serialization/ASTReader.cpp

@@ -1311,13 +1311,11 @@ void LazySpecializationInfoLookupTrait::ReadDataInto(internal_key_type,
   using namespace llvm::support;
 
   for (unsigned NumDecls =
-           DataLen / serialization::reader::LazySpecializationInfo::Length;
+           DataLen / sizeof(serialization::reader::LazySpecializationInfo);
        NumDecls; --NumDecls) {
     LocalDeclID LocalID =
         LocalDeclID::get(Reader, F, endian::readNext<DeclID, llvm::endianness::little, unaligned>(d));
-    const bool IsPartial =
-        endian::readNext<bool, llvm::endianness::little, unaligned>(d);
-    Val.insert({Reader.getGlobalDeclID(F, LocalID), IsPartial});
+    Val.insert(Reader.getGlobalDeclID(F, LocalID));
   }
 }
 
@@ -1410,14 +1408,16 @@ bool ASTReader::ReadVisibleDeclContextStorage(ModuleFile &M,
 }
 
 void ASTReader::AddSpecializations(const Decl *D, const unsigned char *Data,
-                                   ModuleFile &M) {
+                                   ModuleFile &M, bool IsPartial) {
   D = D->getCanonicalDecl();
-  SpecializationsLookups[D].Table.add(
-      &M, Data, reader::LazySpecializationInfoLookupTrait(*this, M));
+  auto &SpecLookups =
+      IsPartial ? PartialSpecializationsLookups : SpecializationsLookups;
+  SpecLookups[D].Table.add(&M, Data,
+                           reader::LazySpecializationInfoLookupTrait(*this, M));
 }
 
 bool ASTReader::ReadSpecializations(ModuleFile &M, BitstreamCursor &Cursor,
-                                    uint64_t Offset, Decl *D) {
+                                    uint64_t Offset, Decl *D, bool IsPartial) {
   assert(Offset != 0);
 
   SavedStreamPosition SavedPosition(Cursor);
@@ -1441,13 +1441,14 @@ bool ASTReader::ReadSpecializations(ModuleFile &M, BitstreamCursor &Cursor,
     return true;
   }
   unsigned RecCode = MaybeRecCode.get();
-  if (RecCode != DECL_SPECIALIZATIONS) {
+  if (RecCode != DECL_SPECIALIZATIONS &&
+      RecCode != DECL_PARTIAL_SPECIALIZATIONS) {
     Error("Expected decl specs block");
     return true;
   }
 
   auto *Data = (const unsigned char *)Blob.data();
-  AddSpecializations(D, Data, M);
+  AddSpecializations(D, Data, M, IsPartial);
   return false;
 }
 
@@ -3556,6 +3557,19 @@ llvm::Error ASTReader::ReadASTBlock(ModuleFile &F,
       break;
     }
 
+    case CXX_ADDED_TEMPLATE_PARTIAL_SPECIALIZATION: {
+      unsigned Idx = 0;
+      GlobalDeclID ID = ReadDeclID(F, Record, Idx);
+      auto *Data = (const unsigned char *)Blob.data();
+      PendingPartialSpecializationsUpdates[ID].push_back(UpdateData{&F, Data});
+      // If we've already loaded the decl, perform the updates when we finish
+      // loading this block.
+      if (Decl *D = GetExistingDecl(ID))
+        PendingUpdateRecords.push_back(
+            PendingUpdateRecord(ID, D, /*JustLoaded=*/false));
+      break;
+    }
+
     case IDENTIFIER_TABLE:
       F.IdentifierTableData =
           reinterpret_cast<const unsigned char *>(Blob.data());
@@ -8136,11 +8150,12 @@ Stmt *ASTReader::GetExternalDeclStmt(uint64_t Offset) {
   return ReadStmtFromStream(*Loc.F);
 }
 
-bool ASTReader::LoadExternalSpecializations(const Decl *D, bool OnlyPartial) {
+bool ASTReader::LoadExternalSpecializationsImpl(SpecLookupTableTy &SpecLookups,
+                                                const Decl *D) {
   assert(D);
 
-  auto It = SpecializationsLookups.find(D);
-  if (It == SpecializationsLookups.end())
+  auto It = SpecLookups.find(D);
+  if (It == SpecLookups.end())
     return false;
 
   // Get Decl may violate the iterator from SpecializationsLookups so we store
@@ -8150,48 +8165,71 @@ bool ASTReader::LoadExternalSpecializations(const Decl *D, bool OnlyPartial) {
 
   // Since we've loaded all the specializations, we can erase it from
   // the lookup table.
-  if (!OnlyPartial)
-    SpecializationsLookups.erase(It);
+  SpecLookups.erase(It);
 
   bool NewSpecsFound = false;
   Deserializing LookupResults(this);
-  for (auto &Info : Infos)
-    if (!OnlyPartial || Info.IsPartial) {
-      if (GetExistingDecl(Info.ID))
-        continue;
-      NewSpecsFound = true;
-      GetDecl(Info.ID);
-    }
+  for (auto &Info : Infos) {
+    if (GetExistingDecl(Info))
+      continue;
+    NewSpecsFound = true;
+    GetDecl(Info);
+  }
 
   return NewSpecsFound;
 }
 
-bool ASTReader::LoadExternalSpecializations(
-    const Decl *D, ArrayRef<TemplateArgument> TemplateArgs) {
+bool ASTReader::LoadExternalSpecializations(const Decl *D, bool OnlyPartial) {
   assert(D);
 
-  auto It = SpecializationsLookups.find(D);
-  if (It == SpecializationsLookups.end())
+  bool NewSpecsFound =
+      LoadExternalSpecializationsImpl(PartialSpecializationsLookups, D);
+  if (OnlyPartial)
+    return NewSpecsFound;
+
+  NewSpecsFound |= LoadExternalSpecializationsImpl(SpecializationsLookups, D);
+  return NewSpecsFound;
+}
+
+bool ASTReader::LoadExternalSpecializationsImpl(
+    SpecLookupTableTy &SpecLookups, const Decl *D,
+    ArrayRef<TemplateArgument> TemplateArgs) {
+  assert(D);
+
+  auto It = SpecLookups.find(D);
+  if (It == SpecLookups.end())
     return false;
 
   Deserializing LookupResults(this);
   auto HashValue = StableHashForTemplateArguments(TemplateArgs);
 
-  // Get Decl may violate the iterator from SpecializationsLookups
+  // Get Decl may violate the iterator from SpecLookups
   llvm::SmallVector<serialization::reader::LazySpecializationInfo, 8> Infos =
       It->second.Table.find(HashValue);
 
   bool NewSpecsFound = false;
   for (auto &Info : Infos) {
-    if (GetExistingDecl(Info.ID))
+    if (GetExistingDecl(Info))
       continue;
     NewSpecsFound = true;
-    GetDecl(Info.ID);
+    GetDecl(Info);
   }
 
   return NewSpecsFound;
 }
 
+bool ASTReader::LoadExternalSpecializations(
+    const Decl *D, ArrayRef<TemplateArgument> TemplateArgs) {
+  assert(D);
+
+  bool NewDeclsFound = LoadExternalSpecializationsImpl(
+      PartialSpecializationsLookups, D, TemplateArgs);
+  NewDeclsFound |=
+      LoadExternalSpecializationsImpl(SpecializationsLookups, D, TemplateArgs);
+
+  return NewDeclsFound;
+}
+
 void ASTReader::FindExternalLexicalDecls(
     const DeclContext *DC, llvm::function_ref<bool(Decl::Kind)> IsKindWeWant,
     SmallVectorImpl<Decl *> &Decls) {
@@ -8372,10 +8410,19 @@ ASTReader::getLoadedLookupTables(DeclContext *Primary) const {
 }
 
 serialization::reader::LazySpecializationInfoLookupTable *
-ASTReader::getLoadedSpecializationsLookupTables(const Decl *D) {
+ASTReader::getLoadedSpecializationsLookupTables(const Decl *D, bool IsPartial) {
+  assert(D->isCanonicalDecl());
+  auto &LookupTable =
+      IsPartial ? PartialSpecializationsLookups : SpecializationsLookups;
+  auto I = LookupTable.find(D);
+  return I == LookupTable.end() ? nullptr : &I->second;
+}
+
+bool ASTReader::haveUnloadedSpecializations(const Decl *D) const {
   assert(D->isCanonicalDecl());
-  auto I = SpecializationsLookups.find(D);
-  return I == SpecializationsLookups.end() ? nullptr : &I->second;
+  return (PartialSpecializationsLookups.find(D) !=
+          PartialSpecializationsLookups.end()) ||
+         (SpecializationsLookups.find(D) != SpecializationsLookups.end());
 }
 
 /// Under non-PCH compilation the consumer receives the objc methods