Split partial specializations from the lookup table

Author: ChuanqiXu9

clang/include/clang/Serialization/ASTBitCodes.h

@@ -739,6 +739,8 @@ enum ASTRecordTypes {
   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

@@ -634,10 +634,16 @@ 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
@@ -654,8 +660,9 @@ 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.
@@ -690,9 +697,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.
@@ -1399,7 +1406,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

@@ -431,6 +431,7 @@ class ASTWriter : public ASTDeserializationListener,
   using SpecializationUpdateMap =
       llvm::MapVector<const NamedDecl *, SmallVector<const Decl *>>;
   SpecializationUpdateMap SpecializationsUpdates;
+  SpecializationUpdateMap PartialSpecializationsUpdates;
 
   using FirstLatestDeclMap = llvm::DenseMap<Decl *, Decl *>;
 
@@ -585,9 +586,9 @@ 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(const DeclContext *DC,
                                llvm::SmallVectorImpl<char> &LookupTable);
@@ -604,7 +605,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(RecordDataImpl &OffsetsRecord);
   void WriteDeclContextVisibleUpdate(const DeclContext *DC);
   void WriteFPPragmaOptions(const FPOptionsOverride &Opts);
@@ -633,6 +634,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

@@ -1321,13 +1321,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));
   }
 }
 
@@ -1420,14 +1418,15 @@ 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(
+  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);
@@ -1451,13 +1450,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;
 }
 
@@ -3565,6 +3565,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());
@@ -8156,11 +8169,11 @@ 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
@@ -8170,48 +8183,67 @@ 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) {
@@ -8392,10 +8424,17 @@ 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

clang/lib/Serialization/ASTReaderDecl.cpp

@@ -304,7 +304,8 @@ class ASTDeclReader : public DeclVisitor<ASTDeclReader, void> {
   static void markIncompleteDeclChainImpl(...);
 
   void ReadSpecializations(ModuleFile &M, Decl *D,
-                           llvm::BitstreamCursor &DeclsCursor);
+                           llvm::BitstreamCursor &DeclsCursor,
+                           bool IsPartial);
 
   void ReadFunctionDefinition(FunctionDecl *FD);
   void Visit(Decl *D);
@@ -2384,9 +2385,10 @@ void ASTDeclReader::VisitRequiresExprBodyDecl(RequiresExprBodyDecl *D) {
 }
 
 void ASTDeclReader::ReadSpecializations(ModuleFile &M, Decl *D,
-                                        llvm::BitstreamCursor &DeclsCursor) {
+                                        llvm::BitstreamCursor &DeclsCursor,
+                                        bool IsPartial) {
   uint64_t Offset = ReadLocalOffset();
-  bool Failed = Reader.ReadSpecializations(M, DeclsCursor, Offset, D);
+  bool Failed = Reader.ReadSpecializations(M, DeclsCursor, Offset, D, IsPartial);
   (void)Failed;
   assert(!Failed);
 }
@@ -2429,7 +2431,8 @@ void ASTDeclReader::VisitClassTemplateDecl(ClassTemplateDecl *D) {
   if (ThisDeclID == Redecl.getFirstID()) {
     // This ClassTemplateDecl owns a CommonPtr; read it to keep track of all of
     // the specializations.
-    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor);
+    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor, /*IsPartial=*/false);
+    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor, /*IsPartial=*/true);
   }
 
   if (D->getTemplatedDecl()->TemplateOrInstantiation) {
@@ -2455,7 +2458,8 @@ void ASTDeclReader::VisitVarTemplateDecl(VarTemplateDecl *D) {
   if (ThisDeclID == Redecl.getFirstID()) {
     // This VarTemplateDecl owns a CommonPtr; read it to keep track of all of
     // the specializations.
-    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor);
+    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor, /*IsPartial=*/false);
+    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor, /*IsPartial=*/true);
   }
 }
 
@@ -2554,7 +2558,7 @@ void ASTDeclReader::VisitFunctionTemplateDecl(FunctionTemplateDecl *D) {
 
   if (ThisDeclID == Redecl.getFirstID()) {
     // This FunctionTemplateDecl owns a CommonPtr; read it.
-    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor);
+    ReadSpecializations(*Loc.F, D, Loc.F->DeclsCursor, /*IsPartial=*/false);
   }
 }
 
@@ -3845,6 +3849,7 @@ Decl *ASTReader::ReadDeclRecord(GlobalDeclID ID) {
   case DECL_CONTEXT_LEXICAL:
   case DECL_CONTEXT_VISIBLE:
   case DECL_SPECIALIZATIONS:
+  case DECL_PARTIAL_SPECIALIZATIONS:
     llvm_unreachable("Record cannot be de-serialized with readDeclRecord");
   case DECL_TYPEDEF:
     D = TypedefDecl::CreateDeserialized(Context, ID);
@@ -4332,7 +4337,17 @@ void ASTReader::loadDeclUpdateRecords(PendingUpdateRecord &Record) {
     PendingSpecializationsUpdates.erase(I);
 
     for (const auto &Update : SpecializationUpdates)
-      AddSpecializations(D, Update.Data, *Update.Mod);
+      AddSpecializations(D, Update.Data, *Update.Mod, /*IsPartial=*/false);
+  }
+
+  // Load the pending specializations update for this decl, if it has any.
+  if (auto I = PendingPartialSpecializationsUpdates.find(ID);
+      I != PendingPartialSpecializationsUpdates.end()) {
+    auto SpecializationUpdates = std::move(I->second);
+    PendingPartialSpecializationsUpdates.erase(I);
+
+    for (const auto &Update : SpecializationUpdates)
+      AddSpecializations(D, Update.Data, *Update.Mod, /*IsPartial=*/true);
   }
 }