1 files changed, 150 insertions, 45 deletions

diff --git a/clang/tools/libclang/CIndex.cpp b/clang/tools/libclang/CIndex.cpp
index 173aa237861..61cc113bf59 100644
--- a/clang/tools/libclang/CIndex.cpp
+++ b/clang/tools/libclang/CIndex.cpp
@@ -4424,7 +4424,7 @@ CXString clang_getTokenSpelling(CXTranslationUnit TU, CXToken CXTok) {
 
   SourceLocation Loc = SourceLocation::getFromRawEncoding(CXTok.int_data[1]);
   std::pair<FileID, unsigned> LocInfo
-    = CXXUnit->getSourceManager().getDecomposedLoc(Loc);
+    = CXXUnit->getSourceManager().getDecomposedSpellingLoc(Loc);
   bool Invalid = false;
   StringRef Buffer
     = CXXUnit->getSourceManager().getBufferData(LocInfo.first, &Invalid);
@@ -4580,6 +4580,16 @@ class AnnotateTokensWorker {
   SourceLocation GetTokenLoc(unsigned tokI) {
     return SourceLocation::getFromRawEncoding(Tokens[tokI].int_data[1]);
   }
+  bool isMacroArgToken(unsigned tokI) const {
+    return Tokens[tokI].int_data[3] != 0;
+  }
+  SourceLocation getMacroArgLoc(unsigned tokI) const {
+    return SourceLocation::getFromRawEncoding(Tokens[tokI].int_data[3]);
+  }
+
+  void annotateAndAdvanceTokens(CXCursor, RangeComparisonResult, SourceRange);
+  void annotateAndAdvanceMacroArgTokens(CXCursor, RangeComparisonResult,
+                                        SourceRange);
 
 public:
   AnnotateTokensWorker(AnnotateTokensData &annotated,
@@ -4632,6 +4642,63 @@ void AnnotateTokensWorker::AnnotateTokens(CXCursor parent) {
   }
 }
 
+/// \brief It annotates and advances tokens with a cursor until the comparison
+//// between the cursor location and the source range is the same as
+/// \arg compResult.
+///
+/// Pass RangeBefore to annotate tokens with a cursor until a range is reached.
+/// Pass RangeOverlap to annotate tokens inside a range.
+void AnnotateTokensWorker::annotateAndAdvanceTokens(CXCursor updateC,
+                                               RangeComparisonResult compResult,
+                                               SourceRange range) {
+  while (MoreTokens()) {
+    const unsigned I = NextToken();
+    if (isMacroArgToken(I))
+      return annotateAndAdvanceMacroArgTokens(updateC, compResult, range);
+
+    SourceLocation TokLoc = GetTokenLoc(I);
+    if (LocationCompare(SrcMgr, TokLoc, range) == compResult) {
+      Cursors[I] = updateC;
+      AdvanceToken();
+      continue;
+    }
+    break;
+  }
+}
+
+/// \brief Special annotation handling for macro argument tokens.
+void AnnotateTokensWorker::annotateAndAdvanceMacroArgTokens(CXCursor updateC,
+                                               RangeComparisonResult compResult,
+                                               SourceRange range) {
+  assert(isMacroArgToken(NextToken()) &&
+         "Should be called only for macro arg tokens");
+
+  // This works differently than annotateAndAdvanceTokens; because expanded
+  // macro arguments can have arbitrary translation-unit source order, we do not
+  // advance the token index one by one until a token fails the range test.
+  // We only advance once past all of the macro arg tokens if all of them
+  // pass the range test. If one of them fails we keep the token index pointing
+  // at the start of the macro arg tokens so that the failing token will be
+  // annotated by a subsequent annotation try.
+
+  bool atLeastOneCompFail = false;
+  
+  unsigned I = NextToken();
+  for (; isMacroArgToken(I); ++I) {
+    SourceLocation TokLoc = getMacroArgLoc(I);
+    if (TokLoc.isFileID())
+      continue; // not macro arg token, it's parens or comma.
+    if (LocationCompare(SrcMgr, TokLoc, range) == compResult) {
+      if (clang_isInvalid(clang_getCursorKind(Cursors[I])))
+        Cursors[I] = updateC;
+    } else
+      atLeastOneCompFail = true;
+  }
+
+  if (!atLeastOneCompFail)
+    TokIdx = I; // All of the tokens were handled, advance beyond all of them.
+}
+
 enum CXChildVisitResult
 AnnotateTokensWorker::Visit(CXCursor cursor, CXCursor parent) {  
   CXSourceLocation Loc = clang_getCursorLocation(cursor);
@@ -4783,20 +4850,7 @@ AnnotateTokensWorker::Visit(CXCursor cursor, CXCursor parent) {
     (clang_isInvalid(K) || K == CXCursor_TranslationUnit)
      ? clang_getNullCursor() : parent;
 
-  while (MoreTokens()) {
-    const unsigned I = NextToken();
-    SourceLocation TokLoc = GetTokenLoc(I);
-    switch (LocationCompare(SrcMgr, TokLoc, cursorRange)) {
-      case RangeBefore:
-        Cursors[I] = updateC;
-        AdvanceToken();
-        continue;
-      case RangeAfter:
-      case RangeOverlap:
-        break;
-    }
-    break;
-  }
+  annotateAndAdvanceTokens(updateC, RangeBefore, cursorRange);
 
   // Avoid having the cursor of an expression "overwrite" the annotation of the
   // variable declaration that it belongs to.
@@ -4821,46 +4875,19 @@ AnnotateTokensWorker::Visit(CXCursor cursor, CXCursor parent) {
   VisitChildren(cursor);
   const unsigned AfterChildren = NextToken();
 
-  // Adjust 'Last' to the last token within the extent of the cursor.
-  while (MoreTokens()) {
-    const unsigned I = NextToken();
-    SourceLocation TokLoc = GetTokenLoc(I);
-    switch (LocationCompare(SrcMgr, TokLoc, cursorRange)) {
-      case RangeBefore:
-        assert(0 && "Infeasible");
-      case RangeAfter:
-        break;
-      case RangeOverlap:
-        Cursors[I] = updateC;
-        AdvanceToken();
-        continue;
-    }
-    break;
-  }
-  const unsigned Last = NextToken();
+  // 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 AST elements within macros, rely on a post-annotate pass to
-  // to correctly annotate the tokens with cursors.  Otherwise we can
-  // get confusing results of having tokens that map to cursors that really
-  // are expanded by an instantiation.
-  if (L.isMacroID())
-    cursor = clang_getNullCursor();
-
   for (unsigned I = BeforeChildren; I != AfterChildren; ++I) {
     if (!clang_isInvalid(clang_getCursorKind(Cursors[I])))
       break;
     
     Cursors[I] = cursor;
   }
-  // Scan the tokens that are at the end of the cursor, but are not captured
-  // but the child cursors.
-  for (unsigned I = AfterChildren; I != Last; ++I)
-    Cursors[I] = cursor;
 
-  TokIdx = Last;
   return CXChildVisit_Continue;
 }
 
@@ -4871,6 +4898,74 @@ static enum CXChildVisitResult AnnotateTokensVisitor(CXCursor cursor,
 }
 
 namespace {
+
+/// \brief Uses the macro expansions in the preprocessing record to find
+/// and mark tokens that are macro arguments. This info is used by the
+/// AnnotateTokensWorker.
+class MarkMacroArgTokensVisitor {
+  SourceManager &SM;
+  CXToken *Tokens;
+  unsigned NumTokens;
+  unsigned CurIdx;
+  
+public:
+  MarkMacroArgTokensVisitor(SourceManager &SM,
+                            CXToken *tokens, unsigned numTokens)
+    : SM(SM), Tokens(tokens), NumTokens(numTokens), CurIdx(0) { }
+
+  CXChildVisitResult visit(CXCursor cursor, CXCursor parent) {
+    if (cursor.kind != CXCursor_MacroExpansion)
+      return CXChildVisit_Continue;
+
+    SourceRange macroRange = getCursorMacroExpansion(cursor)->getSourceRange();
+    if (macroRange.getBegin() == macroRange.getEnd())
+      return CXChildVisit_Continue; // it's not a function macro.
+
+    for (; CurIdx < NumTokens; ++CurIdx) {
+      if (!SM.isBeforeInTranslationUnit(getTokenLoc(CurIdx),
+                                        macroRange.getBegin()))
+        break;
+    }
+    
+    if (CurIdx == NumTokens)
+      return CXChildVisit_Break;
+
+    for (; CurIdx < NumTokens; ++CurIdx) {
+      SourceLocation tokLoc = getTokenLoc(CurIdx);
+      if (!SM.isBeforeInTranslationUnit(tokLoc, macroRange.getEnd()))
+        break;
+
+      setMacroArgExpandedLoc(CurIdx, SM.getMacroArgExpandedLocation(tokLoc));
+    }
+
+    if (CurIdx == NumTokens)
+      return CXChildVisit_Break;
+
+    return CXChildVisit_Continue;
+  }
+
+private:
+  SourceLocation getTokenLoc(unsigned tokI) {
+    return SourceLocation::getFromRawEncoding(Tokens[tokI].int_data[1]);
+  }
+
+  void setMacroArgExpandedLoc(unsigned tokI, SourceLocation loc) {
+    // The third field is reserved and currently not used. Use it here
+    // to mark macro arg expanded tokens with their expanded locations.
+    Tokens[tokI].int_data[3] = loc.getRawEncoding();
+  }
+};
+
+} // end anonymous namespace
+
+static CXChildVisitResult
+MarkMacroArgTokensVisitorDelegate(CXCursor cursor, CXCursor parent,
+                                  CXClientData client_data) {
+  return static_cast<MarkMacroArgTokensVisitor*>(client_data)->visit(cursor,
+                                                                     parent);
+}
+
+namespace {
   struct clang_annotateTokens_Data {
     CXTranslationUnit TU;
     ASTUnit *CXXUnit;
@@ -4960,6 +5055,16 @@ static void clang_annotateTokensImpl(void *UserData) {
     }
   }
   
+  if (CXXUnit->getPreprocessor().getPreprocessingRecord()) {
+    // Search and mark tokens that are macro argument expansions.
+    MarkMacroArgTokensVisitor Visitor(CXXUnit->getSourceManager(),
+                                      Tokens, NumTokens);
+    CursorVisitor MacroArgMarker(TU,
+                                 MarkMacroArgTokensVisitorDelegate, &Visitor,
+                                 Decl::MaxPCHLevel, true, RegionOfInterest);
+    MacroArgMarker.visitPreprocessedEntitiesInRegion();
+  }
+  
   // Annotate all of the source locations in the region of interest that map to
   // a specific cursor.
   AnnotateTokensWorker W(Annotated, Tokens, Cursors, NumTokens,