When using a precompiled preamble with detailed preprocessing records,
trap the serialized preprocessing records (macro definitions, macro
instantiations, macro definitions) from the generation of the
precompiled preamble, then replay those when walking the list of
preprocessed entities. This eliminates a bug where clang_getCursor()
wasn't able to find preprocessed-entity cursors in the preamble.


git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@120396 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index c7db42d..4e27cf3 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/include/clang/Frontend/ASTUnit.h
@@ -110,6 +110,14 @@
   // more scalable search mechanisms.
   std::vector<Decl*> TopLevelDecls;
 
+  /// \brief The list of preprocessed entities which appeared when the ASTUnit
+  /// was loaded.
+  ///
+  /// FIXME: This is just an optimization hack to avoid deserializing large
+  /// parts of a PCH file while performing a walk or search. In the long term,
+  /// we should provide more scalable search mechanisms.
+  std::vector<PreprocessedEntity *> PreprocessedEntities;
+  
   /// The name of the original source file used to generate this ASTUnit.
   std::string OriginalSourceFile;
 
@@ -215,6 +223,10 @@
   /// declarations parsed within the precompiled preamble.
   std::vector<serialization::DeclID> TopLevelDeclsInPreamble;
 
+  /// \brief A list of the offsets into the precompiled preamble which
+  /// correspond to preprocessed entities.
+  std::vector<uint64_t> PreprocessedEntitiesInPreamble;
+  
   /// \brief Whether we should be caching code-completion results.
   bool ShouldCacheCodeCompletionResults;
   
@@ -317,7 +329,8 @@
                                                      bool AllowRebuild = true,
                                                         unsigned MaxLines = 0);
   void RealizeTopLevelDeclsFromPreamble();
-
+  void RealizePreprocessedEntitiesFromPreamble();
+  
 public:
   class ConcurrencyCheck {
     volatile ASTUnit &Self;
@@ -426,6 +439,17 @@
     TopLevelDeclsInPreamble.push_back(D);
   }
 
+  typedef std::vector<PreprocessedEntity *>::iterator pp_entity_iterator;
+  
+  pp_entity_iterator pp_entity_begin();
+  pp_entity_iterator pp_entity_end();
+  
+  /// \brief Add a new preprocessed entity that's stored at the given offset
+  /// in the precompiled preamble.
+  void addPreprocessedEntityFromPreamble(uint64_t Offset) {
+    PreprocessedEntitiesInPreamble.push_back(Offset);
+  }
+  
   /// \brief Retrieve the mapping from File IDs to the preprocessed entities
   /// within that file.
   PreprocessedEntitiesByFileMap &getPreprocessedEntitiesByFile() {
diff --git a/include/clang/Lex/PreprocessingRecord.h b/include/clang/Lex/PreprocessingRecord.h
index 01ac79d..5a3b85d 100644
--- a/include/clang/Lex/PreprocessingRecord.h
+++ b/include/clang/Lex/PreprocessingRecord.h
@@ -248,6 +248,9 @@
     /// \brief Read any preallocated preprocessed entities from the external
     /// source.
     virtual void ReadPreprocessedEntities() = 0;
+    
+    /// \brief Read the preprocessed entity at the given offset.
+    virtual PreprocessedEntity *ReadPreprocessedEntity(uint64_t Offset) = 0;
   };
   
   /// \brief A record of the steps taken while preprocessing a source file,
@@ -302,6 +305,11 @@
     void SetExternalSource(ExternalPreprocessingRecordSource &Source,
                            unsigned NumPreallocatedEntities);
 
+    /// \brief Retrieve the external source for preprocessed entities.
+    ExternalPreprocessingRecordSource *getExternalSource() const {
+      return ExternalSource;
+    }
+    
     unsigned getNumPreallocatedEntities() const {
       return NumPreallocatedEntities;
     }
diff --git a/include/clang/Serialization/ASTReader.h b/include/clang/Serialization/ASTReader.h
index c1470ea..cc3497a 100644
--- a/include/clang/Serialization/ASTReader.h
+++ b/include/clang/Serialization/ASTReader.h
@@ -845,9 +845,12 @@
   /// build prior to including the precompiled header.
   const std::string &getSuggestedPredefines() { return SuggestedPredefines; }
       
-  /// \brief Read preprocessed entities into the 
+  /// \brief Read preprocessed entities into the preprocessing record.
   virtual void ReadPreprocessedEntities();
 
+  /// \brief Read the preprocessed entity at the given offset.
+  virtual PreprocessedEntity *ReadPreprocessedEntity(uint64_t Offset);
+
   void ReadUserDiagnosticMappings(Diagnostic &Diag);
 
   /// \brief Returns the number of source locations found in the chain.
@@ -1154,7 +1157,7 @@
   Expr *ReadSubExpr();
 
   /// \brief Reads the macro record located at the given offset.
-  void ReadMacroRecord(PerFileData &F, uint64_t Offset);
+  PreprocessedEntity *ReadMacroRecord(PerFileData &F, uint64_t Offset);
 
   /// \brief Note that the identifier is a macro whose record will be loaded
   /// from the given AST file at the given (file-local) offset.
diff --git a/include/clang/Serialization/ASTWriter.h b/include/clang/Serialization/ASTWriter.h
index 6b7884d..e0f1d96 100644
--- a/include/clang/Serialization/ASTWriter.h
+++ b/include/clang/Serialization/ASTWriter.h
@@ -23,6 +23,7 @@
 #include "clang/Sema/SemaConsumer.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Bitcode/BitstreamWriter.h"
 #include <map>
 #include <queue>
@@ -37,6 +38,7 @@
 namespace clang {
 
 class ASTContext;
+class ASTSerializationListener;
 class NestedNameSpecifier;
 class CXXBaseSpecifier;
 class CXXBaseOrMemberInitializer;
@@ -44,6 +46,7 @@
 class MacroDefinition;
 class MemorizeStatCalls;
 class ASTReader;
+class PreprocessedEntity;
 class Preprocessor;
 class Sema;
 class SourceManager;
@@ -70,6 +73,10 @@
   /// \brief The reader of existing AST files, if we're chaining.
   ASTReader *Chain;
 
+  /// \brief A listener object that receives notifications when certain 
+  /// entities are serialized.                    
+  ASTSerializationListener *SerializationListener;
+                    
   /// \brief Stores a declaration or a type to be written to the AST file.
   class DeclOrType {
   public:
@@ -334,6 +341,12 @@
   /// the given bitstream.
   ASTWriter(llvm::BitstreamWriter &Stream);
 
+  /// \brief Set the listener that will receive notification of serialization
+  /// events.
+  void SetSerializationListener(ASTSerializationListener *Listener) {
+    SerializationListener = Listener;
+  }
+                    
   /// \brief Write a precompiled header for the given semantic analysis.
   ///
   /// \param SemaRef a reference to the semantic analysis object that processed
@@ -573,6 +586,7 @@
   virtual void InitializeSema(Sema &S) { SemaPtr = &S; }
   virtual void HandleTranslationUnit(ASTContext &Ctx);
   virtual ASTMutationListener *GetASTMutationListener();
+  virtual ASTSerializationListener *GetASTSerializationListener();
   virtual ASTDeserializationListener *GetASTDeserializationListener();
 };
 
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index a43d0d3..586002d 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -27,6 +27,7 @@
 #include "clang/Frontend/FrontendOptions.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Serialization/ASTReader.h"
+#include "clang/Serialization/ASTSerializationListener.h"
 #include "clang/Serialization/ASTWriter.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Preprocessor.h"
@@ -637,10 +638,11 @@
   }
 };
 
-class PrecompilePreambleConsumer : public PCHGenerator {
+class PrecompilePreambleConsumer : public PCHGenerator, 
+                                   public ASTSerializationListener {
   ASTUnit &Unit;
   std::vector<Decl *> TopLevelDecls;
-
+                                     
 public:
   PrecompilePreambleConsumer(ASTUnit &Unit,
                              const Preprocessor &PP, bool Chaining,
@@ -672,6 +674,15 @@
                                       getWriter().getDeclID(TopLevelDecls[I]));
     }
   }
+                                     
+  virtual void SerializedPreprocessedEntity(PreprocessedEntity *Entity,
+                                            uint64_t Offset) {
+    Unit.addPreprocessedEntityFromPreamble(Offset);
+  }
+                                     
+  virtual ASTSerializationListener *GetASTSerializationListener() {
+    return this;
+  }
 };
 
 class PrecompilePreambleAction : public ASTFrontendAction {
@@ -757,6 +768,7 @@
   
   // Clear out old caches and data.
   TopLevelDecls.clear();
+  PreprocessedEntities.clear();
   CleanTemporaryFiles();
   PreprocessedEntitiesByFile.clear();
 
@@ -765,6 +777,7 @@
                     StoredDiagnostics.begin() + NumStoredDiagnosticsFromDriver,
                             StoredDiagnostics.end());
     TopLevelDeclsInPreamble.clear();
+    PreprocessedEntitiesInPreamble.clear();
   }
 
   // Create a file manager object to provide access to and cache the filesystem.
@@ -1237,6 +1250,8 @@
                           StoredDiagnostics.end());
   TopLevelDecls.clear();
   TopLevelDeclsInPreamble.clear();
+  PreprocessedEntities.clear();
+  PreprocessedEntitiesInPreamble.clear();
   
   // Create a file manager object to provide access to and cache the filesystem.
   Clang.setFileManager(new FileManager(Clang.getFileSystemOpts()));
@@ -1269,6 +1284,8 @@
     llvm::sys::Path(FrontendOpts.OutputFile).eraseFromDisk();
     Preamble.clear();
     TopLevelDeclsInPreamble.clear();
+    PreprocessedEntities.clear();
+    PreprocessedEntitiesInPreamble.clear();
     PreambleRebuildCounter = DefaultPreambleRebuildInterval;
     PreprocessorOpts.eraseRemappedFile(
                                PreprocessorOpts.remapped_file_buffer_end() - 1);
@@ -1321,6 +1338,55 @@
   TopLevelDecls.insert(TopLevelDecls.begin(), Resolved.begin(), Resolved.end());
 }
 
+void ASTUnit::RealizePreprocessedEntitiesFromPreamble() {
+  if (!PP)
+    return;
+  
+  PreprocessingRecord *PPRec = PP->getPreprocessingRecord();
+  if (!PPRec)
+    return;
+  
+  ExternalPreprocessingRecordSource *External = PPRec->getExternalSource();
+  if (!External)
+    return;
+
+  for (unsigned I = 0, N = PreprocessedEntitiesInPreamble.size(); I != N; ++I) {
+    if (PreprocessedEntity *PE
+          = External->ReadPreprocessedEntity(PreprocessedEntitiesInPreamble[I]))
+      PreprocessedEntities.push_back(PE);
+  }
+  
+  if (PreprocessedEntities.empty())
+    return;
+  
+  PreprocessedEntities.insert(PreprocessedEntities.end(), 
+                              PPRec->begin(true), PPRec->end(true));
+}
+
+ASTUnit::pp_entity_iterator ASTUnit::pp_entity_begin() {
+  if (!PreprocessedEntitiesInPreamble.empty() &&
+      PreprocessedEntities.empty())
+    RealizePreprocessedEntitiesFromPreamble();
+  
+  if (PreprocessedEntities.empty())
+    if (PreprocessingRecord *PPRec = PP->getPreprocessingRecord())
+      return PPRec->begin(true);
+  
+  return PreprocessedEntities.begin();
+}
+
+ASTUnit::pp_entity_iterator ASTUnit::pp_entity_end() {
+  if (!PreprocessedEntitiesInPreamble.empty() &&
+      PreprocessedEntities.empty())
+    RealizePreprocessedEntitiesFromPreamble();
+  
+  if (PreprocessedEntities.empty())
+    if (PreprocessingRecord *PPRec = PP->getPreprocessingRecord())
+      return PPRec->end(true);
+  
+  return PreprocessedEntities.end();
+}
+
 unsigned ASTUnit::getMaxPCHLevel() const {
   if (!getOnlyLocalDecls())
     return Decl::MaxPCHLevel;
diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index 1cac948..57002eb 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -1339,7 +1339,7 @@
   }
 }
 
-void ASTReader::ReadMacroRecord(PerFileData &F, uint64_t Offset) {
+PreprocessedEntity *ASTReader::ReadMacroRecord(PerFileData &F, uint64_t Offset) {
   assert(PP && "Forgot to set Preprocessor ?");
   llvm::BitstreamCursor &Stream = F.MacroCursor;
 
@@ -1356,14 +1356,14 @@
     unsigned Code = Stream.ReadCode();
     switch (Code) {
     case llvm::bitc::END_BLOCK:
-      return;
+      return 0;
 
     case llvm::bitc::ENTER_SUBBLOCK:
       // No known subblocks, always skip them.
       Stream.ReadSubBlockID();
       if (Stream.SkipBlock()) {
         Error("malformed block record in AST file");
-        return;
+        return 0;
       }
       continue;
 
@@ -1387,12 +1387,12 @@
       // of the definition of the macro we were looking for. We're
       // done.
       if (Macro)
-        return;
+        return 0;
 
       IdentifierInfo *II = DecodeIdentifierInfo(Record[0]);
       if (II == 0) {
         Error("macro must have a name in AST file");
-        return;
+        return 0;
       }
       SourceLocation Loc = ReadSourceLocation(F, Record[1]);
       bool isUsed = Record[2];
@@ -1459,16 +1459,16 @@
       // of the definition of the macro we were looking for. We're
       // done.
       if (Macro)
-        return;
+        return 0;
 
       if (!PP->getPreprocessingRecord()) {
         Error("missing preprocessing record in AST file");
-        return;
+        return 0;
       }
 
       PreprocessingRecord &PPRec = *PP->getPreprocessingRecord();
-      if (PPRec.getPreprocessedEntity(Record[0]))
-        return;
+      if (PreprocessedEntity *PE = PPRec.getPreprocessedEntity(Record[0]))
+        return PE;
 
       MacroInstantiation *MI
         = new (PPRec) MacroInstantiation(DecodeIdentifierInfo(Record[3]),
@@ -1476,7 +1476,7 @@
                                            ReadSourceLocation(F, Record[2])),
                                          getMacroDefinition(Record[4]));
       PPRec.SetPreallocatedEntity(Record[0], MI);
-      return;
+      return MI;
     }
 
     case PP_MACRO_DEFINITION: {
@@ -1484,20 +1484,20 @@
       // of the definition of the macro we were looking for. We're
       // done.
       if (Macro)
-        return;
+        return 0;
 
       if (!PP->getPreprocessingRecord()) {
         Error("missing preprocessing record in AST file");
-        return;
+        return 0;
       }
 
       PreprocessingRecord &PPRec = *PP->getPreprocessingRecord();
-      if (PPRec.getPreprocessedEntity(Record[0]))
-        return;
+      if (PreprocessedEntity *PE = PPRec.getPreprocessedEntity(Record[0]))
+        return PE;
 
       if (Record[1] > MacroDefinitionsLoaded.size()) {
         Error("out-of-bounds macro definition record");
-        return;
+        return 0;
       }
 
       // Decode the identifier info and then check again; if the macro is
@@ -1518,7 +1518,7 @@
           DeserializationListener->MacroDefinitionRead(Record[1], MD);
       }
 
-      return;
+      return MacroDefinitionsLoaded[Record[1] - 1];
     }
 
     case PP_INCLUSION_DIRECTIVE: {
@@ -1526,21 +1526,21 @@
       // of the definition of the macro we were looking for. We're
       // done.
       if (Macro)
-        return;
+        return 0;
 
       if (!PP->getPreprocessingRecord()) {
         Error("missing preprocessing record in AST file");
-        return;
+        return 0;
       }
 
       PreprocessingRecord &PPRec = *PP->getPreprocessingRecord();
-      if (PPRec.getPreprocessedEntity(Record[0]))
-        return;
+      if (PreprocessedEntity *PE = PPRec.getPreprocessedEntity(Record[0]))
+        return PE;
 
       const char *FullFileNameStart = BlobStart + Record[3];
       const FileEntry *File
-      = PP->getFileManager().getFile(llvm::StringRef(FullFileNameStart,
-                                                     BlobLen - Record[3]));
+        = PP->getFileManager().getFile(llvm::StringRef(FullFileNameStart,
+                                                       BlobLen - Record[3]));
 
       // FIXME: Stable encoding
       InclusionDirective::InclusionKind Kind
@@ -1553,10 +1553,12 @@
                                  SourceRange(ReadSourceLocation(F, Record[1]),
                                              ReadSourceLocation(F, Record[2])));
       PPRec.SetPreallocatedEntity(Record[0], ID);
-      return;
+      return ID;
     }
     }
   }
+  
+  return 0;
 }
 
 void ASTReader::SetIdentifierIsMacro(IdentifierInfo *II, PerFileData &F,
@@ -2638,6 +2640,25 @@
   ReadDefinedMacros();
 }
 
+PreprocessedEntity *ASTReader::ReadPreprocessedEntity(uint64_t Offset) {
+  PerFileData *F = 0;  
+  for (unsigned I = 0, N = Chain.size(); I != N; ++I) {
+    if (Offset < Chain[I]->SizeInBits) {
+      F = Chain[I];
+      break;
+    }
+    
+    Offset -= Chain[I]->SizeInBits;
+  }
+  
+  if (!F) {
+    Error("Malformed preprocessed entity offset");
+    return 0;
+  }
+
+  return ReadMacroRecord(*F, Offset);
+}
+
 void ASTReader::ReadUserDiagnosticMappings(Diagnostic &Diag) {
   unsigned Idx = 0;
   while (Idx < UserDiagMappings.size()) {
diff --git a/lib/Serialization/ASTWriter.cpp b/lib/Serialization/ASTWriter.cpp
index b143ce4..9a5085d 100644
--- a/lib/Serialization/ASTWriter.cpp
+++ b/lib/Serialization/ASTWriter.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Serialization/ASTWriter.h"
+#include "clang/Serialization/ASTSerializationListener.h"
 #include "ASTCommon.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Sema/IdentifierResolver.h"
@@ -1358,33 +1359,28 @@
          E != EEnd; ++E) {
       Record.clear();
 
-      if (MacroInstantiation *MI = dyn_cast<MacroInstantiation>(*E)) {
-        Record.push_back(IndexBase + NumPreprocessingRecords++);
-        AddSourceLocation(MI->getSourceRange().getBegin(), Record);
-        AddSourceLocation(MI->getSourceRange().getEnd(), Record);
-        AddIdentifierRef(MI->getName(), Record);
-        Record.push_back(getMacroDefinitionID(MI->getDefinition()));
-        Stream.EmitRecord(PP_MACRO_INSTANTIATION, Record);
-        continue;
-      }
-
       if (MacroDefinition *MD = dyn_cast<MacroDefinition>(*E)) {
         // Record this macro definition's location.
         MacroID ID = getMacroDefinitionID(MD);
-
+        
         // Don't write the macro definition if it is from another AST file.
         if (ID < FirstMacroID)
           continue;
+        
+        // Notify the serialization listener that we're serializing this entity.
+        if (SerializationListener)
+          SerializationListener->SerializedPreprocessedEntity(*E, 
+                                                      Stream.GetCurrentBitNo());
 
         unsigned Position = ID - FirstMacroID;
         if (Position != MacroDefinitionOffsets.size()) {
           if (Position > MacroDefinitionOffsets.size())
             MacroDefinitionOffsets.resize(Position + 1);
-
+          
           MacroDefinitionOffsets[Position] = Stream.GetCurrentBitNo();
         } else
           MacroDefinitionOffsets.push_back(Stream.GetCurrentBitNo());
-
+        
         Record.push_back(IndexBase + NumPreprocessingRecords++);
         Record.push_back(ID);
         AddSourceLocation(MD->getSourceRange().getBegin(), Record);
@@ -1395,6 +1391,21 @@
         continue;
       }
 
+      // Notify the serialization listener that we're serializing this entity.
+      if (SerializationListener)
+        SerializationListener->SerializedPreprocessedEntity(*E, 
+                                                      Stream.GetCurrentBitNo());
+
+      if (MacroInstantiation *MI = dyn_cast<MacroInstantiation>(*E)) {          
+        Record.push_back(IndexBase + NumPreprocessingRecords++);
+        AddSourceLocation(MI->getSourceRange().getBegin(), Record);
+        AddSourceLocation(MI->getSourceRange().getEnd(), Record);
+        AddIdentifierRef(MI->getName(), Record);
+        Record.push_back(getMacroDefinitionID(MI->getDefinition()));
+        Stream.EmitRecord(PP_MACRO_INSTANTIATION, Record);
+        continue;
+      }
+
       if (InclusionDirective *ID = dyn_cast<InclusionDirective>(*E)) {
         Record.push_back(PP_INCLUSION_DIRECTIVE);
         Record.push_back(IndexBase + NumPreprocessingRecords++);
@@ -1409,6 +1420,8 @@
         Stream.EmitRecordWithBlob(InclusionAbbrev, Record, Buffer);
         continue;
       }
+      
+      llvm_unreachable("Unhandled PreprocessedEntity in ASTWriter");
     }
   }
 
@@ -2232,7 +2245,8 @@
 }
 
 ASTWriter::ASTWriter(llvm::BitstreamWriter &Stream)
-  : Stream(Stream), Chain(0), FirstDeclID(1), NextDeclID(FirstDeclID),
+  : Stream(Stream), Chain(0), SerializationListener(0), 
+    FirstDeclID(1), NextDeclID(FirstDeclID),
     FirstTypeID(NUM_PREDEF_TYPE_IDS), NextTypeID(FirstTypeID),
     FirstIdentID(1), NextIdentID(FirstIdentID), FirstSelectorID(1),
     NextSelectorID(FirstSelectorID), FirstMacroID(1), NextMacroID(FirstMacroID),
@@ -3415,3 +3429,5 @@
   Record.push_back(UPD_CXX_ADDED_TEMPLATE_SPECIALIZATION);
   AddDeclRef(D, Record);
 }
+
+ASTSerializationListener::~ASTSerializationListener() { }
diff --git a/lib/Serialization/GeneratePCH.cpp b/lib/Serialization/GeneratePCH.cpp
index 4f6f5ca..3b4f869 100644
--- a/lib/Serialization/GeneratePCH.cpp
+++ b/lib/Serialization/GeneratePCH.cpp
@@ -32,7 +32,6 @@
                            llvm::raw_ostream *OS)
   : PP(PP), isysroot(isysroot), Out(OS), SemaPtr(0),
     StatCalls(0), Stream(Buffer), Writer(Stream), Chaining(Chaining) {
-
   // Install a stat() listener to keep track of all of the stat()
   // calls.
   StatCalls = new MemorizeStatCalls();
@@ -46,6 +45,9 @@
   if (PP.getDiagnostics().hasErrorOccurred())
     return;
 
+  // Set up the serialization listener.
+  Writer.SetSerializationListener(GetASTSerializationListener());
+  
   // Emit the PCH file
   assert(SemaPtr && "No Sema?");
   Writer.WriteAST(*SemaPtr, StatCalls, isysroot);
@@ -66,6 +68,10 @@
   return 0;
 }
 
+ASTSerializationListener *PCHGenerator::GetASTSerializationListener() {
+  return 0;
+}
+
 ASTDeserializationListener *PCHGenerator::GetASTDeserializationListener() {
   return &Writer;
 }
diff --git a/test/Index/c-index-getCursor-pp.c b/test/Index/c-index-getCursor-pp.c
index 5f3c1df..2393965 100644
--- a/test/Index/c-index-getCursor-pp.c
+++ b/test/Index/c-index-getCursor-pp.c
@@ -20,3 +20,9 @@
 // CHECK-5: macro instantiation=DECORATION:2:9
 // RUN: c-index-test -cursor-at=%s:9:10 -I%S/Inputs %s | FileCheck -check-prefix=CHECK-6 %s
 // CHECK-6: inclusion directive=a.h
+
+// Same tests, but with "editing" optimizations
+// RUN: env CINDEXTEST_EDITING=1 c-index-test -cursor-at=%s:1:11 -I%S/Inputs %s | FileCheck -check-prefix=CHECK-1 %s
+// RUN: env CINDEXTEST_EDITING=1 c-index-test -cursor-at=%s:2:14 -I%S/Inputs %s | FileCheck -check-prefix=CHECK-2 %s
+// RUN: env CINDEXTEST_EDITING=1 c-index-test -cursor-at=%s:5:7 -I%S/Inputs %s | FileCheck -check-prefix=CHECK-3 %s
+// RUN: env CINDEXTEST_EDITING=1 c-index-test -cursor-at=%s:9:10 -I%S/Inputs %s | FileCheck -check-prefix=CHECK-6 %s
diff --git a/tools/libclang/CIndex.cpp b/tools/libclang/CIndex.cpp
index 7fc2f6e..cdd7764 100644
--- a/tools/libclang/CIndex.cpp
+++ b/tools/libclang/CIndex.cpp
@@ -410,10 +410,18 @@
   bool OnlyLocalDecls
     = !AU->isMainFileAST() && AU->getOnlyLocalDecls();
   
+  PreprocessingRecord::iterator StartEntity, EndEntity;
+  if (OnlyLocalDecls) {
+    StartEntity = AU->pp_entity_begin();
+    EndEntity = AU->pp_entity_end();
+  } else {
+    StartEntity = PPRec.begin();
+    EndEntity = PPRec.end();
+  }
+  
   // There is no region of interest; we have to walk everything.
   if (RegionOfInterest.isInvalid())
-    return std::make_pair(PPRec.begin(OnlyLocalDecls),
-                          PPRec.end(OnlyLocalDecls));
+    return std::make_pair(StartEntity, EndEntity);
 
   // Find the file in which the region of interest lands.
   SourceManager &SM = AU->getSourceManager();
@@ -424,18 +432,16 @@
   
   // The region of interest spans files; we have to walk everything.
   if (Begin.first != End.first)
-    return std::make_pair(PPRec.begin(OnlyLocalDecls),
-                          PPRec.end(OnlyLocalDecls));
+    return std::make_pair(StartEntity, EndEntity);
     
   ASTUnit::PreprocessedEntitiesByFileMap &ByFileMap
     = AU->getPreprocessedEntitiesByFile();
   if (ByFileMap.empty()) {
     // Build the mapping from files to sets of preprocessed entities.
-    for (PreprocessingRecord::iterator E = PPRec.begin(OnlyLocalDecls),
-                                    EEnd = PPRec.end(OnlyLocalDecls);
-         E != EEnd; ++E) {
+    for (PreprocessingRecord::iterator E = StartEntity; E != EndEntity; ++E) {
       std::pair<FileID, unsigned> P
         = SM.getDecomposedInstantiationLoc((*E)->getSourceRange().getBegin());
+      
       ByFileMap[P.first].push_back(*E);
     }
   }