[clangd] locateMacroAt handles patched macros

[kadircet]

Summary

Depends on D79992.

This patch changes locateMacroAt to perform #line directive substitution
for macro identifier locations.

We first check whether a location is inside a file included through
built-in header. If so we check whether line directive maps it back to
the main file, and afterwards use TokenBuffers to find exact location of
the identifier on the line.

Instead of performing the mapping in locateMacroAt, we could also store
a mapping inside the ParsedAST whenever we use a patched preamble. But
that would imply adding more responsibility to ParsedAST and paying for
the mapping even when it is not going to be used.

====

Go-To-Definition:

Later on these locations are used for serving go-to-definition requests,
this enables jumping to definition inside the preamble section in
presence of patched macros.

=====

Go-To-Refs:

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 preamble patch, any reference to a macro inside the
preamble section will unfortunately have the wrong range. They'll point
into the patch rather than the main file. Hence during findReferences,
we won't get any ranges reported for those.

Fixing those requires:
- 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. We don't do
any patching for references inside the preamble patch but get any
reference inside the main file for free.

[kadircet]

• Polish

[kadircet]

• Rebase

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

975

OK, this is fairly horrible...

I want to say the preamble patch location translation should be a separate function rather than coupled to macro-specific stuff.
(Then we don't even need the extra struct member, but we can still keep it to "remind" people the translation is needed, if we like).

But of course we didn't preserve the formatting in the preamble patch, so getPresumedLoc() doesn't give us the right location, it just gives us something on the right line that we then need to re-parse...

This really isn't looking like a great tradeoff to me anymore (and I know I suggested it).
How much work do you think it is (compared to say, the logic here plus doing it in ~2 more places) to give the preamblepatch the invariant that the PresumedLoc has a usable col as well.
The original implementation did padding with spaces, but I guess we could as well have the preamble patching stuff splice the actual source code...

clang-tools-extra/clangd/SourceCode.h

295

It's not obvious to me that this comment describes a special case but the field is valid in general.
Maybe

/// Location of the identifier that names the macro.
/// Unlike Info->Location, this translates preamble-patch locations to main-file locations.

297

nit: call this NameLoc or NameLocation?
ident is *slightly* vague

[kadircet]

Instead of re-lexing at use time, put macro identifier on correct column while
patching and make use of presumed locations at use time.

[sammccall]

clang-tools-extra/clangd/Preamble.cpp

123

Spells directive in \p DirectiveRange while prepending it with \p Prefix

I wouldn't understand what this meant if I didn't already know :-) Maybe:

// Formats a PP directive consisting of Prefix (e.g. "#define ") and Body ("X 10").
// The formatting is copied so that the tokens in Body have PresumedLocs with
// correct columns and lines.

127

nit: move out-param to end?

133

This is only going to give you something useful for file IDs, not macro expansions. Whose responsibility is it to make sure there's no macro IDs here?

134

this name seems a bit confusing to me, since it only covers one case.
TargetColumn?
if (TargetColumn < Prefix.size()) reads a lot celarer I think.

136

again this comment only applyes to one branch

139

Rather than "pad from the start of a new line" it might be clearer to give an example here:

... We produce e.g.:
#line N-1
#define \
   X 10

149

this asserts isValidFileRange, so we need this to be true :-)

clang-tools-extra/clangd/SourceCode.cpp

976

I think this block can now be lifted to a function in SourceCode or Preamble or so?

[sammccall]

(BTW this looks miles better to me, thanks for persistence!)

[kadircet]

• Extract preamble-patch location translation to a helper:
  • Migrate usage in include collector
  • Make preamble patch header name part of the check.
• Handle macro ids in spellDirective
• Tweak comments/names

[kadircet]

clang-tools-extra/clangd/Preamble.cpp

133

For define directives this is already true as MacroInfo's definition range is a file range(as macro bodies are lexed without macro expansions).

For any other directive with macro expansions, we only care about the expansion locations as we want to spell out the directive as written. So we can just map back to expansion locs here.

136

why ? we pad in both cases, starting from the beginning of the line in first branch and starting after Prefix in the second.

clang-tools-extra/clangd/SourceCode.cpp

975

as discussed offline, changed patching logic to put directive in correct column in addition to line. Now we can just make use of presumedloc.

[sammccall]

Looks great, ship it!

clang-tools-extra/clangd/Preamble.cpp

153

nit: not "before the directive", just before the DirectiveRange.begin()

156

nit: I think it'd be easier to understand the other case if this one came first, as it's the fallback

clang-tools-extra/clangd/Preamble.h

136 ↗

(On Diff #266111)

Can we have a unit-test for this function?
I think the minimal thing (least entangled with PreamblePatch construction) is to manually write a snippet that looks like a preamble patch like

# line 3 main.cpp
int foo();

map it into a TestTU with AdditionalFiles, set ExtraArgs = {"-include","patch.h"} and then get the location with findDecl

[kadircet]

• Add tests for preamble-patch translation logic

[kadircet]

• Rebase

[kadircet]

• Rebase

clang-tools-extra/clangd/Headers.cpp

 #include "clang/Lex/HeaderSearch.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
-bool isMainFile(llvm::StringRef FileName, const SourceManager &SM) {
-  auto FE = SM.getFileManager().getFile(FileName);
-  return FE && *FE == SM.getFileEntryForID(SM.getMainFileID());
-}
-
 class RecordHeaders : public PPCallbacks {
 public:
   RecordHeaders(const SourceManager &SM, IncludeStructure *Out)
       : SM(SM), Out(Out) {}
 
   // Record existing #includes - both written and resolved paths. Only #includes
   // in the main file are collected.
   void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,

clang-tools-extra/clangd/Hover.cpp

llvm::Optional<std::string> printExprValue(const SelectionTree::Node *N,
   }
   return llvm::None;
 }
 
 llvm::Optional<StringRef> fieldName(const Expr *E) {
   const auto *ME = llvm::dyn_cast<MemberExpr>(E->IgnoreCasts());
   if (!ME || !llvm::isa<CXXThisExpr>(ME->getBase()->IgnoreCasts()))
     return llvm::None;
-  const auto *Field =
-      llvm::dyn_cast<FieldDecl>(ME->getMemberDecl());
+  const auto *Field = llvm::dyn_cast<FieldDecl>(ME->getMemberDecl());
   if (!Field || !Field->getDeclName().isIdentifier())
     return llvm::None;
   return Field->getDeclName().getAsIdentifierInfo()->getName();
 }
 
 // If CMD is of the form T foo() { return FieldName; } then returns "FieldName".
 llvm::Optional<StringRef> getterVariableName(const CXXMethodDecl *CMD) {
   assert(CMD->hasBody());
HoverInfo getHoverContents(const DefinedMacro &Macro, ParsedAST &AST) {
   HI.Name = std::string(Macro.Name);
   HI.Kind = index::SymbolKind::Macro;
   // FIXME: Populate documentation
   // FIXME: Populate parameters
 
   // Try to get the full definition, not just the name
   SourceLocation StartLoc = Macro.Info->getDefinitionLoc();
   SourceLocation EndLoc = Macro.Info->getDefinitionEndLoc();
-  if (EndLoc.isValid()) {
+  // Ensure that EndLoc is a valid offset. For example it might come from
+  // preamble, and source file might've changed, in such a scenario EndLoc still
+  // stays valid, but getLocForEndOfToken will fail as it is no longer a valid
+  // offset.
+  // Note that this check is just to ensure there's text data inside the range.
+  // It will still succeed even when the data inside the range is irrelevant to
+  // macro definition.
+  if (SM.getPresumedLoc(EndLoc, /*UseLineDirectives=*/false).isValid()) {
     EndLoc = Lexer::getLocForEndOfToken(EndLoc, 0, SM, AST.getLangOpts());
     bool Invalid;
     StringRef Buffer = SM.getBufferData(SM.getFileID(StartLoc), &Invalid);
     if (!Invalid) {
       unsigned StartOffset = SM.getFileOffset(StartLoc);
       unsigned EndOffset = SM.getFileOffset(EndLoc);
       if (EndOffset <= Buffer.size() && StartOffset < EndOffset)
         HI.Definition =
llvm::Optional<HoverInfo> getHover(ParsedAST &AST, Position Pos,
   llvm::Optional<HoverInfo> HI;
   // Macros and deducedtype only works on identifiers and auto/decltype keywords
   // respectively. Therefore they are only trggered on whichever works for them,
   // similar to SelectionTree::create().
   for (const auto &Tok : TokensTouchingCursor) {
     if (Tok.kind() == tok::identifier) {
       // Prefer the identifier token as a fallback highlighting range.
       HighlightRange = Tok.range(SM).toCharRange(SM);
-      if (auto M = locateMacroAt(Tok, AST.getPreprocessor())) {
+      if (auto M = locateMacroAt(Tok, AST.getPreprocessor(), AST.getTokens())) {
         HI = getHoverContents(*M, AST);
         break;
       }
     } else if (Tok.kind() == tok::kw_auto || Tok.kind() == tok::kw_decltype) {
       if (auto Deduced = getDeducedType(AST.getASTContext(), Tok.location())) {
         HI = getHoverContents(*Deduced, AST.getASTContext(), Index);
         HighlightRange = Tok.range(SM).toCharRange(SM);
         break;
if (Contents.empty() || isWhitespace(Contents.front()) ||
     return llvm::None;
 
   // The close-quote is usually followed by whitespace or punctuation.
   llvm::StringRef Suffix = Line.substr(Next + 1);
   constexpr llvm::StringLiteral AfterEndChars = " \t)=.,;:";
   if (!Suffix.empty() && !AfterEndChars.contains(Suffix.front()))
     return llvm::None;
 
-  return Line.slice(Offset, Next+1);
+  return Line.slice(Offset, Next + 1);
 }
 
 void parseDocumentationLine(llvm::StringRef Line, markup::Paragraph &Out) {
   // Probably this is appendText(Line), but scan for something interesting.
   for (unsigned I = 0; I < Line.size(); ++I) {
     switch (Line[I]) {
-      case '`':
-        if (auto Range = getBacktickQuoteRange(Line, I)) {
-          Out.appendText(Line.substr(0, I));
-          Out.appendCode(Range->trim("`"), /*Preserve=*/true);
-          return parseDocumentationLine(Line.substr(I+Range->size()), Out);
-        }
-        break;
+        return parseDocumentationLine(Line.substr(I + Range->size()), Out);
     }
   }
   Out.appendText(Line).appendSpace();
 }
 
 void parseDocumentation(llvm::StringRef Input, markup::Document &Output) {
   std::vector<llvm::StringRef> ParagraphLines;
   auto FlushParagraph = [&] {

clang-tools-extra/clangd/Preamble.cpp

struct TextualPPDirective {
   std::string Text;
 
   bool operator==(const TextualPPDirective &RHS) const {
     return std::tie(DirectiveLine, Text) ==
            std::tie(RHS.DirectiveLine, RHS.Text);
   }
 };
 
+// Spells directive in \p DirectiveRange while prepending it with \p Prefix.
+// Returned string can be used with a #line directive on line \p DirectiveLine
+// to correctly align with \p DirectiveRange.
+std::string spellDirective(llvm::StringRef Prefix, SourceRange DirectiveRange,
+                           unsigned &DirectiveLine, const SourceManager &SM) {
+  std::string SpelledDirective;
+  llvm::raw_string_ostream OS(SpelledDirective);
+  OS << Prefix;
+
+  auto DecompLoc = SM.getDecomposedLoc(DirectiveRange.getBegin());
+  DirectiveLine = SM.getLineNumber(DecompLoc.first, DecompLoc.second);
+  auto SpaceBefore = SM.getColumnNumber(DecompLoc.first, DecompLoc.second) - 1;
+
+  // Pad with spaces before DirectiveRange to make sure it will be on right
+  // column when patched.
+  if (SpaceBefore < Prefix.size()) {
+    // Prefix was longer than the space we had. Pad from start of a new line.
+    OS << "\\\n" << std::string(SpaceBefore, ' ');
+    // Decrement because returned string has a line break before directive now.
+    --DirectiveLine;
+  } else {
+    // There is enough space for Prefix and space before directive, use it.
+    // We try to squeeze the Prefix into the same line whenever we can, as
+    // putting onto a separate line won't work at the beginning of the file.
+    OS << std::string(SpaceBefore - Prefix.size(), ' ');
+  }
+  OS << toSourceCode(SM, DirectiveRange);
+  return OS.str();
+}
+
 // Collects #define directives inside the main file.
 struct DirectiveCollector : public PPCallbacks {
   DirectiveCollector(const Preprocessor &PP,
                      std::vector<TextualPPDirective> &TextualDirectives)
       : LangOpts(PP.getLangOpts()), SM(PP.getSourceManager()),
         TextualDirectives(TextualDirectives) {}
 
   void FileChanged(SourceLocation Loc, FileChangeReason Reason,
                    SrcMgr::CharacteristicKind FileType,
                    FileID PrevFID) override {
     InMainFile = SM.isWrittenInMainFile(Loc);
   }
 
   void MacroDefined(const Token &MacroNameTok,
                     const MacroDirective *MD) override {
     if (!InMainFile)
       return;
     TextualDirectives.emplace_back();
     TextualPPDirective &TD = TextualDirectives.back();
 
-    auto DecompLoc = SM.getDecomposedLoc(MacroNameTok.getLocation());
-    TD.DirectiveLine = SM.getLineNumber(DecompLoc.first, DecompLoc.second);
-
     SourceRange DefRange(
         MD->getMacroInfo()->getDefinitionLoc(),
         Lexer::getLocForEndOfToken(MD->getMacroInfo()->getDefinitionEndLoc(), 0,
                                    SM, LangOpts));
-    llvm::raw_string_ostream OS(TD.Text);
-    OS << "#define " << toSourceCode(SM, DefRange);
+    TD.Text = spellDirective("#define ", DefRange, TD.DirectiveLine, SM);
   }
 
 private:
   bool InMainFile = true;
   const LangOptions &LangOpts;
   const SourceManager &SM;
   std::vector<TextualPPDirective> &TextualDirectives;
 };
if (DirectivesChanged) {
     //
     // If we've patched only the first directive, the macro definition would've
     // been wrong for the rest of the file, since patch is applied after the
     // baseline preamble.
     //
     // Note that we deliberately ignore conditional directives and undefs to
     // reduce complexity. The former might cause problems because scanning is
     // imprecise and might pick directives from disabled regions.
-    for (const auto &TD : ModifiedScan->TextualDirectives)
+    for (const auto &TD : ModifiedScan->TextualDirectives) {
+      Patch << "#line " << TD.DirectiveLine << '\n';
       Patch << TD.Text << '\n';
     }
+  }
   dlog("Created preamble patch: {0}", Patch.str());
   Patch.flush();
   return PP;
 }
 
 void PreamblePatch::apply(CompilerInvocation &CI) const {
   // No need to map an empty file.
   if (PatchContents.empty())

clang-tools-extra/clangd/SourceCode.h

 /// \p FullyQualifiedName should not contain anonymous namespaces.
 EligibleRegion getEligiblePoints(llvm::StringRef Code,
                                  llvm::StringRef FullyQualifiedName,
                                  const LangOptions &LangOpts);
 
 struct DefinedMacro {
   llvm::StringRef Name;
   const MacroInfo *Info;
+  /// Location of the identifier that names the macro.
+  /// Unlike Info->Location, this translates preamble-patch locations to
+  /// main-file locations.
+  SourceLocation NameLoc;
 };
 /// 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);
+                                           Preprocessor &PP,
+                                           const syntax::TokenBuffer &TB);
 
 /// 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

Result.LikelyIdentifier =
       tok::isAnyIdentifier(
           IdentifierTable(LangOpts).get(Result.Text).getTokenID());
   for (const auto &T : Touching)
     if (T.location() <= Result.Location)
       Result.PartOfSpelledToken = &T;
   return Result;
 }
 
+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) {
+                                           Preprocessor &PP,
+                                           const syntax::TokenBuffer &TB) {
   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);
-  if (auto *MI = MacroDef.getMacroInfo())
-    return DefinedMacro{IdentifierInfo->getName(), MI};
+  auto *MI = MacroDef.getMacroInfo();
+  if (!MI)
     return None;
+
+  SourceLocation NameLoc = MI->getDefinitionLoc();
+  // Macro definitions could be injected through preamble patch. These contain
+  // #line directives to hint their original location in main file.
+  auto DefFile = SM.getFileID(NameLoc);
+  auto IncludeLoc = SM.getIncludeLoc(DefFile);
+  // Preamble patch is included inside the builtin file.
+  if (IncludeLoc.isValid() && SM.isWrittenInBuiltinFile(IncludeLoc)) {
+    auto Presumed = SM.getPresumedLoc(NameLoc);
+    // Check that line directive is pointing at main file.
+    if (Presumed.isValid() && Presumed.getFileID().isInvalid() &&
+        isMainFile(Presumed.getFilename(), SM)) {
+      NameLoc = SM.translateLineCol(SM.getMainFileID(), Presumed.getLine(),
+                                    Presumed.getColumn());
+    }
+  }
+
+  return DefinedMacro{IdentifierInfo->getName(), MI, NameLoc};
 }
 
 llvm::Expected<std::string> Edit::apply() const {
   return tooling::applyAllReplacements(InitialCode, Replacements);
 }
 
 std::vector<TextEdit> Edit::asTextEdits() const {
   return replacementsToEdits(InitialCode, Replacements);

clang-tools-extra/clangd/XRefs.cpp

llvm::Optional<LocatedSymbol> locateFileReferent(const Position &Pos,
   return llvm::None;
 }
 
 // Macros are simple: there's no declaration/definition distinction.
 // As a consequence, there's no need to look them up in the index either.
 llvm::Optional<LocatedSymbol>
 locateMacroReferent(const syntax::Token &TouchedIdentifier, ParsedAST &AST,
                     llvm::StringRef MainFilePath) {
-  if (auto M = locateMacroAt(TouchedIdentifier, AST.getPreprocessor())) {
-    if (auto Loc = makeLocation(AST.getASTContext(),
-                                M->Info->getDefinitionLoc(), MainFilePath)) {
+  if (auto M = locateMacroAt(TouchedIdentifier, AST.getPreprocessor(),
+                             AST.getTokens())) {
+    if (auto Loc =
+            makeLocation(AST.getASTContext(), M->NameLoc, MainFilePath)) {
       LocatedSymbol Macro;
       Macro.Name = std::string(M->Name);
       Macro.PreferredDeclaration = *Loc;
       Macro.Definition = Loc;
       return Macro;
     }
   }
   return llvm::None;
ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   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());
+    Macro = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor(),
+                          AST.getTokens());
   }
 
   RefsRequest Req;
   if (Macro) {
     // Handle references to macro.
     if (auto MacroSID = getSymbolID(Macro->Name, Macro->Info, SM)) {
       // Collect macro references from main file.
       const auto &IDToRefs = AST.getMacros().MacroRefs;
for (const NamedDecl *D : getDeclAtPosition(AST, *CurLoc, Relations)) {
     Results.push_back(std::move(NewSymbol));
   }
 
   const auto *IdentifierAtCursor =
       syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens());
   if (!IdentifierAtCursor)
     return Results;
 
-  if (auto M = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor())) {
+  if (auto M = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor(),
+                             AST.getTokens())) {
     SymbolDetails NewMacro;
     NewMacro.name = std::string(M->Name);
     llvm::SmallString<32> USR;
     if (!index::generateUSRForMacro(NewMacro.name, M->Info->getDefinitionLoc(),
                                     SM, USR)) {
       NewMacro.USR = std::string(USR.str());
       NewMacro.ID = SymbolID(NewMacro.USR);
     }

clang-tools-extra/clangd/refactor/Rename.cpp

if (!Loc)
     return Loc.takeError();
   const syntax::Token *IdentifierToken =
       spelledIdentifierTouching(*Loc, AST.getTokens());
   // Renames should only triggered on identifiers.
   if (!IdentifierToken)
     return makeError(ReasonToReject::NoSymbolFound);
   // FIXME: Renaming macros is not supported yet, the macro-handling code should
   // be moved to rename tooling library.
-  if (locateMacroAt(*IdentifierToken, AST.getPreprocessor()))
+  if (locateMacroAt(*IdentifierToken, AST.getPreprocessor(), AST.getTokens()))
     return makeError(ReasonToReject::UnsupportedSymbol);
 
   auto DeclsUnderCursor = locateDeclAt(AST, IdentifierToken->location());
   if (DeclsUnderCursor.empty())
     return makeError(ReasonToReject::NoSymbolFound);
   if (DeclsUnderCursor.size() > 1)
     return makeError(ReasonToReject::AmbiguousSymbol);
 

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

for (int I = 1;; I++) {
       const auto ExpectedRefs = T.ranges(llvm::to_string(I));
       if (ExpectedRefs.empty())
         break;
 
       auto Loc = sourceLocationInMainFile(SM, ExpectedRefs.begin()->start);
       ASSERT_TRUE(bool(Loc));
       const auto *Id = syntax::spelledIdentifierTouching(*Loc, AST.getTokens());
       ASSERT_TRUE(Id);
-      auto Macro = locateMacroAt(*Id, PP);
+      auto Macro = locateMacroAt(*Id, PP, AST.getTokens());
       assert(Macro);
       auto SID = getSymbolID(Macro->Name, Macro->Info, SM);
 
       EXPECT_THAT(ExpectedRefs,
                   UnorderedElementsAreArray(ActualMacroRefs.MacroRefs[*SID]))
           << "Annotation=" << I << ", MacroName=" << Macro->Name
           << ", Test = " << Test;
     }

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;
 using testing::MatchesRegex;
 
 namespace clang {
 namespace clangd {
 namespace {
 
 MATCHER_P2(Distance, File, D, "") {
   return arg.first() == File && arg.second == D;
llvm::Optional<ParsedAST> createPatchedAST(llvm::StringRef Baseline,
 
   IgnoreDiagnostics Diags;
   auto TU = TestTU::withCode(Modified);
   auto CI = buildCompilerInvocation(TU.inputs(), Diags);
   if (!CI) {
     ADD_FAILURE() << "Failed to build compiler invocation";
     return llvm::None;
   }
-  return ParsedAST::build(testPath("main.cpp"), TU.inputs(), std::move(CI), {},
+  return ParsedAST::build(testPath(TU.Filename), TU.inputs(), std::move(CI), {},
                           BaselinePreamble);
 }
 
 std::string getPreamblePatch(llvm::StringRef Baseline,
                              llvm::StringRef Modified) {
   auto BaselinePreamble = TestTU::withCode(Baseline).preamble();
   if (!BaselinePreamble) {
     ADD_FAILURE() << "Failed to build baseline preamble";
struct {
     llvm::StringLiteral Contents;
     llvm::StringLiteral ExpectedPatch;
   } Cases[] = {
       {
           R"cpp(
         #define BAR
         [[BAR]])cpp",
           R"cpp(#line 0 ".*main.cpp"
-#define BAR
+#line 2
+#define         BAR
 )cpp",
       },
       // multiline macro
       {
           R"cpp(
         #define BAR \
 
         [[BAR]])cpp",
           R"cpp(#line 0 ".*main.cpp"
-#define BAR
+#line 2
+#define         BAR
 )cpp",
       },
       // multiline macro
       {
           R"cpp(
         #define \
                 BAR
         [[BAR]])cpp",
           R"cpp(#line 0 ".*main.cpp"
-#define BAR
+#line 3
+#define         BAR
 )cpp",
       },
   };
 
   for (const auto &Case : Cases) {
     SCOPED_TRACE(Case.Contents);
     Annotations Modified(Case.Contents);
     EXPECT_THAT(getPreamblePatch("", Modified.code()),
TEST(PreamblePatchTest, OrderingPreserved) {
   llvm::StringLiteral Baseline = "#define BAR(X) X";
   Annotations Modified(R"cpp(
     #define BAR(X, Y) X Y
     #define BAR(X) X
     [[BAR]](int y);
   )cpp");
 
   llvm::StringLiteral ExpectedPatch(R"cpp(#line 0 ".*main.cpp"
-#define BAR\(X, Y\) X Y
-#define BAR\(X\) X
+#line 2
+#define     BAR\(X, Y\) X Y
+#line 3
+#define     BAR\(X\) X
 )cpp");
   EXPECT_THAT(getPreamblePatch(Baseline, Modified.code()),
               MatchesRegex(ExpectedPatch.str()));
 
   auto AST = createPatchedAST(Baseline, Modified.code());
   ASSERT_TRUE(AST);
   EXPECT_THAT(AST->getDiagnostics(),
               Not(Contains(Field(&Diag::Range, Modified.range()))));
 }
+
+TEST(PreamblePatchTest, LocateMacroAtWorks) {
+  struct {
+    llvm::StringLiteral Baseline;
+    llvm::StringLiteral Modified;
+  } Cases[] = {
+      // Addition of new directive
+      {
+          "",
+          R"cpp(
+            #define $def^FOO
+            $use^FOO)cpp",
+      },
+      // Available inside preamble section
+      {
+          "",
+          R"cpp(
+            #define $def^FOO
+            #undef $use^FOO)cpp",
+      },
+      // Available after undef, as we don't patch those
+      {
+          "",
+          R"cpp(
+            #define $def^FOO
+            #undef FOO
+            $use^FOO)cpp",
+      },
+      // Identifier on a different line
+      {
+          "",
+          R"cpp(
+            #define \
+              $def^FOO
+            $use^FOO)cpp",
+      },
+      // In presence of comment tokens
+      {
+          "",
+          R"cpp(
+            #\
+              define /* FOO */\
+              /* FOO */ $def^FOO
+            $use^FOO)cpp",
+      },
+      // Moved around
+      {
+          "#define FOO",
+          R"cpp(
+            #define BAR
+            #define $def^FOO
+            $use^FOO)cpp",
+      },
+  };
+  for (const auto &Case : Cases) {
+    SCOPED_TRACE(Case.Modified);
+    llvm::Annotations Modified(Case.Modified);
+    auto AST = createPatchedAST(Case.Baseline, Modified.code());
+    ASSERT_TRUE(AST);
+
+    const auto &SM = AST->getSourceManager();
+    auto *MacroTok = AST->getTokens().spelledTokenAt(
+        SM.getComposedLoc(SM.getMainFileID(), Modified.point("use")));
+    ASSERT_TRUE(MacroTok);
+
+    auto FoundMacro =
+        locateMacroAt(*MacroTok, AST->getPreprocessor(), AST->getTokens());
+    ASSERT_TRUE(FoundMacro);
+    EXPECT_THAT(FoundMacro->Name, "FOO");
+
+    auto MacroLoc = FoundMacro->NameLoc;
+    EXPECT_EQ(SM.getFileID(MacroLoc), SM.getMainFileID());
+    EXPECT_EQ(SM.getFileOffset(MacroLoc), Modified.point("def"));
+  }
+}
+
+TEST(PreamblePatchTest, LocateMacroAtDeletion) {
+  {
+    // We don't patch deleted define directives, make sure we don't crash.
+    llvm::StringLiteral Baseline = "#define FOO";
+    llvm::Annotations Modified("^FOO");
+
+    auto AST = createPatchedAST(Baseline, Modified.code());
+    ASSERT_TRUE(AST);
+
+    const auto &SM = AST->getSourceManager();
+    auto *MacroTok = AST->getTokens().spelledTokenAt(
+        SM.getComposedLoc(SM.getMainFileID(), Modified.point()));
+    ASSERT_TRUE(MacroTok);
+
+    auto FoundMacro =
+        locateMacroAt(*MacroTok, AST->getPreprocessor(), AST->getTokens());
+    ASSERT_TRUE(FoundMacro);
+    EXPECT_THAT(FoundMacro->Name, "FOO");
+    auto HI =
+        getHover(*AST, offsetToPosition(Modified.code(), Modified.point()),
+                 format::getLLVMStyle(), nullptr);
+    ASSERT_TRUE(HI);
+    EXPECT_THAT(HI->Definition, testing::IsEmpty());
+  }
+
+  {
+    // Offset is valid, but underlying text is different.
+    llvm::StringLiteral Baseline = "#define FOO";
+    Annotations Modified(R"cpp(#define BAR
+    ^FOO")cpp");
+
+    auto AST = createPatchedAST(Baseline, Modified.code());
+    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

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

Annotations Code(R"cpp(
      int abc = MA^CRO;
    )cpp");
   TestTU TU = TestTU::withCode(Code.code());
   auto AST = TU.build();
   auto CurLoc = sourceLocationInMainFile(AST.getSourceManager(), Code.point());
   ASSERT_TRUE(bool(CurLoc));
   const auto *Id = syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens());
   ASSERT_TRUE(Id);
-  auto Result = locateMacroAt(*Id, AST.getPreprocessor());
+  auto Result = locateMacroAt(*Id, AST.getPreprocessor(), AST.getTokens());
   ASSERT_TRUE(Result);
   EXPECT_THAT(*Result, MacroName("MACRO"));
 }
 
 TEST(SourceCodeTests, WorksAtBeginOfFile) {
   Annotations Code("^MACRO");
   TestTU TU = TestTU::withCode(Code.code());
   TU.HeaderCode = "#define MACRO int x;";
   auto AST = TU.build();
   auto CurLoc = sourceLocationInMainFile(AST.getSourceManager(), Code.point());
   ASSERT_TRUE(bool(CurLoc));
   const auto *Id = syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens());
   ASSERT_TRUE(Id);
-  auto Result = locateMacroAt(*Id, AST.getPreprocessor());
+  auto Result = locateMacroAt(*Id, AST.getPreprocessor(), AST.getTokens());
   ASSERT_TRUE(Result);
   EXPECT_THAT(*Result, MacroName("MACRO"));
 }
 
 TEST(SourceCodeTests, IsInsideMainFile) {
   TestTU TU;
   TU.HeaderCode = R"cpp(
     #define DEFINE_CLASS(X) class X {};