[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

974

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

124

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.

128

nit: move out-param to end?

134

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?

135

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.

137

again this comment only applyes to one branch

140

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

150

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

clang-tools-extra/clangd/SourceCode.cpp

975

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

134

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.

137

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

974

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

154

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

157

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

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

 //===--- Headers.cpp - Include headers ---------------------------*- 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 "Headers.h"
 #include "Compiler.h"
+#include "Preamble.h"
 #include "SourceCode.h"
 #include "support/Logger.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #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,
                           llvm::StringRef FileName, bool IsAngled,
                           CharSourceRange /*FilenameRange*/,
                           const FileEntry *File, llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
                           const Module * /*Imported*/,
                           SrcMgr::CharacteristicKind FileKind) override {
     auto MainFID = SM.getMainFileID();
     // If an include is part of the preamble patch, translate #line directives.
-    if (InBuiltinFile) {
-      auto Presumed = SM.getPresumedLoc(HashLoc);
-      // Presumed locations will have an invalid file id when #line directive
-      // changes the filename.
-      if (Presumed.getFileID().isInvalid() &&
-          isMainFile(Presumed.getFilename(), SM)) {
-        // Now we'll hit the case below.
-        HashLoc = SM.translateLineCol(MainFID, Presumed.getLine(),
-                                      Presumed.getColumn());
-      }
-    }
+    if (InBuiltinFile)
+      HashLoc = translatePreamblePatchLocation(HashLoc, SM);
 
     // Record main-file inclusions (including those mapped from the preamble
     // patch).
     if (isInsideMainFile(HashLoc, SM)) {
       Out->MainFileIncludes.emplace_back();
       auto &Inc = Out->MainFileIncludes.back();
       Inc.Written =
           (IsAngled ? "<" + FileName + ">" : "\"" + FileName + "\"").str();

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 =
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.h

private:
   PreamblePatch() = default;
   std::string PatchContents;
   std::string PatchFileName;
   /// Includes that are present in both \p Baseline and \p Modified. Used for
   /// patching includes of baseline preamble.
   std::vector<Inclusion> PreambleIncludes;
 };
 
+/// Translates locations inside preamble patch to their main-file equivalent
+/// using presumed locations. Returns \p Loc if it isn't inside preamble patch.
+SourceLocation translatePreamblePatchLocation(SourceLocation Loc,
+                                              const SourceManager &SM);
+
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

 #include <string>
 #include <system_error>
 #include <utility>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
+constexpr llvm::StringLiteral PreamblePatchHeaderName = "__preamble_patch__.h";
 
 bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,
                              const tooling::CompileCommand &RHS) {
   // We don't check for Output, it should not matter to clangd.
   return LHS.Directory == RHS.Directory && LHS.Filename == RHS.Filename &&
          llvm::makeArrayRef(LHS.CommandLine).equals(RHS.CommandLine);
 }
 
struct TextualPPDirective {
   std::string Text;
 
   bool operator==(const TextualPPDirective &RHS) const {
     return std::tie(DirectiveLine, Text) ==
            std::tie(RHS.DirectiveLine, RHS.Text);
   }
 };
 
+// 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.
+std::string spellDirective(llvm::StringRef Prefix,
+                           CharSourceRange DirectiveRange,
+                           const LangOptions &LangOpts, const SourceManager &SM,
+                           unsigned &DirectiveLine) {
+  std::string SpelledDirective;
+  llvm::raw_string_ostream OS(SpelledDirective);
+  OS << Prefix;
+
+  // Make sure DirectiveRange is a char range and doesn't contain macro ids.
+  DirectiveRange = SM.getExpansionRange(DirectiveRange);
+  if (DirectiveRange.isTokenRange()) {
+    DirectiveRange.setEnd(
+        Lexer::getLocForEndOfToken(DirectiveRange.getEnd(), 0, SM, LangOpts));
+  }
+
+  auto DecompLoc = SM.getDecomposedLoc(DirectiveRange.getBegin());
+  DirectiveLine = SM.getLineNumber(DecompLoc.first, DecompLoc.second);
+  auto TargetColumn = SM.getColumnNumber(DecompLoc.first, DecompLoc.second) - 1;
+
+  // Pad with spaces before DirectiveRange to make sure it will be on right
+  // column when patched.
+  if (Prefix.size() <= TargetColumn) {
+    // 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(TargetColumn - Prefix.size(), ' ');
+  } else {
+    // Prefix was longer than the space we had. We produce e.g.:
+    // #line N-1
+    // #define \
+    //    X 10
+    OS << "\\\n" << std::string(TargetColumn, ' ');
+    // Decrement because we put an additional line break before
+    // DirectiveRange.begin().
+    --DirectiveLine;
+  }
+  OS << toSourceCode(SM, DirectiveRange.getAsRange());
+  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);
+    const auto *MI = MD->getMacroInfo();
+    TD.Text =
+        spellDirective("#define ",
+                       CharSourceRange::getTokenRange(
+                           MI->getDefinitionLoc(), MI->getDefinitionEndLoc()),
+                       LangOpts, SM, TD.DirectiveLine);
   }
 
 private:
   bool InMainFile = true;
   const LangOptions &LangOpts;
   const SourceManager &SM;
   std::vector<TextualPPDirective> &TextualDirectives;
 };
case tok::pp_import:
     return "import";
   case tok::pp_include_next:
     return "include_next";
   default:
     break;
   }
   llvm_unreachable("not an include directive");
 }
+
+// Checks whether \p FileName is a valid spelling of main file.
+bool isMainFile(llvm::StringRef FileName, const SourceManager &SM) {
+  auto FE = SM.getFileManager().getFile(FileName);
+  return FE && *FE == SM.getFileEntryForID(SM.getMainFileID());
+}
+
 } // namespace
 
 PreambleData::PreambleData(const ParseInputs &Inputs,
                            PrecompiledPreamble Preamble,
                            std::vector<Diag> Diags, IncludeStructure Includes,
                            MainFileMacros Macros,
                            std::unique_ptr<PreambleFileStatusCache> StatCache,
                            CanonicalIncludes CanonIncludes)
bool DirectivesChanged =
       BaselineScan->TextualDirectives != ModifiedScan->TextualDirectives;
   if (!IncludesChanged && !DirectivesChanged)
     return PreamblePatch::unmodified(Baseline);
 
   PreamblePatch PP;
   // This shouldn't coincide with any real file name.
   llvm::SmallString<128> PatchName;
   llvm::sys::path::append(PatchName, llvm::sys::path::parent_path(FileName),
-                          "__preamble_patch__.h");
+                          PreamblePatchHeaderName);
   PP.PatchFileName = PatchName.str().str();
 
   llvm::raw_string_ostream Patch(PP.PatchContents);
   // Set default filename for subsequent #line directives
   Patch << "#line 0 \"";
   // FileName part of a line directive is subject to backslash escaping, which
   // might lead to problems on windows especially.
   escapeBackslashAndQuotes(FileName, Patch);
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())

 }
 
 PreamblePatch PreamblePatch::unmodified(const PreambleData &Preamble) {
   PreamblePatch PP;
   PP.PreambleIncludes = Preamble.Includes.MainFileIncludes;
   return PP;
 }
 
+SourceLocation translatePreamblePatchLocation(SourceLocation Loc,
+                                              const SourceManager &SM) {
+  auto DefFile = SM.getFileID(Loc);
+  if (auto *FE = SM.getFileEntryForID(DefFile)) {
+    auto IncludeLoc = SM.getIncludeLoc(DefFile);
+    // Preamble patch is included inside the builtin file.
+    if (IncludeLoc.isValid() && SM.isWrittenInBuiltinFile(IncludeLoc) &&
+        FE->getName().endswith(PreamblePatchHeaderName)) {
+      auto Presumed = SM.getPresumedLoc(Loc);
+      // Check that line directive is pointing at main file.
+      if (Presumed.isValid() && Presumed.getFileID().isInvalid() &&
+          isMainFile(Presumed.getFilename(), SM)) {
+        Loc = SM.translateLineCol(SM.getMainFileID(), Presumed.getLine(),
+                                  Presumed.getColumn());
+      }
+    }
+  }
+  return Loc;
+}
 } // namespace clangd
 } // namespace clang

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);
 
 /// Infers whether this is a header from the FileName and LangOpts (if
 /// presents).

clang-tools-extra/clangd/SourceCode.cpp

 //===--- SourceCode.h - Manipulating source code as strings -----*- 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 "SourceCode.h"
 
 #include "FuzzyMatch.h"
+#include "Preamble.h"
 #include "Protocol.h"
 #include "refactor/Tweak.h"
 #include "support/Context.h"
 #include "support/Logger.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
llvm::Optional<DefinedMacro> locateMacroAt(const syntax::Token &SpelledTok,
 
   // 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};
+    return DefinedMacro{
+        IdentifierInfo->getName(), MI,
+        translatePreamblePatchLocation(MI->getDefinitionLoc(), SM)};
   return None;
 }
 
 llvm::Expected<std::string> Edit::apply() const {
   return tooling::applyAllReplacements(InitialCode, Replacements);
 }
 
 std::vector<TextEdit> Edit::asTextEdits() const {
   return replacementsToEdits(InitialCode, Replacements);
 }
 
 bool Edit::canApplyTo(llvm::StringRef Code) const {
   // Create line iterators, since line numbers are important while applying our
   // edit we cannot skip blank lines.
   auto LHS = llvm::MemoryBuffer::getMemBuffer(Code);
   llvm::line_iterator LHSIt(*LHS, /*SkipBlanks=*/false);

clang-tools-extra/clangd/XRefs.cpp

 }
 
 // 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 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;

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

EXPECT_EQ(Inserter.calculateIncludePath(Verbatim, MainFile),
             std::string("<x>"));
   EXPECT_EQ(Inserter.shouldInsertInclude(HeaderPath, Verbatim), true);
 
   EXPECT_EQ(Inserter.calculateIncludePath(Inserting, "sub2/main2.cpp"),
             llvm::None);
 }
 
 TEST_F(HeadersTest, PresumedLocations) {
-  std::string HeaderFile = "implicit_include.h";
+  std::string HeaderFile = "__preamble_patch__.h";
 
   // Line map inclusion back to main file.
   std::string HeaderContents =
       llvm::formatv("#line 0 \"{0}\"", llvm::sys::path::filename(MainFile));
   HeaderContents += R"cpp(
 #line 3
 #include <a.h>)cpp";
   FS.Files[HeaderFile] = HeaderContents;
 
   // Including through non-builtin file has no effects.
-  FS.Files[MainFile] = "#include \"implicit_include.h\"\n\n";
+  FS.Files[MainFile] = "#include \"__preamble_patch__.h\"\n\n";
   EXPECT_THAT(collectIncludes().MainFileIncludes,
               Not(Contains(Written("<a.h>"))));
 
   // Now include through built-in file.
   CDB.ExtraClangFlags = {"-include", testPath(HeaderFile)};
   EXPECT_THAT(collectIncludes().MainFileIncludes,
               Contains(AllOf(IncludeLine(2), Written("<a.h>"))));
 }
 } // namespace
 } // namespace clangd
 } // namespace clang

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());
+    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());
+    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));
+    }
+  }
+}
+
+TEST(TranslatePreamblePatchLocation, Simple) {
+  auto TU = TestTU::withHeaderCode(R"cpp(
+    #line 3 "main.cpp"
+    int foo();)cpp");
+  // Presumed line/col needs to be valid in the main file.
+  TU.Code = R"cpp(// line 1
+    // line 2
+    // line 3
+    // line 4)cpp";
+  TU.Filename = "main.cpp";
+  TU.HeaderFilename = "__preamble_patch__.h";
+  TU.ImplicitHeaderGuard = false;
+
+  auto AST = TU.build();
+  auto &SM = AST.getSourceManager();
+  auto &ND = findDecl(AST, "foo");
+  EXPECT_NE(SM.getFileID(ND.getLocation()), SM.getMainFileID());
+
+  auto TranslatedLoc = translatePreamblePatchLocation(ND.getLocation(), SM);
+  auto DecompLoc = SM.getDecomposedLoc(TranslatedLoc);
+  EXPECT_EQ(DecompLoc.first, SM.getMainFileID());
+  EXPECT_EQ(SM.getLineNumber(DecompLoc.first, DecompLoc.second), 3U);
+}
 } // namespace
 } // namespace clangd
 } // namespace clang