1 files changed, 82 insertions, 14 deletions

diff --git a/clang/tools/libclang/CIndex.cpp b/clang/tools/libclang/CIndex.cpp
index 17b2316bede..4bf1c5286c4 100644
--- a/clang/tools/libclang/CIndex.cpp
+++ b/clang/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();
   }