From decdbc1155f5120554269319b1c77675bac9151c Mon Sep 17 00:00:00 2001
From: Haojian Wu <hokein.wu@gmail.com>
Date: Mon, 9 Dec 2019 12:00:42 +0100
Subject: [PATCH] [clangd] Use expansion location when the ref is inside
 macros.

Summary:
Previously, xrefs has inconsistent behavior when the reference is inside
macro body:
- AST-based xrefs (for main file) uses the expansion location;
- our index uses the spelling location;

This patch makes our index use file locations for references, which is
consistent with AST-based xrefs, and kythe as well.

After this patch, memory usage of static index on LLVM increases ~5%.

Reviewers: ilya-biryukov

Subscribers: merge_guards_bot, MaskRay, jkorous, arphaman, kadircet, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D70480
---
 clang-tools-extra/clangd/index/SymbolCollector.cpp | 12 +++++++----
 .../clangd/unittests/SymbolCollectorTests.cpp      | 23 ++++++++++++++++++++++
 clang-tools-extra/clangd/unittests/XRefsTests.cpp  | 13 ++++++++++++
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index 191cd68..6775ccc 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -276,10 +276,9 @@ bool SymbolCollector::handleDeclOccurence(
   // Mark D as referenced if this is a reference coming from the main file.
   // D may not be an interesting symbol, but it's cheaper to check at the end.
   auto &SM = ASTCtx->getSourceManager();
-  auto SpellingLoc = SM.getSpellingLoc(Loc);
   if (Opts.CountReferences &&
       (Roles & static_cast<unsigned>(index::SymbolRole::Reference)) &&
-      SM.getFileID(SpellingLoc) == SM.getMainFileID())
+      SM.getFileID(SM.getSpellingLoc(Loc)) == SM.getMainFileID())
     ReferencedDecls.insert(ND);
 
   auto ID = getSymbolID(ND);
@@ -312,9 +311,14 @@ bool SymbolCollector::handleDeclOccurence(
       !shouldCollectSymbol(*ND, *ASTCtx, Opts, IsMainFileOnly))
     return true;
   // Do not store references to main-file symbols.
+  // Unlike other fields, e.g. Symbols (which use spelling locations), we use
+  // file locations for references (as it aligns the behavior of clangd's
+  // AST-based xref).
+  // FIXME: we should try to use the file locations for other fields.
   if (CollectRef && !IsMainFileOnly && !isa<NamespaceDecl>(ND) &&
-      (Opts.RefsInHeaders || SM.getFileID(SpellingLoc) == SM.getMainFileID()))
-    DeclRefs[ND].emplace_back(SpellingLoc, Roles);
+      (Opts.RefsInHeaders ||
+       SM.getFileID(SM.getFileLoc(Loc)) == SM.getMainFileID()))
+    DeclRefs[ND].emplace_back(SM.getFileLoc(Loc), Roles);
   // Don't continue indexing if this is a mere reference.
   if (IsOnlyRef)
     return true;
diff --git a/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp b/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
index abc7aa3..b5876b15 100644
--- a/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
+++ b/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
@@ -718,6 +718,29 @@ TEST_F(SymbolCollectorTest, NameReferences) {
                                   HaveRanges(Header.ranges()))));
 }
 
+TEST_F(SymbolCollectorTest, RefsOnMacros) {
+  // Refs collected from SymbolCollector behave in the same way as
+  // AST-based xrefs.
+  CollectorOpts.RefFilter = RefKind::All;
+  CollectorOpts.RefsInHeaders = true;
+  Annotations Header(R"(
+    #define TYPE(X) X
+    #define FOO Foo
+    #define CAT(X, Y) X##Y
+    class [[Foo]] {};
+    void test() {
+      TYPE([[Foo]]) foo;
+      [[FOO]] foo2;
+      TYPE(TYPE([[Foo]])) foo3;
+      [[CAT]](Fo, o) foo4;
+    }
+  )");
+  CollectorOpts.RefFilter = RefKind::All;
+  runSymbolCollector(Header.code(), "");
+  EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "Foo").ID,
+                                  HaveRanges(Header.ranges()))));
+}
+
 TEST_F(SymbolCollectorTest, HeaderAsMainFile) {
   CollectorOpts.RefFilter = RefKind::All;
   Annotations Header(R"(
diff --git a/clang-tools-extra/clangd/unittests/XRefsTests.cpp b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
index ee23522..e009c1c 100644
--- a/clang-tools-extra/clangd/unittests/XRefsTests.cpp
+++ b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
@@ -811,6 +811,19 @@ TEST(FindReferences, WithinAST) {
         } // namespace ns
         int main() { [[^ns]]::Foo foo; }
       )cpp",
+
+      R"cpp(// Macros
+        #define TYPE(X) X
+        #define FOO Foo
+        #define CAT(X, Y) X##Y
+        class [[Fo^o]] {};
+        void test() {
+          TYPE([[Foo]]) foo;
+          [[FOO]] foo2;
+          TYPE(TYPE([[Foo]])) foo3;
+          [[CAT]](Fo, o) foo4;
+        }
+      )cpp",
   };
   for (const char *Test : Tests) {
     Annotations T(Test);
-- 
2.7.4