ASTUnit/CIndex: Explicitly track the top-level decls when using an ASTUnit made

from a source file.
 - This allows CIndex to avoid iterating over all the top-level decls when using
   a PCH, which means we deserialize far fewer decls.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@90559 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 71 insertions, 14 deletions

diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index 3bc815446f..29d7b1aa2f 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/include/clang/Frontend/ASTUnit.h
@@ -19,6 +19,8 @@
 #include "clang/Basic/FileManager.h"
 #include "clang/Index/ASTLocation.h"
 #include <string>
+#include <vector>
+#include <cassert>
 
 namespace clang {
 class ASTContext;
@@ -45,21 +47,30 @@ class ASTUnit {
   llvm::OwningPtr<Preprocessor>     PP;
   llvm::OwningPtr<ASTContext>       Ctx;
   bool                              tempFile;
-  
+
   // OnlyLocalDecls - when true, walking this AST should only visit declarations
   // that come from the AST itself, not from included precompiled headers.
   // FIXME: This is temporary; eventually, CIndex will always do this.
   bool                              OnlyLocalDecls;
 
-  // Track whether the main file was loaded from an AST or not.
+  /// Track whether the main file was loaded from an AST or not.
   bool MainFileIsAST;
 
+  /// Track the top-level decls which appeared in an ASTUnit which was loaded
+  /// from a source file.
+  //
+  // FIXME: This is just an optimization hack to avoid deserializing large parts
+  // of a PCH file when using the Index library on an ASTUnit loaded from
+  // source. In the long term we should make the Index library use efficient and
+  // more scalable search mechanisms.
+  std::vector<Decl*> TopLevelDecls;
+
   /// The name of the original source file used to generate this ASTUnit.
   std::string OriginalSourceFile;
 
   // Critical optimization when using clang_getCursor().
   ASTLocation LastLoc;
-  
+
   ASTUnit(const ASTUnit&); // DO NOT IMPLEMENT
   ASTUnit &operator=(const ASTUnit &); // DO NOT IMPLEMENT
 
@@ -80,17 +91,26 @@ public:
 
   const FileManager &getFileManager() const { return FileMgr; }
         FileManager &getFileManager()       { return FileMgr; }
-  
+
   const std::string &getOriginalSourceFileName();
   const std::string &getPCHFileName();
 
   void unlinkTemporaryFile() { tempFile = true; }
-  
+
   bool getOnlyLocalDecls() const { return OnlyLocalDecls; }
-  
+
   void setLastASTLocation(ASTLocation ALoc) { LastLoc = ALoc; }
   ASTLocation getLastASTLocation() const { return LastLoc; }
-  
+
+  std::vector<Decl*> &getTopLevelDecls() {
+    assert(!isMainFileAST() && "Invalid call for AST based ASTUnit!");
+    return TopLevelDecls;
+  }
+  const std::vector<Decl*> &getTopLevelDecls() const {
+    assert(!isMainFileAST() && "Invalid call for AST based ASTUnit!");
+    return TopLevelDecls;
+  }
+
   /// \brief Create a ASTUnit from a PCH file.
   ///
   /// \param Filename - The PCH file to load.
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index 2afbcd133e..8c2731b895 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -180,13 +180,30 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
 
 namespace {
 
-class NullAction : public ASTFrontendAction {
+class TopLevelDeclTrackerConsumer : public ASTConsumer {
+  ASTUnit &Unit;
+
+public:
+  TopLevelDeclTrackerConsumer(ASTUnit &_Unit) : Unit(_Unit) {}
+
+  void HandleTopLevelDecl(DeclGroupRef D) {
+    for (DeclGroupRef::iterator it = D.begin(), ie = D.end(); it != ie; ++it)
+      Unit.getTopLevelDecls().push_back(*it);
+  }
+};
+
+class TopLevelDeclTrackerAction : public ASTFrontendAction {
+public:
+  ASTUnit &Unit;
+
   virtual ASTConsumer *CreateASTConsumer(CompilerInstance &CI,
                                          llvm::StringRef InFile) {
-    return new ASTConsumer();
+    return new TopLevelDeclTrackerConsumer(Unit);
   }
 
 public:
+  TopLevelDeclTrackerAction(ASTUnit &_Unit) : Unit(_Unit) {}
+
   virtual bool hasCodeCompletionSupport() const { return false; }
 };
 
@@ -198,7 +215,7 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocation(const CompilerInvocation &CI,
   // Create the compiler instance to use for building the AST.
   CompilerInstance Clang;
   llvm::OwningPtr<ASTUnit> AST;
-  NullAction Act;
+  llvm::OwningPtr<TopLevelDeclTrackerAction> Act;
 
   Clang.getInvocation() = CI;
 
@@ -237,11 +254,12 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocation(const CompilerInvocation &CI,
   // Create the preprocessor.
   Clang.createPreprocessor();
 
-  if (!Act.BeginSourceFile(Clang, Clang.getFrontendOpts().Inputs[0].second,
+  Act.reset(new TopLevelDeclTrackerAction(*AST));
+  if (!Act->BeginSourceFile(Clang, Clang.getFrontendOpts().Inputs[0].second,
                            /*IsAST=*/false))
     goto error;
 
-  Act.Execute();
+  Act->Execute();
 
   // Steal the created target, context, and preprocessor, and take back the
   // source and file managers.
@@ -251,7 +269,7 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocation(const CompilerInvocation &CI,
   Clang.takeFileManager();
   AST->Target.reset(Clang.takeTarget());
 
-  Act.EndSourceFile();
+  Act->EndSourceFile();
 
   Clang.takeDiagnosticClient();
   Clang.takeDiagnostics();
diff --git a/test/Index/cindex-from-source.m b/test/Index/cindex-from-source.m
new file mode 100644
index 0000000000..8d6a1de4bb
--- /dev/null
+++ b/test/Index/cindex-from-source.m
@@ -0,0 +1,9 @@
+// RUN: echo 'typedef int t0;' > %t.pfx.h
+// RUN: clang -x objective-c-header %t.pfx.h -o %t.pfx.h.gch
+// RUN: c-index-test -test-load-source local %s -include %t.pfx.h > %t
+// RUN: FileCheck %s < %t
+// CHECK: cindex-from-source.m:{{.*}}:{{.*}}: StructDecl=s0:{{.*}}:{{.*}} [Context=cindex-from-source.m]
+// CHECK: cindex-from-source.m:{{.*}}:{{.*}}: VarDecl=g0:{{.*}}:{{.*}} [Context=cindex-from-source.m]
+
+struct s0 {};
+t0 g0;
diff --git a/tools/CIndex/CIndex.cpp b/tools/CIndex/CIndex.cpp
index 432f4a09ee..12c3935ca1 100644
--- a/tools/CIndex/CIndex.cpp
+++ b/tools/CIndex/CIndex.cpp
@@ -613,7 +613,17 @@ void clang_loadTranslationUnit(CXTranslationUnit CTUnit,
   }
 
   TUVisitor DVisit(CTUnit, callback, CData, PCHLevel);
-  DVisit.Visit(Ctx.getTranslationUnitDecl());
+
+  // If using a non-AST based ASTUnit, iterate over the stored list of top-level
+  // decls.
+  if (!CXXUnit->isMainFileAST() && CXXUnit->getOnlyLocalDecls()) {
+    const std::vector<Decl*> &TLDs = CXXUnit->getTopLevelDecls();
+    for (std::vector<Decl*>::const_iterator it = TLDs.begin(),
+           ie = TLDs.end(); it != ie; ++it) {
+      DVisit.Visit(*it);
+    }
+  } else
+    DVisit.Visit(Ctx.getTranslationUnitDecl());
 }
 
 void clang_loadDeclaration(CXDecl Dcl,