1 files changed, 67 insertions, 9 deletions

diff --git a/tools/libclang/CIndex.cpp b/tools/libclang/CIndex.cpp
index 9e261de964..ae488abf39 100644
--- a/tools/libclang/CIndex.cpp
+++ b/tools/libclang/CIndex.cpp
@@ -182,8 +182,13 @@ bool CursorVisitor::Visit(CXCursor Cursor, bool CheckedRegionOfInterest) {
   case CXChildVisit_Continue:
     return false;
 
-  case CXChildVisit_Recurse:
-    return VisitChildren(Cursor);
+  case CXChildVisit_Recurse: {
+    bool ret = VisitChildren(Cursor);
+    if (PostChildrenVisitor)
+      if (PostChildrenVisitor(Cursor, ClientData))
+        return true;
+    return ret;
+  }
   }
 
   llvm_unreachable("Invalid CXChildVisitResult!");
@@ -1632,6 +1637,7 @@ DEF_JOB(ExplicitTemplateArgsVisit, ASTTemplateArgumentListInfo,
         ExplicitTemplateArgsVisitKind)
 DEF_JOB(SizeOfPackExprParts, SizeOfPackExpr, SizeOfPackExprPartsKind)
 DEF_JOB(LambdaExprParts, LambdaExpr, LambdaExprPartsKind)
+DEF_JOB(PostChildrenVisit, void, PostChildrenVisitKind)
 #undef DEF_JOB
 
 class DeclVisit : public VisitorJob {
@@ -2208,6 +2214,8 @@ bool CursorVisitor::RunVisitorWorkList(VisitorWorkList &WL) {
           case CXChildVisit_Break: return true;
           case CXChildVisit_Continue: break;
           case CXChildVisit_Recurse:
+            if (PostChildrenVisitor)
+              WL.push_back(PostChildrenVisit(0, Cursor));
             EnqueueWorkList(WL, S);
             break;
         }
@@ -2324,6 +2332,11 @@ bool CursorVisitor::RunVisitorWorkList(VisitorWorkList &WL) {
         }
         break;
       }
+
+      case VisitorJob::PostChildrenVisitKind:
+        if (PostChildrenVisitor(Parent, ClientData))
+          return true;
+        break;
     }
   }
   return false;
@@ -4819,6 +4832,9 @@ typedef llvm::DenseMap<unsigned, CXCursor> AnnotateTokensData;
 static enum CXChildVisitResult AnnotateTokensVisitor(CXCursor cursor,
                                                      CXCursor parent,
                                                      CXClientData client_data);
+static bool AnnotateTokensPostChildrenVisitor(CXCursor cursor,
+                                              CXClientData client_data);
+
 namespace {
 class AnnotateTokensWorker {
   AnnotateTokensData &Annotated;
@@ -4830,6 +4846,13 @@ class AnnotateTokensWorker {
   CursorVisitor AnnotateVis;
   SourceManager &SrcMgr;
   bool HasContextSensitiveKeywords;
+
+  struct PostChildrenInfo {
+    CXCursor Cursor;
+    SourceRange CursorRange;
+    unsigned BeforeChildrenTokenIdx;
+  };
+  llvm::SmallVector<PostChildrenInfo, 8> PostChildrenInfos;
   
   bool MoreTokens() const { return TokIdx < NumTokens; }
   unsigned NextToken() const { return TokIdx; }
@@ -4858,12 +4881,15 @@ public:
                   AnnotateTokensVisitor, this,
                   /*VisitPreprocessorLast=*/true,
                   /*VisitIncludedEntities=*/false,
-                  RegionOfInterest),
+                  RegionOfInterest,
+                  /*VisitDeclsOnly=*/false,
+                  AnnotateTokensPostChildrenVisitor),
       SrcMgr(static_cast<ASTUnit*>(tu->TUData)->getSourceManager()),
       HasContextSensitiveKeywords(false) { }
 
   void VisitChildren(CXCursor C) { AnnotateVis.VisitChildren(C); }
   enum CXChildVisitResult Visit(CXCursor cursor, CXCursor parent);
+  bool postVisitChildren(CXCursor cursor);
   void AnnotateTokens();
   
   /// \brief Determine whether the annotator saw any cursors that have 
@@ -4871,6 +4897,10 @@ public:
   bool hasContextSensitiveKeywords() const {
     return HasContextSensitiveKeywords;
   }
+
+  ~AnnotateTokensWorker() {
+    assert(PostChildrenInfos.empty());
+  }
 };
 }
 
@@ -5131,25 +5161,47 @@ AnnotateTokensWorker::Visit(CXCursor cursor, CXCursor parent) {
     }
   }
 
-  // Visit children to get their cursor information.
-  const unsigned BeforeChildren = NextToken();
-  VisitChildren(cursor);
+  // Before recursing into the children keep some state that we are going
+  // to use in the AnnotateTokensWorker::postVisitChildren callback to do some
+  // extra work after the child nodes are visited.
+  // Note that we don't call VisitChildren here to avoid traversing statements
+  // code-recursively which can blow the stack.
+
+  PostChildrenInfo Info;
+  Info.Cursor = cursor;
+  Info.CursorRange = cursorRange;
+  Info.BeforeChildrenTokenIdx = NextToken();
+  PostChildrenInfos.push_back(Info);
+
+  return CXChildVisit_Recurse;
+}
+
+bool AnnotateTokensWorker::postVisitChildren(CXCursor cursor) {
+  if (PostChildrenInfos.empty())
+    return false;
+  const PostChildrenInfo &Info = PostChildrenInfos.back();
+  if (!clang_equalCursors(Info.Cursor, cursor))
+    return false;
+
+  const unsigned BeforeChildren = Info.BeforeChildrenTokenIdx;
   const unsigned AfterChildren = NextToken();
+  SourceRange cursorRange = Info.CursorRange;
 
   // Scan the tokens that are at the end of the cursor, but are not captured
   // but the child cursors.
   annotateAndAdvanceTokens(cursor, RangeOverlap, cursorRange);
-  
+
   // Scan the tokens that are at the beginning of the cursor, but are not
   // capture by the child cursors.
   for (unsigned I = BeforeChildren; I != AfterChildren; ++I) {
     if (!clang_isInvalid(clang_getCursorKind(Cursors[I])))
       break;
-    
+
     Cursors[I] = cursor;
   }
 
-  return CXChildVisit_Continue;
+  PostChildrenInfos.pop_back();
+  return false;
 }
 
 static enum CXChildVisitResult AnnotateTokensVisitor(CXCursor cursor,
@@ -5158,6 +5210,12 @@ static enum CXChildVisitResult AnnotateTokensVisitor(CXCursor cursor,
   return static_cast<AnnotateTokensWorker*>(client_data)->Visit(cursor, parent);
 }
 
+static bool AnnotateTokensPostChildrenVisitor(CXCursor cursor,
+                                              CXClientData client_data) {
+  return static_cast<AnnotateTokensWorker*>(client_data)->
+                                                      postVisitChildren(cursor);
+}
+
 namespace {
 
 /// \brief Uses the macro expansions in the preprocessing record to find