2 files changed, 55 insertions, 69 deletions

diff --git a/clang/lib/Serialization/ASTReader.cpp b/clang/lib/Serialization/ASTReader.cpp
index c12894d773f..acb19a2cc7e 100644
--- a/clang/lib/Serialization/ASTReader.cpp
+++ b/clang/lib/Serialization/ASTReader.cpp
@@ -973,9 +973,13 @@ bool ASTReader::ReadLexicalDeclContextStorage(ModuleFile &M,
     return true;
   }
 
-  M.DeclContextInfos[DC].LexicalDecls = llvm::makeArrayRef(
-      reinterpret_cast<const KindDeclIDPair *>(Blob.data()),
-      Blob.size() / sizeof(KindDeclIDPair));
+  assert(!isa<TranslationUnitDecl>(DC) &&
+         "expected a TU_UPDATE_LEXICAL record for TU");
+  // FIXME: Once we remove RewriteDecl, assert that we didn't already have
+  // lexical decls for this context.
+  LexicalDecls[DC] = llvm::makeArrayRef(
+      reinterpret_cast<const llvm::support::unaligned_uint32_t *>(Blob.data()),
+      Blob.size() / 4);
   DC->setHasExternalLexicalStorage(true);
   return false;
 }
@@ -2498,10 +2502,11 @@ ASTReader::ReadASTBlock(ModuleFile &F, unsigned ClientLoadCapabilities) {
         
     case TU_UPDATE_LEXICAL: {
       DeclContext *TU = Context.getTranslationUnitDecl();
-      DeclContextInfo &Info = F.DeclContextInfos[TU];
-      Info.LexicalDecls = llvm::makeArrayRef(
-          reinterpret_cast<const KindDeclIDPair *>(Blob.data()),
-          static_cast<unsigned int>(Blob.size() / sizeof(KindDeclIDPair)));
+      LexicalContents Contents(
+          reinterpret_cast<const llvm::support::unaligned_uint32_t *>(
+              Blob.data()),
+          static_cast<unsigned int>(Blob.size() / 4));
+      TULexicalDecls.push_back(std::make_pair(&F, Contents));
       TU->setHasExternalLexicalStorage(true);
       break;
     }
@@ -6175,69 +6180,45 @@ Stmt *ASTReader::GetExternalDeclStmt(uint64_t Offset) {
   return ReadStmtFromStream(*Loc.F);
 }
 
-namespace {
-  class FindExternalLexicalDeclsVisitor {
-    ASTReader &Reader;
-    const DeclContext *DC;
-    llvm::function_ref<bool(Decl::Kind)> IsKindWeWant;
-    
-    SmallVectorImpl<Decl*> &Decls;
-    bool PredefsVisited[NUM_PREDEF_DECL_IDS];
-
-  public:
-    FindExternalLexicalDeclsVisitor(
-        ASTReader &Reader, const DeclContext *DC,
-        llvm::function_ref<bool(Decl::Kind)> IsKindWeWant,
-        SmallVectorImpl<Decl *> &Decls)
-        : Reader(Reader), DC(DC), IsKindWeWant(IsKindWeWant), Decls(Decls) {
-      for (unsigned I = 0; I != NUM_PREDEF_DECL_IDS; ++I)
-        PredefsVisited[I] = false;
-    }
+void ASTReader::FindExternalLexicalDecls(
+    const DeclContext *DC, llvm::function_ref<bool(Decl::Kind)> IsKindWeWant,
+    SmallVectorImpl<Decl *> &Decls) {
+  bool PredefsVisited[NUM_PREDEF_DECL_IDS] = {};
 
-    static bool visitPostorder(ModuleFile &M, void *UserData) {
-      return (*static_cast<FindExternalLexicalDeclsVisitor*>(UserData))(M);
-    }
+  auto Visit = [&] (ModuleFile *M, const LexicalContents &LexicalDecls) {
+    assert(LexicalDecls.size() % 2 == 0 && "expected an even number of entries");
+    for (int I = 0, N = LexicalDecls.size(); I != N; I += 2) {
+      auto K = (Decl::Kind)+LexicalDecls[I];
+      if (!IsKindWeWant(K))
+        continue;
 
-    bool operator()(ModuleFile &M) {
-      ModuleFile::DeclContextInfosMap::iterator Info =
-          M.DeclContextInfos.find(DC);
-      if (Info == M.DeclContextInfos.end() || Info->second.LexicalDecls.empty())
-        return false;
+      auto ID = (serialization::DeclID)+LexicalDecls[I + 1];
 
-      // Load all of the declaration IDs
-      for (const KindDeclIDPair &P : Info->second.LexicalDecls) {
-        if (!IsKindWeWant((Decl::Kind)P.first))
+      // Don't add predefined declarations to the lexical context more
+      // than once.
+      if (ID < NUM_PREDEF_DECL_IDS) {
+        if (PredefsVisited[ID])
           continue;
 
-        // Don't add predefined declarations to the lexical context more
-        // than once.
-        if (P.second < NUM_PREDEF_DECL_IDS) {
-          if (PredefsVisited[P.second])
-            continue;
-
-          PredefsVisited[P.second] = true;
-        }
-
-        if (Decl *D = Reader.GetLocalDecl(M, P.second)) {
-          if (!DC->isDeclInLexicalTraversal(D))
-            Decls.push_back(D);
-        }
+        PredefsVisited[ID] = true;
       }
 
-      return false;
+      if (Decl *D = GetLocalDecl(*M, ID)) {
+        if (!DC->isDeclInLexicalTraversal(D))
+          Decls.push_back(D);
+      }
     }
   };
-}
 
-void ASTReader::FindExternalLexicalDecls(
-    const DeclContext *DC, llvm::function_ref<bool(Decl::Kind)> IsKindWeWant,
-    SmallVectorImpl<Decl *> &Decls) {
-  // There might be lexical decls in multiple modules, for the TU at
-  // least. FIXME: Only look in multiple module files in the very rare
-  // cases where this can actually happen.
-  FindExternalLexicalDeclsVisitor Visitor(*this, DC, IsKindWeWant, Decls);
-  ModuleMgr.visitDepthFirst(
-      nullptr, &FindExternalLexicalDeclsVisitor::visitPostorder, &Visitor);
+  if (isa<TranslationUnitDecl>(DC)) {
+    for (auto Lexical : TULexicalDecls)
+      Visit(Lexical.first, Lexical.second);
+  } else {
+    auto I = LexicalDecls.find(DC);
+    if (I != LexicalDecls.end())
+      Visit(getOwningModuleFile(cast<Decl>(DC)), I->second);
+  }
+
   ++NumLexicalDeclContextsRead;
 }
 
diff --git a/clang/lib/Serialization/ASTWriter.cpp b/clang/lib/Serialization/ASTWriter.cpp
index 1199803a9c0..366d4b39fc5 100644
--- a/clang/lib/Serialization/ASTWriter.cpp
+++ b/clang/lib/Serialization/ASTWriter.cpp
@@ -2738,12 +2738,15 @@ uint64_t ASTWriter::WriteDeclContextLexicalBlock(ASTContext &Context,
   uint64_t Offset = Stream.GetCurrentBitNo();
   RecordData Record;
   Record.push_back(DECL_CONTEXT_LEXICAL);
-  SmallVector<KindDeclIDPair, 64> Decls;
-  for (const auto *D : DC->decls())
-    Decls.push_back(std::make_pair(D->getKind(), GetDeclRef(D)));
+  SmallVector<uint32_t, 128> KindDeclPairs;
+  for (const auto *D : DC->decls()) {
+    KindDeclPairs.push_back(D->getKind());
+    KindDeclPairs.push_back(GetDeclRef(D));
+  }
 
   ++NumLexicalDeclContexts;
-  Stream.EmitRecordWithBlob(DeclContextLexicalAbbrev, Record, bytes(Decls));
+  Stream.EmitRecordWithBlob(DeclContextLexicalAbbrev, Record,
+                            bytes(KindDeclPairs));
   return Offset;
 }
 
@@ -4288,10 +4291,12 @@ void ASTWriter::WriteASTCore(Sema &SemaRef,
   // Create a lexical update block containing all of the declarations in the
   // translation unit that do not come from other AST files.
   const TranslationUnitDecl *TU = Context.getTranslationUnitDecl();
-  SmallVector<KindDeclIDPair, 64> NewGlobalDecls;
-  for (const auto *I : TU->noload_decls()) {
-    if (!I->isFromASTFile())
-      NewGlobalDecls.push_back(std::make_pair(I->getKind(), GetDeclRef(I)));
+  SmallVector<uint32_t, 128> NewGlobalKindDeclPairs;
+  for (const auto *D : TU->noload_decls()) {
+    if (!D->isFromASTFile()) {
+      NewGlobalKindDeclPairs.push_back(D->getKind());
+      NewGlobalKindDeclPairs.push_back(GetDeclRef(D));
+    }
   }
   
   llvm::BitCodeAbbrev *Abv = new llvm::BitCodeAbbrev();
@@ -4301,7 +4306,7 @@ void ASTWriter::WriteASTCore(Sema &SemaRef,
   Record.clear();
   Record.push_back(TU_UPDATE_LEXICAL);
   Stream.EmitRecordWithBlob(TuUpdateLexicalAbbrev, Record,
-                            bytes(NewGlobalDecls));
+                            bytes(NewGlobalKindDeclPairs));
   
   // And a visible updates block for the translation unit.
   Abv = new llvm::BitCodeAbbrev();