[libclang] Make clang_findReferencesInFile also work on macros (find all expansions/definition

of a macro in a file).

As a bonus, also make searching for declarations more efficient by ignoring
preprocessing entities when we know that we are looking for a declaration.

Fixes rdar://10427411.

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

4 files changed, 136 insertions, 11 deletions

diff --git a/test/Index/file-macro-refs.c b/test/Index/file-macro-refs.c
new file mode 100644
index 0000000000..b37daaa433
--- /dev/null
+++ b/test/Index/file-macro-refs.c
@@ -0,0 +1,12 @@
+#define FOO
+
+FOO
+FOO
+
+// RUN: c-index-test -file-refs-at=%s:3:2 %s | FileCheck %s
+// RUN: CINDEXTEST_EDITING=1 c-index-test -file-refs-at=%s:3:2 %s | FileCheck %s
+
+// CHECK:      macro expansion=FOO:1:9
+// CHECK-NEXT: macro definition=FOO =[1:9 - 1:12]
+// CHECK-NEXT: macro expansion=FOO:1:9 =[3:1 - 3:4]
+// CHECK-NEXT: macro expansion=FOO:1:9 =[4:1 - 4:4]
diff --git a/tools/libclang/CIndex.cpp b/tools/libclang/CIndex.cpp
index 72abfa76ea..beb50a7c87 100644
--- a/tools/libclang/CIndex.cpp
+++ b/tools/libclang/CIndex.cpp
@@ -226,14 +226,13 @@ void CursorVisitor::visitFileRegion() {
   unsigned Offset = Begin.second;
   unsigned Length = End.second - Begin.second;
 
-  if (!VisitPreprocessorLast &&
-      Unit->getPreprocessor().getPreprocessingRecord())
-    visitPreprocessedEntitiesInRegion();
+  if (!VisitDeclsOnly && !VisitPreprocessorLast)
+    if (visitPreprocessedEntitiesInRegion())
+      return; // visitation break.
 
   visitDeclsFromFileRegion(File, Offset, Length);
 
-  if (VisitPreprocessorLast &&
-      Unit->getPreprocessor().getPreprocessingRecord())
+  if (!VisitDeclsOnly && VisitPreprocessorLast)
     visitPreprocessedEntitiesInRegion();
 }
 
@@ -354,6 +353,9 @@ void CursorVisitor::visitDeclsFromFileRegion(FileID File,
 }
 
 bool CursorVisitor::visitPreprocessedEntitiesInRegion() {
+  if (!AU->getPreprocessor().getPreprocessingRecord())
+    return false;
+
   PreprocessingRecord &PPRec
     = *AU->getPreprocessor().getPreprocessingRecord();
   SourceManager &SM = AU->getSourceManager();
diff --git a/tools/libclang/CIndexHigh.cpp b/tools/libclang/CIndexHigh.cpp
index 0b36af3c95..a2a585454e 100644
--- a/tools/libclang/CIndexHigh.cpp
+++ b/tools/libclang/CIndexHigh.cpp
@@ -168,7 +168,8 @@ static enum CXChildVisitResult findFileIdRefVisit(CXCursor cursor,
     if (SelIdLoc.isValid())
       Loc = SelIdLoc;
 
-    SourceManager &SM = data->getASTContext().getSourceManager();
+    ASTContext &Ctx = data->getASTContext();
+    SourceManager &SM = Ctx.getSourceManager();
     bool isInMacroDef = false;
     if (Loc.isMacroID()) {
       bool isMacroArg;
@@ -184,11 +185,11 @@ static enum CXChildVisitResult findFileIdRefVisit(CXCursor cursor,
     if (isInMacroDef) {
       // FIXME: For a macro definition make sure that all expansions
       // of it expand to the same reference before allowing to point to it.
-      Loc = SourceLocation();
+      return CXChildVisit_Recurse;
     }
 
     data->visitor.visit(data->visitor.context, cursor,
-                        cxloc::translateSourceRange(D->getASTContext(), Loc));
+                        cxloc::translateSourceRange(Ctx, Loc));
   }
   return CXChildVisit_Recurse;
 }
@@ -217,10 +218,104 @@ static void findIdRefsInFile(CXTranslationUnit TU, CXCursor declCursor,
                                   findFileIdRefVisit, &data,
                                   /*VisitPreprocessorLast=*/true,
                                   /*VisitIncludedEntities=*/false,
-                                  Range);
+                                  Range,
+                                  /*VisitDeclsOnly=*/true);
   FindIdRefsVisitor.visitFileRegion();
 }
 
+namespace {
+
+struct FindFileMacroRefVisitData {
+  ASTUnit &Unit;
+  const FileEntry *File;
+  const IdentifierInfo *Macro;
+  CXCursorAndRangeVisitor visitor;
+
+  FindFileMacroRefVisitData(ASTUnit &Unit, const FileEntry *File,
+                            const IdentifierInfo *Macro,
+                            CXCursorAndRangeVisitor visitor)
+    : Unit(Unit), File(File), Macro(Macro), visitor(visitor) { }
+
+  ASTContext &getASTContext() const {
+    return Unit.getASTContext();
+  }
+};
+
+} // anonymous namespace
+
+static enum CXChildVisitResult findFileMacroRefVisit(CXCursor cursor,
+                                                     CXCursor parent,
+                                                     CXClientData client_data) {
+  const IdentifierInfo *Macro = 0;
+  if (cursor.kind == CXCursor_MacroDefinition)
+    Macro = getCursorMacroDefinition(cursor)->getName();
+  else if (cursor.kind == CXCursor_MacroExpansion)
+    Macro = getCursorMacroExpansion(cursor)->getName();
+  if (!Macro)
+    return CXChildVisit_Continue;
+
+  FindFileMacroRefVisitData *data = (FindFileMacroRefVisitData *)client_data;
+  if (data->Macro != Macro)
+    return CXChildVisit_Continue;
+
+  SourceLocation
+    Loc = cxloc::translateSourceLocation(clang_getCursorLocation(cursor));
+
+  ASTContext &Ctx = data->getASTContext();
+  SourceManager &SM = Ctx.getSourceManager();
+  bool isInMacroDef = false;
+  if (Loc.isMacroID()) {
+    bool isMacroArg;
+    Loc = getFileSpellingLoc(SM, Loc, isMacroArg);
+    isInMacroDef = !isMacroArg;
+  }
+
+  // We are looking for identifiers in a specific file.
+  std::pair<FileID, unsigned> LocInfo = SM.getDecomposedLoc(Loc);
+  if (SM.getFileEntryForID(LocInfo.first) != data->File)
+    return CXChildVisit_Continue;
+
+  if (isInMacroDef) {
+    // FIXME: For a macro definition make sure that all expansions
+    // of it expand to the same reference before allowing to point to it.
+    return CXChildVisit_Continue;
+  }
+
+  data->visitor.visit(data->visitor.context, cursor,
+                      cxloc::translateSourceRange(Ctx, Loc));
+  return CXChildVisit_Continue;
+}
+
+static void findMacroRefsInFile(CXTranslationUnit TU, CXCursor Cursor,
+                                const FileEntry *File,
+                                CXCursorAndRangeVisitor Visitor) {
+  if (Cursor.kind != CXCursor_MacroDefinition &&
+      Cursor.kind != CXCursor_MacroExpansion)
+    return;
+
+  ASTUnit *Unit = static_cast<ASTUnit*>(TU->TUData);
+  SourceManager &SM = Unit->getSourceManager();
+
+  FileID FID = SM.translateFile(File);
+  const IdentifierInfo *Macro = 0;
+  if (Cursor.kind == CXCursor_MacroDefinition)
+    Macro = getCursorMacroDefinition(Cursor)->getName();
+  else
+    Macro = getCursorMacroExpansion(Cursor)->getName();
+  if (!Macro)
+    return;
+
+  FindFileMacroRefVisitData data(*Unit, File, Macro, Visitor);
+
+  SourceRange Range(SM.getLocForStartOfFile(FID), SM.getLocForEndOfFile(FID));
+  CursorVisitor FindMacroRefsVisitor(TU,
+                                  findFileMacroRefVisit, &data,
+                                  /*VisitPreprocessorLast=*/false,
+                                  /*VisitIncludedEntities=*/false,
+                                  Range);
+  FindMacroRefsVisitor.visitPreprocessedEntitiesInRegion();
+}
+
 
 //===----------------------------------------------------------------------===//
 // libclang public APIs.
@@ -248,6 +343,15 @@ void clang_findReferencesInFile(CXCursor cursor, CXFile file,
     return;
   }
 
+  if (cursor.kind == CXCursor_MacroDefinition ||
+      cursor.kind == CXCursor_MacroExpansion) {
+    findMacroRefsInFile(cxcursor::getCursorTU(cursor),
+                        cursor,
+                        static_cast<const FileEntry *>(file),
+                        visitor);
+    return;
+  }
+
   // We are interested in semantics of identifiers so for C++ constructor exprs
   // prefer type references, e.g.:
   //
diff --git a/tools/libclang/CursorVisitor.h b/tools/libclang/CursorVisitor.h
index d3fa70d8a2..ec08655154 100644
--- a/tools/libclang/CursorVisitor.h
+++ b/tools/libclang/CursorVisitor.h
@@ -83,6 +83,10 @@ class CursorVisitor : public DeclVisitor<CursorVisitor, bool>,
   /// its search.
   SourceRange RegionOfInterest;
 
+  /// \brief Whether we should only visit declarations and not preprocessing
+  /// record entries.
+  bool VisitDeclsOnly;
+
   // FIXME: Eventually remove.  This part of a hack to support proper
   // iteration over all Decls contained lexically within an ObjC container.
   DeclContext::decl_iterator *DI_current;
@@ -131,12 +135,15 @@ public:
                 CXClientData ClientData,
                 bool VisitPreprocessorLast,
                 bool VisitIncludedPreprocessingEntries = false,
-                SourceRange RegionOfInterest = SourceRange())
+                SourceRange RegionOfInterest = SourceRange(),
+                bool VisitDeclsOnly = false)
     : TU(TU), AU(static_cast<ASTUnit*>(TU->TUData)),
       Visitor(Visitor), ClientData(ClientData),
       VisitPreprocessorLast(VisitPreprocessorLast),
       VisitIncludedEntities(VisitIncludedPreprocessingEntries),
-      RegionOfInterest(RegionOfInterest), DI_current(0), FileDI_current(0)
+      RegionOfInterest(RegionOfInterest),
+      VisitDeclsOnly(VisitDeclsOnly),
+      DI_current(0), FileDI_current(0)
   {
     Parent.kind = CXCursor_NoDeclFound;
     Parent.data[0] = 0;