4 files changed, 55 insertions, 42 deletions

diff --git a/include/clang/AST/ExternalASTSource.h b/include/clang/AST/ExternalASTSource.h
index def9ced94c..4a6377eaa9 100644
--- a/include/clang/AST/ExternalASTSource.h
+++ b/include/clang/AST/ExternalASTSource.h
@@ -58,6 +58,20 @@ public:
 
   virtual ~ExternalASTSource();
 
+  /// \brief RAII class for safely pairing a StartedDeserializing call
+  /// with FinishedDeserializing.
+  class Deserializing {
+    ExternalASTSource *Source;
+  public:
+    explicit Deserializing(ExternalASTSource *source) : Source(source) {
+      assert(Source);
+      Source->StartedDeserializing();
+    }
+    ~Deserializing() {
+      Source->FinishedDeserializing();
+    }
+  };
+
   /// \brief Resolve a declaration ID into a declaration, potentially
   /// building a new declaration.
   ///
@@ -100,6 +114,19 @@ public:
   virtual bool FindExternalLexicalDecls(const DeclContext *DC,
                                 llvm::SmallVectorImpl<Decl*> &Result) = 0;
 
+  /// \brief Notify ExternalASTSource that we started deserialization of
+  /// a decl or type so until FinishedDeserializing is called there may be
+  /// decls that are initializing. Must be paired with FinishedDeserializing.
+  ///
+  /// The default implementation of this method is a no-op.
+  virtual void StartedDeserializing() { }
+
+  /// \brief Notify ExternalASTSource that we finished the deserialization of
+  /// a decl or type. Must be paired with StartedDeserializing.
+  ///
+  /// The default implementation of this method is a no-op.
+  virtual void FinishedDeserializing() { }
+
   /// \brief Function that will be invoked when we begin parsing a new
   /// translation unit involving this external AST source.
   ///
diff --git a/include/clang/Frontend/PCHReader.h b/include/clang/Frontend/PCHReader.h
index 8e8cf25fb9..57af01033a 100644
--- a/include/clang/Frontend/PCHReader.h
+++ b/include/clang/Frontend/PCHReader.h
@@ -468,26 +468,9 @@ private:
 
   /// Number of visible decl contexts read/total.
   unsigned NumVisibleDeclContextsRead, TotalVisibleDeclContexts;
-
-  /// \brief When a type or declaration is being loaded from the PCH file, an
-  /// instantance of this RAII object will be available on the stack to
-  /// indicate when we are in a recursive-loading situation.
-  class LoadingTypeOrDecl {
-    PCHReader &Reader;
-    LoadingTypeOrDecl *Parent;
-
-    LoadingTypeOrDecl(const LoadingTypeOrDecl&); // do not implement
-    LoadingTypeOrDecl &operator=(const LoadingTypeOrDecl&); // do not implement
-
-  public:
-    explicit LoadingTypeOrDecl(PCHReader &Reader);
-    ~LoadingTypeOrDecl();
-  };
-  friend class LoadingTypeOrDecl;
-
-  /// \brief If we are currently loading a type or declaration, points to the
-  /// most recent LoadingTypeOrDecl object on the stack.
-  LoadingTypeOrDecl *CurrentlyLoadingTypeOrDecl;
+  
+  /// \brief Number of Decl/types that are currently deserializing.
+  unsigned NumCurrentElementsDeserializing;
 
   /// \brief An IdentifierInfo that has been loaded but whose top-level
   /// declarations of the same name have not (yet) been loaded.
@@ -757,6 +740,15 @@ public:
   virtual bool FindExternalLexicalDecls(const DeclContext *DC,
                                         llvm::SmallVectorImpl<Decl*> &Decls);
 
+  /// \brief Notify PCHReader that we started deserialization of
+  /// a decl or type so until FinishedDeserializing is called there may be
+  /// decls that are initializing. Must be paired with FinishedDeserializing.
+  virtual void StartedDeserializing() { ++NumCurrentElementsDeserializing; }
+
+  /// \brief Notify PCHReader that we finished the deserialization of
+  /// a decl or type. Must be paired with StartedDeserializing.
+  virtual void FinishedDeserializing();
+
   /// \brief Function that will be invoked when we begin parsing a new
   /// translation unit involving this external AST source.
   ///
diff --git a/lib/Frontend/PCHReader.cpp b/lib/Frontend/PCHReader.cpp
index 6fa7294f87..5a239e4f04 100644
--- a/lib/Frontend/PCHReader.cpp
+++ b/lib/Frontend/PCHReader.cpp
@@ -426,7 +426,7 @@ PCHReader::PCHReader(Preprocessor &PP, ASTContext *Context,
     TotalNumStatements(0), NumMacrosRead(0), NumMethodPoolSelectorsRead(0),
     NumMethodPoolMisses(0), TotalNumMacros(0), NumLexicalDeclContextsRead(0),
     TotalLexicalDeclContexts(0), NumVisibleDeclContextsRead(0),
-    TotalVisibleDeclContexts(0), CurrentlyLoadingTypeOrDecl(0) {
+    TotalVisibleDeclContexts(0), NumCurrentElementsDeserializing(0) {
   RelocatablePCH = false;
 }
 
@@ -443,7 +443,7 @@ PCHReader::PCHReader(SourceManager &SourceMgr, FileManager &FileMgr,
     TotalNumStatements(0), NumMacrosRead(0), NumMethodPoolSelectorsRead(0),
     NumMethodPoolMisses(0), TotalNumMacros(0), NumLexicalDeclContextsRead(0),
     TotalLexicalDeclContexts(0), NumVisibleDeclContextsRead(0),
-    TotalVisibleDeclContexts(0), CurrentlyLoadingTypeOrDecl(0) {
+    TotalVisibleDeclContexts(0), NumCurrentElementsDeserializing(0) {
   RelocatablePCH = false;
 }
 
@@ -2242,7 +2242,7 @@ QualType PCHReader::ReadTypeRecord(unsigned Index) {
   ReadingKindTracker ReadingKind(Read_Type, *this);
   
   // Note that we are loading a type record.
-  LoadingTypeOrDecl Loading(*this);
+  Deserializing AType(this);
 
   DeclsCursor.JumpToBit(Loc.second);
   RecordData Record;
@@ -3235,7 +3235,7 @@ void
 PCHReader::SetGloballyVisibleDecls(IdentifierInfo *II,
                               const llvm::SmallVectorImpl<uint32_t> &DeclIDs,
                                    bool Nonrecursive) {
-  if (CurrentlyLoadingTypeOrDecl && !Nonrecursive) {
+  if (NumCurrentElementsDeserializing && !Nonrecursive) {
     PendingIdentifierInfos.push_back(PendingIdentifierInfo());
     PendingIdentifierInfo &PII = PendingIdentifierInfos.back();
     PII.II = II;
@@ -3677,28 +3677,22 @@ void PCHReader::SetLabelOf(AddrLabelExpr *S, unsigned ID) {
   }
 }
 
-
-PCHReader::LoadingTypeOrDecl::LoadingTypeOrDecl(PCHReader &Reader)
-  : Reader(Reader), Parent(Reader.CurrentlyLoadingTypeOrDecl) {
-  Reader.CurrentlyLoadingTypeOrDecl = this;
-}
-
-PCHReader::LoadingTypeOrDecl::~LoadingTypeOrDecl() {
-  if (!Parent) {
+void PCHReader::FinishedDeserializing() {
+  assert(NumCurrentElementsDeserializing &&
+         "FinishedDeserializing not paired with StartedDeserializing");
+  if (NumCurrentElementsDeserializing == 1) {
     // If any identifiers with corresponding top-level declarations have
     // been loaded, load those declarations now.
-    while (!Reader.PendingIdentifierInfos.empty()) {
-      Reader.SetGloballyVisibleDecls(Reader.PendingIdentifierInfos.front().II,
-                                 Reader.PendingIdentifierInfos.front().DeclIDs,
-                                     true);
-      Reader.PendingIdentifierInfos.pop_front();
+    while (!PendingIdentifierInfos.empty()) {
+      SetGloballyVisibleDecls(PendingIdentifierInfos.front().II,
+                              PendingIdentifierInfos.front().DeclIDs, true);
+      PendingIdentifierInfos.pop_front();
     }
 
     // We are not in recursive loading, so it's safe to pass the "interesting"
     // decls to the consumer.
-    if (Reader.Consumer)
-      Reader.PassInterestingDeclsToConsumer();
+    if (Consumer)
+      PassInterestingDeclsToConsumer();
   }
-
-  Reader.CurrentlyLoadingTypeOrDecl = Parent;
+  --NumCurrentElementsDeserializing;
 }
diff --git a/lib/Frontend/PCHReaderDecl.cpp b/lib/Frontend/PCHReaderDecl.cpp
index 03e7e1a093..e9428cc0b4 100644
--- a/lib/Frontend/PCHReaderDecl.cpp
+++ b/lib/Frontend/PCHReaderDecl.cpp
@@ -1314,7 +1314,7 @@ Decl *PCHReader::ReadDeclRecord(unsigned Index) {
   ReadingKindTracker ReadingKind(Read_Decl, *this);
 
   // Note that we are loading a declaration record.
-  LoadingTypeOrDecl Loading(*this);
+  Deserializing ADecl(this);
 
   DeclsCursor.JumpToBit(Loc.second);
   RecordData Record;