[clangd] Handle references for patched macros

[kadircet]

Summary

Depends on D80198.

Macro references in main file are collected separetely and stored as a
map from macro's symbol id to reference ranges. Those ranges are
computed inside PPCallbacks, hence we don't have access to TokenBuffer.

In presence of patched macro definitions, this symbol id is computed
from the definition inside the patch, rather than the definition inside
the preamble section. Also the ranges for occurences inside the patch
would be off for similar reasons.

Hence when we try to findReferences for a macro defined in main file we
either need to:
- Make use of definition inside the patch, rather than the real
  definition inside the preamble section.
- Lex preamble section while collecting main file references to figure
  out macro's real location and attribute references to it.

The former implies that we will miss all the references inside the
patch. But it doesn't introduce any new logic. We basically get all the
references inside main file section for free.

The latter doesn't trade any accuracy, we'll get all the references
including the ones inside the preamble section. But requires either:
- Lexing the preamble section to figure out "real range" of a patched
  macro definition
- Postponing range/location calculations until a later step in which we
  have access to tokenbuffers.

This patch trades some accuracy in favor of code complexity, and just
uses the definition inside the patch when we are looking for references
of a macro.

[kadircet]

• Polish

clang-tools-extra/clangd/XRefs.cpp

ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
   if (const auto *IdentifierAtCursor =
           syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens())) {
     Macro = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor(),
                           AST.getTokens());
   }
 
   RefsRequest Req;
   if (Macro) {
-    // Handle references to macro.
-    if (auto MacroSID = getSymbolID(Macro->Name, Macro->IdentLoc, SM)) {
+    // Handle references to macro, we are using DefRange instead of IdentLoc as
+    // references are collected without #line directive mappings.
+    if (auto MacroSID =
+            getSymbolID(Macro->Name, Macro->DefRange.getBegin(), SM)) {
       // Collect macro references from main file.
       const auto &IDToRefs = AST.getMacros().MacroRefs;
       auto Refs = IDToRefs.find(*MacroSID);
       if (Refs != IDToRefs.end()) {
         for (const auto Ref : Refs->second) {
           Location Result;
           Result.range = Ref;
           Result.uri = URIMainFile;

clang-tools-extra/clangd/unittests/PreambleTests.cpp

 //===--- PreambleTests.cpp --------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
 #include "Compiler.h"
 #include "Headers.h"
 #include "Hover.h"
 #include "Preamble.h"
 #include "SourceCode.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "XRefs.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <clang/Frontend/FrontendActions.h>
 #include <memory>
 #include <string>
 #include <vector>
 
 using testing::Contains;
 using testing::Field;
+using testing::Matcher;
 
 namespace clang {
 namespace clangd {
 namespace {
 
 MATCHER_P2(Distance, File, D, "") {
   return arg.first() == File && arg.second == D;
 }
TEST(PreamblePatchTest, LocateMacroAtDeletion) {
     ASSERT_TRUE(AST);
 
     auto HI = getHover(*AST, Modified.point(), format::getLLVMStyle(), nullptr);
     ASSERT_TRUE(HI);
     EXPECT_THAT(HI->Definition, "#define BAR");
   }
 }
 
+TEST(PreamblePatchTest, RefsToMacros) {
+  struct {
+    llvm::StringLiteral Baseline;
+    llvm::StringLiteral Modified;
+  } Cases[] = {
+      // Newly added
+      {
+          "",
+          R"cpp(
+            #define ^FOO
+            ^[[FOO]])cpp",
+      },
+      // Moved around
+      {
+          "#define FOO",
+          R"cpp(
+            #define BAR
+            #define ^FOO
+            ^[[FOO]])cpp",
+      },
+      // Ref in preamble section
+      {
+          "",
+          R"cpp(
+            #define ^FOO
+            #undef ^FOO)cpp",
+      },
+  };
+
+  for (const auto &Case : Cases) {
+    Annotations Modified(Case.Modified);
+    auto AST = createPatchedAST("", Modified.code());
+    ASSERT_TRUE(AST);
+
+    const auto &SM = AST->getSourceManager();
+    std::vector<Matcher<Location>> ExpectedLocations;
+    for (const auto &R : Modified.ranges())
+      ExpectedLocations.push_back(Field(&Location::range, R));
+
+    for (const auto &P : Modified.points()) {
+      auto *MacroTok = AST->getTokens().spelledTokenAt(SM.getComposedLoc(
+          SM.getMainFileID(),
+          llvm::cantFail(positionToOffset(Modified.code(), P))));
+      ASSERT_TRUE(MacroTok);
+      EXPECT_THAT(findReferences(*AST, P, 0).References,
+                  testing::ElementsAreArray(ExpectedLocations));
+    }
+  }
+}
 } // namespace
 } // namespace clangd
 } // namespace clang