3 files changed, 168 insertions, 21 deletions

diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index c153435225..0cdcfe87c3 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -2397,6 +2397,50 @@ SourceLocation ASTUnit::mapLocationToPreamble(SourceLocation Loc) {
   return Loc;
 }
 
+bool ASTUnit::isInPreambleFileID(SourceLocation Loc) {
+  FileID FID;
+  if (SourceMgr)
+    FID = SourceMgr->getPreambleFileID();
+  
+  if (Loc.isInvalid() || FID.isInvalid())
+    return false;
+  
+  return SourceMgr->isInFileID(Loc, FID);
+}
+
+bool ASTUnit::isInMainFileID(SourceLocation Loc) {
+  FileID FID;
+  if (SourceMgr)
+    FID = SourceMgr->getMainFileID();
+  
+  if (Loc.isInvalid() || FID.isInvalid())
+    return false;
+  
+  return SourceMgr->isInFileID(Loc, FID);
+}
+
+SourceLocation ASTUnit::getEndOfPreambleFileID() {
+  FileID FID;
+  if (SourceMgr)
+    FID = SourceMgr->getPreambleFileID();
+  
+  if (FID.isInvalid())
+    return SourceLocation();
+
+  return SourceMgr->getLocForEndOfFile(FID);
+}
+
+SourceLocation ASTUnit::getStartOfMainFileID() {
+  FileID FID;
+  if (SourceMgr)
+    FID = SourceMgr->getMainFileID();
+  
+  if (FID.isInvalid())
+    return SourceLocation();
+  
+  return SourceMgr->getLocForStartOfFile(FID);
+}
+
 void ASTUnit::PreambleData::countLines() const {
   NumLines = 0;
   if (empty())
diff --git a/lib/Lex/PreprocessingRecord.cpp b/lib/Lex/PreprocessingRecord.cpp
index 2816609d5f..770a69d19b 100644
--- a/lib/Lex/PreprocessingRecord.cpp
+++ b/lib/Lex/PreprocessingRecord.cpp
@@ -49,35 +49,108 @@ PreprocessingRecord::PreprocessingRecord(SourceManager &SM,
 std::pair<PreprocessingRecord::iterator, PreprocessingRecord::iterator>
 PreprocessingRecord::getPreprocessedEntitiesInRange(SourceRange Range) {
   if (Range.isInvalid())
-    return std::make_pair(iterator(this, 0), iterator(this, 0));
-  assert(!SourceMgr.isBeforeInTranslationUnit(Range.getEnd(),Range.getBegin()));
+    return std::make_pair(iterator(), iterator());
+
+  if (CachedRangeQuery.Range == Range) {
+    return std::make_pair(iterator(this, CachedRangeQuery.Result.first),
+                          iterator(this, CachedRangeQuery.Result.second));
+  }
+
+  std::pair<PPEntityID, PPEntityID>
+    Res = getPreprocessedEntitiesInRangeSlow(Range);
+  
+  CachedRangeQuery.Range = Range;
+  CachedRangeQuery.Result = Res;
+  
+  return std::make_pair(iterator(this, Res.first), iterator(this, Res.second));
+}
+
+static bool isPreprocessedEntityIfInFileID(PreprocessedEntity *PPE, FileID FID,
+                                           SourceManager &SM) {
+  assert(!FID.isInvalid());
+  if (!PPE)
+    return false;
+
+  SourceLocation Loc = PPE->getSourceRange().getBegin();
+  if (Loc.isInvalid())
+    return false;
+  
+  if (SM.isInFileID(SM.getFileLoc(Loc), FID))
+    return true;
+  else
+    return false;
+}
+
+/// \brief Returns true if the preprocessed entity that \arg PPEI iterator
+/// points to is coming from the file \arg FID.
+///
+/// Can be used to avoid implicit deserializations of preallocated
+/// preprocessed entities if we only care about entities of a specific file
+/// and not from files #included in the range given at
+/// \see getPreprocessedEntitiesInRange.
+bool PreprocessingRecord::isEntityInFileID(iterator PPEI, FileID FID) {
+  if (FID.isInvalid())
+    return false;
+
+  PPEntityID PPID = PPEI.Position;
+  if (PPID < 0) {
+    assert(unsigned(-PPID-1) < LoadedPreprocessedEntities.size() &&
+           "Out-of bounds loaded preprocessed entity");
+    assert(ExternalSource && "No external source to load from");
+    unsigned LoadedIndex = LoadedPreprocessedEntities.size()+PPID;
+    if (PreprocessedEntity *PPE = LoadedPreprocessedEntities[LoadedIndex])
+      return isPreprocessedEntityIfInFileID(PPE, FID, SourceMgr);
+
+    // See if the external source can see if the entity is in the file without
+    // deserializing it.
+    llvm::Optional<bool>
+      IsInFile = ExternalSource->isPreprocessedEntityInFileID(LoadedIndex, FID);
+    if (IsInFile.hasValue())
+      return IsInFile.getValue();
+
+    // The external source did not provide a definite answer, go and deserialize
+    // the entity to check it.
+    return isPreprocessedEntityIfInFileID(
+                                       getLoadedPreprocessedEntity(LoadedIndex),
+                                          FID, SourceMgr);
+  }
+
+  assert(unsigned(PPID) < PreprocessedEntities.size() &&
+         "Out-of bounds local preprocessed entity");
+  return isPreprocessedEntityIfInFileID(PreprocessedEntities[PPID],
+                                        FID, SourceMgr);
+}
 
+/// \brief Returns a pair of [Begin, End) iterators of preprocessed entities
+/// that source range \arg R encompasses.
+std::pair<PreprocessingRecord::PPEntityID, PreprocessingRecord::PPEntityID>
+PreprocessingRecord::getPreprocessedEntitiesInRangeSlow(SourceRange Range) {
+  assert(Range.isValid());
+  assert(!SourceMgr.isBeforeInTranslationUnit(Range.getEnd(),Range.getBegin()));
+  
   std::pair<unsigned, unsigned>
     Local = findLocalPreprocessedEntitiesInRange(Range);
-
+  
   // Check if range spans local entities.
   if (!ExternalSource || SourceMgr.isLocalSourceLocation(Range.getBegin()))
-    return std::make_pair(iterator(this, Local.first),
-                          iterator(this, Local.second));
-
+    return std::make_pair(Local.first, Local.second);
+  
   std::pair<unsigned, unsigned>
     Loaded = ExternalSource->findPreprocessedEntitiesInRange(Range);
-
+  
   // Check if range spans local entities.
   if (Loaded.first == Loaded.second)
-    return std::make_pair(iterator(this, Local.first),
-                          iterator(this, Local.second));
-
+    return std::make_pair(Local.first, Local.second);
+  
   unsigned TotalLoaded = LoadedPreprocessedEntities.size();
-
+  
   // Check if range spans loaded entities.
   if (Local.first == Local.second)
-    return std::make_pair(iterator(this, int(Loaded.first)-TotalLoaded),
-                          iterator(this, int(Loaded.second)-TotalLoaded));
-
+    return std::make_pair(int(Loaded.first)-TotalLoaded,
+                          int(Loaded.second)-TotalLoaded);
+  
   // Range spands loaded and local entities.
-  return std::make_pair(iterator(this, int(Loaded.first)-TotalLoaded),
-                        iterator(this, Local.second));
+  return std::make_pair(int(Loaded.first)-TotalLoaded, Local.second);
 }
 
 std::pair<unsigned, unsigned>
diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index cf14c211fe..7afb3e43de 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -2778,14 +2778,22 @@ bool ASTReader::ParseLanguageOptions(
   return false;
 }
 
-PreprocessedEntity *ASTReader::ReadPreprocessedEntity(unsigned Index) {
-  PreprocessedEntityID PPID = Index+1;
+std::pair<Module *, unsigned>
+ASTReader::getModulePreprocessedEntity(unsigned GlobalIndex) {
   GlobalPreprocessedEntityMapType::iterator
-    I = GlobalPreprocessedEntityMap.find(Index);
+  I = GlobalPreprocessedEntityMap.find(GlobalIndex);
   assert(I != GlobalPreprocessedEntityMap.end() && 
          "Corrupted global preprocessed entity map");
-  Module &M = *I->second;
-  unsigned LocalIndex = Index - M.BasePreprocessedEntityID;
+  Module *M = I->second;
+  unsigned LocalIndex = GlobalIndex - M->BasePreprocessedEntityID;
+  return std::make_pair(M, LocalIndex);
+}
+
+PreprocessedEntity *ASTReader::ReadPreprocessedEntity(unsigned Index) {
+  PreprocessedEntityID PPID = Index+1;
+  std::pair<Module *, unsigned> PPInfo = getModulePreprocessedEntity(Index);
+  Module &M = *PPInfo.first;
+  unsigned LocalIndex = PPInfo.second;
   const PPEntityOffset &PPOffs = M.PreprocessedEntityOffsets[LocalIndex];
 
   SavedStreamPosition SavedPosition(M.PreprocessorDetailCursor);  
@@ -3022,6 +3030,28 @@ std::pair<unsigned, unsigned>
   return std::make_pair(BeginID, EndID);
 }
 
+/// \brief Optionally returns true or false if the preallocated preprocessed
+/// entity with index \arg Index came from file \arg FID.
+llvm::Optional<bool> ASTReader::isPreprocessedEntityInFileID(unsigned Index,
+                                                             FileID FID) {
+  if (FID.isInvalid())
+    return false;
+
+  std::pair<Module *, unsigned> PPInfo = getModulePreprocessedEntity(Index);
+  Module &M = *PPInfo.first;
+  unsigned LocalIndex = PPInfo.second;
+  const PPEntityOffset &PPOffs = M.PreprocessedEntityOffsets[LocalIndex];
+  
+  SourceLocation Loc = ReadSourceLocation(M, PPOffs.Begin);
+  if (Loc.isInvalid())
+    return false;
+  
+  if (SourceMgr.isInFileID(SourceMgr.getFileLoc(Loc), FID))
+    return true;
+  else
+    return false;
+}
+
 namespace {
   /// \brief Visitor used to search for information about a header file.
   class HeaderFileInfoVisitor {