[PCH] When visiting preprocessed entities, make it possible to avoid deserializing

preprocessed entities that are #included in the range that we are interested.

This is useful when we are interested in preprocessed entities of a specific file, e.g
when we are annotating tokens. There is also an optimization where we cache the last
result of PreprocessingRecord::getPreprocessedEntitiesInRange and we re-use it if
there is a call with the same range as before.

rdar://10313365

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

include/clang/Basic/SourceManager.h

@@ -807,6 +807,18 @@ public:
     unsigned FileOffset = Entry.getOffset();
     return SourceLocation::getFileLoc(FileOffset);
   }
+  
+  /// \brief Return the source location corresponding to the last byte of the
+  /// specified file.
+  SourceLocation getLocForEndOfFile(FileID FID) const {
+    bool Invalid = false;
+    const SrcMgr::SLocEntry &Entry = getSLocEntry(FID, &Invalid);
+    if (Invalid || !Entry.isFile())
+      return SourceLocation();
+    
+    unsigned FileOffset = Entry.getOffset();
+    return SourceLocation::getFileLoc(FileOffset + getFileIDSize(FID) - 1);
+  }
 
   /// \brief Returns the include location if \arg FID is a #include'd file
   /// otherwise it returns an invalid location.

include/clang/Frontend/ASTUnit.h

@@ -546,6 +546,11 @@ public:
   /// preamble, otherwise it returns \arg Loc.
   SourceLocation mapLocationToPreamble(SourceLocation Loc);
 
+  bool isInPreambleFileID(SourceLocation Loc);
+  bool isInMainFileID(SourceLocation Loc);
+  SourceLocation getStartOfMainFileID();
+  SourceLocation getEndOfPreambleFileID();
+
   /// \brief \see mapLocationFromPreamble.
   SourceRange mapRangeFromPreamble(SourceRange R) {
     return SourceRange(mapLocationFromPreamble(R.getBegin()),

include/clang/Lex/PreprocessingRecord.h

@@ -18,6 +18,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/IdentifierTable.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/Support/Allocator.h"
 #include <vector>
 
@@ -269,6 +270,13 @@ namespace clang {
     /// preprocessed entities that \arg Range encompasses.
     virtual std::pair<unsigned, unsigned>
         findPreprocessedEntitiesInRange(SourceRange Range) = 0;
+
+    /// \brief Optionally returns true or false if the preallocated preprocessed
+    /// entity with index \arg Index came from file \arg FID.
+    virtual llvm::Optional<bool> isPreprocessedEntityInFileID(unsigned Index,
+                                                              FileID FID) {
+      return llvm::Optional<bool>();
+    }
   };
   
   /// \brief A record of the steps taken while preprocessing a source file,
@@ -386,7 +394,7 @@ namespace clang {
       
       iterator() : Self(0), Position(0) { }
       
-      iterator(PreprocessingRecord *Self, int Position) 
+      iterator(PreprocessingRecord *Self, PPEntityID Position) 
         : Self(Self), Position(Position) { }
       
       value_type operator*() const {
@@ -471,6 +479,7 @@ namespace clang {
         X.Position -= D;
         return X;
       }
+      friend class PreprocessingRecord;
     };
     friend class iterator;
 
@@ -496,8 +505,20 @@ namespace clang {
 
     /// \brief Returns a pair of [Begin, End) iterators of preprocessed entities
     /// that source range \arg R encompasses.
+    ///
+    /// \param R the range to look for preprocessed entities.
+    ///
     std::pair<iterator, iterator> getPreprocessedEntitiesInRange(SourceRange R);
 
+    /// \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 isEntityInFileID(iterator PPEI, FileID FID);
+
     /// \brief Add a new preprocessed entity to this record.
     void addPreprocessedEntity(PreprocessedEntity *Entity);
     
@@ -526,6 +547,17 @@ namespace clang {
                                     StringRef SearchPath,
                                     StringRef RelativePath);
 
+  private:
+    /// \brief Cached result of the last \see getPreprocessedEntitiesInRange
+    /// query.
+    struct {
+      SourceRange Range;
+      std::pair<PPEntityID, PPEntityID> Result;
+    } CachedRangeQuery;
+
+    std::pair<PPEntityID, PPEntityID>
+      getPreprocessedEntitiesInRangeSlow(SourceRange R);
+
     friend class ASTReader;
     friend class ASTWriter;
   };

include/clang/Serialization/ASTReader.h

@@ -703,6 +703,11 @@ private:
     findNextPreprocessedEntity(
                         GlobalSLocOffsetMapType::const_iterator SLocMapI) const;
 
+  /// \brief Returns (Module, Local index) pair for \arg GlobalIndex of a
+  /// preprocessed entity.
+  std::pair<Module *, unsigned>
+    getModulePreprocessedEntity(unsigned GlobalIndex);
+
   void PassInterestingDeclsToConsumer();
 
   /// \brief Produce an error diagnostic and return true.
@@ -800,6 +805,11 @@ public:
   /// preprocessed entities that \arg Range encompasses.
   virtual std::pair<unsigned, unsigned>
       findPreprocessedEntitiesInRange(SourceRange Range);
+  
+  /// \brief Optionally returns true or false if the preallocated preprocessed
+  /// entity with index \arg Index came from file \arg FID.
+  virtual llvm::Optional<bool> isPreprocessedEntityInFileID(unsigned Index,
+                                                            FileID FID);
 
   /// \brief Read the header file information for the given file entry.
   virtual HeaderFileInfo GetHeaderFileInfo(const FileEntry *FE);

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())

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>

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 {

tools/libclang/CIndex.cpp

@@ -192,6 +192,10 @@ class CursorVisitor : public DeclVisitor<CursorVisitor, bool>,
   /// \brief Whether we should visit the preprocessing record entries last, 
   /// after visiting other declarations.
   bool VisitPreprocessorLast;
+
+  /// \brief Whether we should visit the preprocessing record entries that are
+  /// #included inside the \arg RegionOfInterest.
+  bool VisitIncludedPreprocessingEntries;
   
   /// \brief When valid, a source range to which the cursor should restrict
   /// its search.
@@ -240,10 +244,12 @@ public:
   CursorVisitor(CXTranslationUnit TU, CXCursorVisitor Visitor,
                 CXClientData ClientData,
                 bool VisitPreprocessorLast,
+                bool VisitIncludedPreprocessingEntries = false,
                 SourceRange RegionOfInterest = SourceRange())
     : TU(TU), AU(static_cast<ASTUnit*>(TU->TUData)),
       Visitor(Visitor), ClientData(ClientData),
       VisitPreprocessorLast(VisitPreprocessorLast),
+      VisitIncludedPreprocessingEntries(VisitIncludedPreprocessingEntries),
       RegionOfInterest(RegionOfInterest), DI_current(0)
   {
     Parent.kind = CXCursor_NoDeclFound;
@@ -268,8 +274,14 @@ public:
   
   bool visitPreprocessedEntitiesInRegion();
 
+  bool shouldVisitIncludedPreprocessingEntries() const {
+    return VisitIncludedPreprocessingEntries;
+  }
+
   template<typename InputIterator>
-  bool visitPreprocessedEntities(InputIterator First, InputIterator Last);
+  bool visitPreprocessedEntities(InputIterator First, InputIterator Last,
+                                 PreprocessingRecord &PPRec,
+                                 FileID FID = FileID());
 
   bool VisitChildren(CXCursor Parent);
 
@@ -399,45 +411,95 @@ bool CursorVisitor::Visit(CXCursor Cursor, bool CheckedRegionOfInterest) {
   return false;
 }
 
+static bool visitPreprocessedEntitiesInRange(SourceRange R,
+                                             PreprocessingRecord &PPRec,
+                                             CursorVisitor &Visitor) {
+  SourceManager &SM = Visitor.getASTUnit()->getSourceManager();
+  FileID FID;
+  
+  if (!Visitor.shouldVisitIncludedPreprocessingEntries()) {
+    // If the begin/end of the range lie in the same FileID, do the optimization
+    // where we skip preprocessed entities that do not come from the same FileID.
+    FID = SM.getFileID(R.getBegin());
+    if (FID != SM.getFileID(R.getEnd()))
+      FID = FileID();
+  }
+
+  std::pair<PreprocessingRecord::iterator, PreprocessingRecord::iterator>
+    Entities = PPRec.getPreprocessedEntitiesInRange(R);
+  return Visitor.visitPreprocessedEntities(Entities.first, Entities.second,
+                                           PPRec, FID);
+}
+
 bool CursorVisitor::visitPreprocessedEntitiesInRegion() {
   PreprocessingRecord &PPRec
     = *AU->getPreprocessor().getPreprocessingRecord();
+  SourceManager &SM = AU->getSourceManager();
   
   if (RegionOfInterest.isValid()) {
     SourceRange MappedRange = AU->mapRangeToPreamble(RegionOfInterest);
-    std::pair<PreprocessingRecord::iterator, PreprocessingRecord::iterator>
-      Entities = PPRec.getPreprocessedEntitiesInRange(MappedRange);
-    return visitPreprocessedEntities(Entities.first, Entities.second);
+    SourceLocation B = MappedRange.getBegin();
+    SourceLocation E = MappedRange.getEnd();
+
+    if (AU->isInPreambleFileID(B)) {
+      if (SM.isLoadedSourceLocation(E))
+        return visitPreprocessedEntitiesInRange(SourceRange(B, E),
+                                                 PPRec, *this);
+
+      // Beginning of range lies in the preamble but it also extends beyond
+      // it into the main file. Split the range into 2 parts, one covering
+      // the preamble and another covering the main file. This allows subsequent
+      // calls to visitPreprocessedEntitiesInRange to accept a source range that
+      // lies in the same FileID, allowing it to skip preprocessed entities that
+      // do not come from the same FileID.
+      bool breaked =
+        visitPreprocessedEntitiesInRange(
+                                   SourceRange(B, AU->getEndOfPreambleFileID()),
+                                          PPRec, *this);
+      if (breaked) return true;
+      return visitPreprocessedEntitiesInRange(
+                                    SourceRange(AU->getStartOfMainFileID(), E),
+                                        PPRec, *this);
+    }
+
+    return visitPreprocessedEntitiesInRange(SourceRange(B, E), PPRec, *this);
   }
 
   bool OnlyLocalDecls
     = !AU->isMainFileAST() && AU->getOnlyLocalDecls(); 
   
   if (OnlyLocalDecls)
-    return visitPreprocessedEntities(PPRec.local_begin(), PPRec.local_end());
+    return visitPreprocessedEntities(PPRec.local_begin(), PPRec.local_end(),
+                                     PPRec);
 
-  return visitPreprocessedEntities(PPRec.begin(), PPRec.end());
+  return visitPreprocessedEntities(PPRec.begin(), PPRec.end(), PPRec);
 }
 
 template<typename InputIterator>
 bool CursorVisitor::visitPreprocessedEntities(InputIterator First,
-                                              InputIterator Last) {
+                                              InputIterator Last,
+                                              PreprocessingRecord &PPRec,
+                                              FileID FID) {
   for (; First != Last; ++First) {
-    if (MacroExpansion *ME = dyn_cast<MacroExpansion>(*First)) {
+    if (!FID.isInvalid() && !PPRec.isEntityInFileID(First, FID))
+      continue;
+
+    PreprocessedEntity *PPE = *First;
+    if (MacroExpansion *ME = dyn_cast<MacroExpansion>(PPE)) {
       if (Visit(MakeMacroExpansionCursor(ME, TU)))
         return true;
       
       continue;
     }
     
-    if (MacroDefinition *MD = dyn_cast<MacroDefinition>(*First)) {
+    if (MacroDefinition *MD = dyn_cast<MacroDefinition>(PPE)) {
       if (Visit(MakeMacroDefinitionCursor(MD, TU)))
         return true;
       
       continue;
     }
     
-    if (InclusionDirective *ID = dyn_cast<InclusionDirective>(*First)) {
+    if (InclusionDirective *ID = dyn_cast<InclusionDirective>(PPE)) {
       if (Visit(MakeInclusionDirectiveCursor(ID, TU)))
         return true;
       
@@ -3007,8 +3069,8 @@ extern "C" {
 unsigned clang_visitChildren(CXCursor parent,
                              CXCursorVisitor visitor,
                              CXClientData client_data) {
-  CursorVisitor CursorVis(getCursorTU(parent), visitor, client_data, 
-                          false);
+  CursorVisitor CursorVis(getCursorTU(parent), visitor, client_data,
+                          /*VisitPreprocessorLast=*/false);
   return CursorVis.VisitChildren(parent);
 }
 
@@ -3895,6 +3957,7 @@ CXCursor cxcursor::getCursor(CXTranslationUnit TU, SourceLocation SLoc) {
     CXCursor Parent = clang_getTranslationUnitCursor(TU);
     CursorVisitor CursorVis(TU, GetCursorVisitor, &ResultData,
                             /*VisitPreprocessorLast=*/true, 
+                            /*VisitIncludedPreprocessingEntries=*/false,
                             SourceLocation(SLoc));
     CursorVis.VisitChildren(Parent);
   }
@@ -4722,7 +4785,10 @@ public:
     : Annotated(annotated), Tokens(tokens), Cursors(cursors),
       NumTokens(numTokens), TokIdx(0), PreprocessingTokIdx(0),
       AnnotateVis(tu,
-                  AnnotateTokensVisitor, this, true, RegionOfInterest),
+                  AnnotateTokensVisitor, this,
+                  /*VisitPreprocessorLast=*/true,
+                  /*VisitIncludedPreprocessingEntries=*/false,
+                  RegionOfInterest),
       SrcMgr(static_cast<ASTUnit*>(tu->TUData)->getSourceManager()),
       HasContextSensitiveKeywords(false) { }
 
@@ -5197,7 +5263,9 @@ static void clang_annotateTokensImpl(void *UserData) {
                                       Tokens, NumTokens);
     CursorVisitor MacroArgMarker(TU,
                                  MarkMacroArgTokensVisitorDelegate, &Visitor,
-                                 true, RegionOfInterest);
+                                 /*VisitPreprocessorLast=*/true,
+                                 /*VisitIncludedPreprocessingEntries=*/false,
+                                 RegionOfInterest);
     MacroArgMarker.visitPreprocessedEntitiesInRegion();
   }