[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/SourceCode.h

 struct DefinedMacro {
   llvm::StringRef Name;
   CharSourceRange DefRange;
 };
 /// Gets the macro referenced by \p SpelledTok. It must be a spelled token
 /// aligned to the beginning of an identifier.
 llvm::Optional<DefinedMacro> locateMacroAt(const syntax::Token &SpelledTok,
                                            Preprocessor &PP,
-                                           const syntax::TokenBuffer &TB);
+                                           const syntax::TokenBuffer &TB,
+                                           bool UseLineDirectives = true);
 
 /// Infers whether this is a header from the FileName and LangOpts (if
 /// presents).
 bool isHeaderFile(llvm::StringRef FileName,
                   llvm::Optional<LangOptions> LangOpts = llvm::None);
 
 /// Returns true if the given location is in a generated protobuf file.
 bool isProtoFile(SourceLocation Loc, const SourceManager &SourceMgr);
 
 /// Checks whether \p FileName is a valid spelling of main file.
 bool isMainFile(llvm::StringRef FileName, const SourceManager &SM);
 
 } // namespace clangd
 } // namespace clang
 #endif

clang-tools-extra/clangd/SourceCode.cpp

 
 bool isMainFile(llvm::StringRef FileName, const SourceManager &SM) {
   auto FE = SM.getFileManager().getFile(FileName);
   return FE && *FE == SM.getFileEntryForID(SM.getMainFileID());
 }
 
 llvm::Optional<DefinedMacro> locateMacroAt(const syntax::Token &SpelledTok,
                                            Preprocessor &PP,
-                                           const syntax::TokenBuffer &TB) {
+                                           const syntax::TokenBuffer &TB,
+                                           bool UseLineDirectives) {
   SourceLocation Loc = SpelledTok.location();
   assert(Loc.isFileID());
   const auto &SM = PP.getSourceManager();
   IdentifierInfo *IdentifierInfo = PP.getIdentifierInfo(SpelledTok.text(SM));
   if (!IdentifierInfo || !IdentifierInfo->hadMacroDefinition())
     return None;
 
   // Get the definition just before the searched location so that a macro
   // referenced in a '#undef MACRO' can still be found. Note that we only do
   // that if Loc is not pointing at start of file.
   if (SM.getLocForStartOfFile(SM.getFileID(Loc)) != Loc)
     Loc = Loc.getLocWithOffset(-1);
   MacroDefinition MacroDef = PP.getMacroDefinitionAtLoc(IdentifierInfo, Loc);
   auto *MI = MacroDef.getMacroInfo();
   if (!MI)
     return None;
 
   SourceRange DefRange(MI->getDefinitionLoc(), MI->getDefinitionEndLoc());
+  if (!UseLineDirectives)
+    return DefinedMacro{IdentifierInfo->getName(),
+                        CharSourceRange::getTokenRange(DefRange)};
 
   // Macro definitions could be injected through preamble patch. These contain
   // line directives to hint their original location in main file.
   auto DefFile = SM.getFileID(DefRange.getBegin());
   auto IncludeLoc = SM.getIncludeLoc(DefFile);
   // Preamble patch is included inside the builtin file.
   if (IncludeLoc.isValid() && SM.isWrittenInBuiltinFile(IncludeLoc)) {
     auto Presumed = SM.getPresumedLoc(DefRange.getBegin());

clang-tools-extra/clangd/XRefs.cpp

ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());
     return {};
   }
   llvm::Optional<DefinedMacro> Macro;
   if (const auto *IdentifierAtCursor =
           syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens())) {
     Macro = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor(),
-                          AST.getTokens());
+                          AST.getTokens(), false);
   }
 
   RefsRequest Req;
   if (Macro) {
     // Handle references to macro.
     if (auto MacroSID =
             getSymbolID(Macro->Name, Macro->DefRange.getBegin(), SM)) {
       // Collect macro references from main file.

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 "Preamble.h"
+#include "SourceCode.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "XRefs.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(MacroTok);
 
   auto FoundMacro =
       locateMacroAt(*MacroTok, AST->getPreprocessor(), AST->getTokens());
   ASSERT_TRUE(FoundMacro);
   EXPECT_THAT(FoundMacro->Name, "FOO");
 }
 
+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",
+      },
+  };
+
+  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