1 files changed, 131 insertions, 1 deletions

diff --git a/tools/libclang/CIndex.cpp b/tools/libclang/CIndex.cpp
index 0affc17117..b05e6f34a8 100644
--- a/tools/libclang/CIndex.cpp
+++ b/tools/libclang/CIndex.cpp
@@ -6744,11 +6744,18 @@ class AnnotateTokensWorker {
   SourceManager &SrcMgr;
   bool HasContextSensitiveKeywords;
 
+  struct PostChildrenAction {
+    CXCursor cursor;
+    enum Action { Invalid, Ignore, Postpone } action;
+  };
+  using PostChildrenActions = SmallVector<PostChildrenAction, 0>;
+
   struct PostChildrenInfo {
     CXCursor Cursor;
     SourceRange CursorRange;
     unsigned BeforeReachingCursorIdx;
     unsigned BeforeChildrenTokenIdx;
+    PostChildrenActions ChildActions;
   };
   SmallVector<PostChildrenInfo, 8> PostChildrenInfos;
 
@@ -6794,7 +6801,13 @@ public:
 
   void VisitChildren(CXCursor C) { AnnotateVis.VisitChildren(C); }
   enum CXChildVisitResult Visit(CXCursor cursor, CXCursor parent);
+  bool IsIgnoredChildCursor(CXCursor cursor) const;
+  PostChildrenActions DetermineChildActions(CXCursor Cursor) const;
+
   bool postVisitChildren(CXCursor cursor);
+  void HandlePostPonedChildCursors(const PostChildrenInfo &Info);
+  void HandlePostPonedChildCursor(CXCursor Cursor, unsigned StartTokenIndex);
+
   void AnnotateTokens();
   
   /// \brief Determine whether the annotator saw any cursors that have 
@@ -6815,6 +6828,67 @@ void AnnotateTokensWorker::AnnotateTokens() {
   AnnotateVis.visitFileRegion();
 }
 
+bool AnnotateTokensWorker::IsIgnoredChildCursor(CXCursor cursor) const {
+  if (PostChildrenInfos.empty())
+    return false;
+
+  for (const auto &ChildAction : PostChildrenInfos.back().ChildActions) {
+    if (ChildAction.cursor == cursor &&
+        ChildAction.action == PostChildrenAction::Ignore) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+const CXXOperatorCallExpr *GetSubscriptOrCallOperator(CXCursor Cursor) {
+  if (!clang_isExpression(Cursor.kind))
+    return nullptr;
+
+  const Expr *E = getCursorExpr(Cursor);
+  if (const auto *OCE = dyn_cast<CXXOperatorCallExpr>(E)) {
+    const OverloadedOperatorKind Kind = OCE->getOperator();
+    if (Kind == OO_Call || Kind == OO_Subscript)
+      return OCE;
+  }
+
+  return nullptr;
+}
+
+AnnotateTokensWorker::PostChildrenActions
+AnnotateTokensWorker::DetermineChildActions(CXCursor Cursor) const {
+  PostChildrenActions actions;
+
+  // The DeclRefExpr of CXXOperatorCallExpr refering to the custom operator is
+  // visited before the arguments to the operator call. For the Call and
+  // Subscript operator the range of this DeclRefExpr includes the whole call
+  // expression, so that all tokens in that range would be mapped to the
+  // operator function, including the tokens of the arguments. To avoid that,
+  // ensure to visit this DeclRefExpr as last node.
+  if (const auto *OCE = GetSubscriptOrCallOperator(Cursor)) {
+    const Expr *Callee = OCE->getCallee();
+    if (const ImplicitCastExpr *ICE = dyn_cast<ImplicitCastExpr>(Callee)) {
+      const Expr *SubExpr = ICE->getSubExpr();
+      if (const DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(SubExpr)) {
+        const Decl *parentDecl = getCursorParentDecl(Cursor);
+        CXTranslationUnit TU = clang_Cursor_getTranslationUnit(Cursor);
+
+        // Visit the DeclRefExpr as last.
+        CXCursor cxChild = MakeCXCursor(DRE, parentDecl, TU);
+        actions.push_back({cxChild, PostChildrenAction::Postpone});
+
+        // The parent of the DeclRefExpr, an ImplicitCastExpr, has an equally
+        // wide range as the DeclRefExpr. We can skip visiting this entirely.
+        cxChild = MakeCXCursor(ICE, parentDecl, TU);
+        actions.push_back({cxChild, PostChildrenAction::Ignore});
+      }
+    }
+  }
+
+  return actions;
+}
+
 static inline void updateCursorAnnotation(CXCursor &Cursor,
                                           const CXCursor &updateC) {
   if (clang_isInvalid(updateC.kind) || !clang_isInvalid(Cursor.kind))
@@ -6891,7 +6965,10 @@ AnnotateTokensWorker::Visit(CXCursor cursor, CXCursor parent) {
   SourceRange cursorRange = getRawCursorExtent(cursor);
   if (cursorRange.isInvalid())
     return CXChildVisit_Recurse;
-      
+
+  if (IsIgnoredChildCursor(cursor))
+    return CXChildVisit_Continue;
+
   if (!HasContextSensitiveKeywords) {
     // Objective-C properties can have context-sensitive keywords.
     if (cursor.kind == CXCursor_ObjCPropertyDecl) {
@@ -7039,6 +7116,7 @@ AnnotateTokensWorker::Visit(CXCursor cursor, CXCursor parent) {
   Info.CursorRange = cursorRange;
   Info.BeforeReachingCursorIdx = BeforeReachingCursorIdx;
   Info.BeforeChildrenTokenIdx = NextToken();
+  Info.ChildActions = DetermineChildActions(cursor);
   PostChildrenInfos.push_back(Info);
 
   return CXChildVisit_Recurse;
@@ -7051,6 +7129,8 @@ bool AnnotateTokensWorker::postVisitChildren(CXCursor cursor) {
   if (!clang_equalCursors(Info.Cursor, cursor))
     return false;
 
+  HandlePostPonedChildCursors(Info);
+
   const unsigned BeforeChildren = Info.BeforeChildrenTokenIdx;
   const unsigned AfterChildren = NextToken();
   SourceRange cursorRange = Info.CursorRange;
@@ -7077,6 +7157,56 @@ bool AnnotateTokensWorker::postVisitChildren(CXCursor cursor) {
   return false;
 }
 
+void AnnotateTokensWorker::HandlePostPonedChildCursors(
+    const PostChildrenInfo &Info) {
+  for (const auto &ChildAction : Info.ChildActions) {
+    if (ChildAction.action == PostChildrenAction::Postpone) {
+      HandlePostPonedChildCursor(ChildAction.cursor,
+                                 Info.BeforeChildrenTokenIdx);
+    }
+  }
+}
+
+void AnnotateTokensWorker::HandlePostPonedChildCursor(
+    CXCursor Cursor, unsigned StartTokenIndex) {
+  const auto flags = CXNameRange_WantQualifier | CXNameRange_WantQualifier;
+  unsigned I = StartTokenIndex;
+
+  // The bracket tokens of a Call or Subscript operator are mapped to
+  // CallExpr/CXXOperatorCallExpr because we skipped visiting the corresponding
+  // DeclRefExpr. Remap these tokens to the DeclRefExpr cursors.
+  for (unsigned RefNameRangeNr = 0; I < NumTokens; RefNameRangeNr++) {
+    const CXSourceRange CXRefNameRange =
+        clang_getCursorReferenceNameRange(Cursor, flags, RefNameRangeNr);
+    if (clang_Range_isNull(CXRefNameRange))
+      break; // All ranges handled.
+
+    SourceRange RefNameRange = cxloc::translateCXSourceRange(CXRefNameRange);
+    while (I < NumTokens) {
+      const SourceLocation TokenLocation = GetTokenLoc(I);
+      if (!TokenLocation.isValid())
+        break;
+
+      // Adapt the end range, because LocationCompare() reports
+      // RangeOverlap even for the not-inclusive end location.
+      const SourceLocation fixedEnd =
+          RefNameRange.getEnd().getLocWithOffset(-1);
+      RefNameRange = SourceRange(RefNameRange.getBegin(), fixedEnd);
+
+      const RangeComparisonResult ComparisonResult =
+          LocationCompare(SrcMgr, TokenLocation, RefNameRange);
+
+      if (ComparisonResult == RangeOverlap) {
+        Cursors[I++] = Cursor;
+      } else if (ComparisonResult == RangeBefore) {
+        ++I; // Not relevant token, check next one.
+      } else if (ComparisonResult == RangeAfter) {
+        break; // All tokens updated for current range, check next.
+      }
+    }
+  }
+}
+
 static enum CXChildVisitResult AnnotateTokensVisitor(CXCursor cursor,
                                                      CXCursor parent,
                                                      CXClientData client_data) {