[clangd] Provide patched diagnostics with preamble patch

[kadircet]

Summary

Translates diagnostics from baseline preamble to relevant modified
contents.

Translation is done by looking for a set of lines that have the same
contents in diagnostic/note/fix ranges inside baseline and modified
contents.

A diagnostic is preserved if its main range is outside of main file or
there's a translation from baseline to modified contents. Later on fixes
and notes attached to that diagnostic with relevant ranges are also
translated and preserved.

Depends on D143095

[sammccall]

It looks like this fixes up the location only of diagnostics attached to particular directives (#include) based on some "deep" idea about the content of the directive (the spelled header name).

Some shortcomings:

• This misses diagnostics attached to other directives / continued lines in macro definitions / etc
• it allows diagnostics to be translated across lines even if the directive spelling changed (in which case the ranges are incorrect since only line numbers are updated.)
• it requires modelling the directive content in some way that needs to be extended if we find another place that diagnostics can be attached

We discussed the idea of attaching to the text of the line, which seems pretty generic. Needs some handling of duplicate line content but seems like something pretty naive (closest line number with matching content?) would work well, be just as simple and more generic. Any reason this alternative was rejected?

clang-tools-extra/clangd/Preamble.h

162

this looks like a gratuitous copy - ParsedAST copies again

[kadircet]

In D143096#4099662, @sammccall wrote:

It looks like this fixes up the location only of diagnostics attached to particular directives (#include) based on some "deep" idea about the content of the directive (the spelled header name).

Some shortcomings:

• This misses diagnostics attached to other directives / continued lines in macro definitions / etc
• it allows diagnostics to be translated across lines even if the directive spelling changed (in which case the ranges are incorrect since only line numbers are updated.)
• it requires modelling the directive content in some way that needs to be extended if we find another place that diagnostics can be attached

We discussed the idea of attaching to the text of the line, which seems pretty generic. Needs some handling of duplicate line content but seems like something pretty naive (closest line number with matching content?) would work well, be just as simple and more generic. Any reason this alternative was rejected?

agreed this is better, also somewhat more elegant on the implementation side. switching to a line-content based translation, while keeping the logic around limiting diagnostics to single lines. as multiple lines seem more wonky and unclear how common they're.

[kadircet]

• Change to a line based translation logic

[sammccall]

clang-tools-extra/clangd/ParsedAST.cpp

678

llvm::append_range(Diags, Patch->patchedDiags()) ?

clang-tools-extra/clangd/Preamble.cpp

369

llvm::append_range(SP.Lines, llvm::split(Contents, "\n")) ?

631

nit: comment sense is the opposite of the function name

maybe say why this is an important property, or give the function a higher-level name?

632

this function + translateDiag look like a validate/parse pair that need to be kept in sync.

they could be a single translateDiag() function that could fail, instead.

637

this check looks dubious to me:

• File is a human-readable string with ill-defined format, comparing them for equality isn't safe, you want InsideMainFile instead
• if the note points to another file, then it's not that translation is impossible, but rather no translation is needed
• if the note points to the same file, then why can't we apply the same translation?

I think this should rather be something like:

bool translateDiag(DiagBase& D, bool IsMainDiag) {
  if (IsMainDiag) {
    for (const auto& N : cast<Diag>(D).Notes) {
      if (!translateDiag(N, /*IsMainDiag=*/false))
        return false;
    }
  }
  if (D.InsideMainFile) {
    // attempt patching range of D itself
  }
}

672

nit: name seems backwards (but as noted below I think this function can go away entirely)

697

If I'm understanding the algorithm right:

• we make a map of line content => original line number
• we iterate through each line of the modified preamble, and:
  • find the corresponding original line
  • look for diagnostics to translate

Why not loop over diagnostics instead of lines in the modified preamble? There should be a lot fewer, and you avoid the lines-to-diags map.

699

large preamble with no diagnostics seems pretty plausible, maybe bail out early to avoid building the content map?

762–763

I think this change and the IncludeKey class can be reverted now?

clang-tools-extra/clangd/Preamble.h

160

the stuff about #include/#define doesn't seem to correspond to anything in the code (obsolete?)
and below

174–175

if you want to undoxygen this comment, also remove \p?

[sammccall]

clang-tools-extra/clangd/Preamble.cpp

640

Oh yeah, one more thing :-)

Currently if the note/fix aren't at this line we're bailing out of the diagnostic entirely.

We have three behaviors available:

1. patch the note/fix range as best we can and emit it
2. drop the note/fix but keep the diagnostic
3. drop the diagnostic entirely

I think we should generally prefer 1 when we can, maybe there's a case for choosing 2 sometimes (heuristics about when diagnostics have probably been invalidated, with comments), and we should ~never do 3.

[kadircet]

• Change patching logic to iterate over diags instead of preamble lines
• Change translation logic to:
  • Preserve a diagnostic if its range can be translated.
  • Preserve all the notes whose range can be translated.
  • Preserve all the fixes whose edit ranges can be translated.
• Drop restrictions on ranges being a single line, make sure contents for all the lines are matched between modified and basline contents.
• Add couple new tests to demonstrate what's preserved.

[kadircet]

• Reflow comments

[kadircet]

• Update commit message

[sammccall]

clang-tools-extra/clangd/Preamble.cpp

633

I think this function is too long with too many local lambdas, consider converting it to a class instead (call me old-fashioned!)

653

we're assuming that the ranges are within the preamble and everyone agrees about the bounds. If not, BaselineStart may not be a valid index into BaselineScan.Lines

This assumption seems broadly reasonable, but is *very much* not locally enforced. I'd suggest a defensive check or at least an assert.

656

this doesn't look right: you're first deciding which possible starting point is closest, and then deciding whether it matches. So a range that matches can be masked by a range where only the first line matches, if the latter is closer.

691

probably doesn't matter, but you're inconsistent about moving vs copying range

714

reasoning about the fields you're *not* rewriting feels fragile. (Here and in TranslateFix).

Consider copying the whole object and mutating in place (bool TranslateDiag(Diag&) together with erase_if)

clang-tools-extra/clangd/Preamble.h

161

nit: blank line after

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

765

i'm confused about what this comment is getting at - a note without a diag is not even representable

[kadircet]

• Move patch translation logic to a separate class
• Perform eager copies and use bool translateX(X&) type of APIs instead of returning optionals.
• Perform multi line content check for each alternative, not just the closest match.

clang-tools-extra/clangd/Preamble.cpp

633

well it's gonna be an equally long class, but sure :D

653

oops, thanks! both here and also the multi-ine check below.

656

i thought because this would be rare it's fine to not do that, but the same applies in the other direction as well :D
moved the content based matching inside the loop to consider a line as alternative only after it matches the contents.

714

well i am a little worried about copying around notes/fixes just to drop them again (which can have quite some strings). but probably rare enough in practice to not matter.

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

765

well, it is to make sure we're not promoting an orphaned note to a main diag. updated the comment but happy to drop if you think it's confusing rather than being helpful.

[sammccall]

LG thanks!

clang-tools-extra/clangd/Preamble.cpp

490

up to you, but I find avoiding iterators easier to read

for (int AlternateLine : CurrentContentsToLine.lookup(OldLines[OldStart]))

(or with extracted variables)

503

why not just if (RangeContents != CurrentLines.slice(...))? I feel like I'm missing something subtle :-)

Also if you replace the slice(a, b) with drop_front(a).take_front(b) then you don't need the CurrentEnd > CurrentLines.size() check

511

if you make Closest the delta rather than the line number this gets a bit easier to follow IMO:
abs(Closest) < abs(Delta) here
R.start.line += Closest; R.end.line += Closest; below

You need to make Closest an optional<int> but that's also clearer than a sentinel -1 for me

562

FWIW if we're willing to pay a copy for a successful patch, I don't think it's worth any hassle to avoid the copy for a failed patch. Up to you though.

[kadircet]

• Use direct comparison of ArrayRefs
• optional for closest
• loop over lookup results rather than through iterator

[kadircet]

thanks for the review!

clang-tools-extra/clangd/Preamble.cpp

503

why not just if (RangeContents != CurrentLines.slice(...))? I feel like I'm missing something subtle :-)

no that's absolutely right, i guess i had a brain fart 😂

[uabelho]

Hello,

I noticed that if I compile the compiler with ubsan, then lots of Clangd tests start failing with this patch:

Failed Tests (165):
  Clangd :: ast-no-range.test
  Clangd :: ast.test
  Clangd :: call-hierarchy.test
  Clangd :: check-lines.test
  Clangd :: check.test
  Clangd :: code-action-request.test
  Clangd :: completion-auto-trigger.test
  Clangd :: completion-snippets.test
  Clangd :: completion.test
  Clangd :: config.test
  Clangd :: crash-parse.test
  Clangd :: dependency-output.test
  Clangd :: diagnostic-category.test
  Clangd :: diagnostics-no-tidy.test
  Clangd :: diagnostics-notes.test
  Clangd :: diagnostics-tidy.test
  Clangd :: did-change-configuration-params.test
  Clangd :: execute-command.test
  Clangd :: filestatus.test
  Clangd :: fixits-codeaction.test
  Clangd :: fixits-command.test
  Clangd :: fixits-embed-in-diagnostic.test
  Clangd :: folding-range.test
  Clangd :: formatting.test
  Clangd :: hover.test
  Clangd :: implementations.test
  Clangd :: inlayHints.test
  Clangd :: memory_tree.test
  Clangd :: metrics.test
  Clangd :: protocol.test
  Clangd :: references-container.test
  Clangd :: references.test
  Clangd :: rename.test
  Clangd :: request-reply.test
  Clangd :: selection-range.test
  Clangd :: semantic-tokens-refresh.test
  Clangd :: semantic-tokens.test
  Clangd :: signature-help-with-offsets.test
  Clangd :: signature-help.test
  Clangd :: symbol-info.test
  Clangd :: symbols.test
  Clangd :: target_info.test
  Clangd :: test-uri-posix.test
  Clangd :: textdocument-didchange-fail.test
  Clangd :: trace.test
  Clangd :: tweaks-format.test
  Clangd :: type-definition.test
  Clangd :: type-hierarchy-ext.test
  Clangd :: type-hierarchy.test
  Clangd :: unsupported-method.test
  Clangd :: utf8.test
  Clangd :: version.test
  Clangd :: xrefs.test
  Clangd Unit Tests :: ./ClangdTests/0/146
  [...]
  Clangd Unit Tests :: ./ClangdTests/99/146


Testing Time: 362.21s
  Skipped          :    47
  Unsupported      :  1094
  Passed           : 89208
  Expectedly Failed:   192
  Failed           :   165

Not sure if all fail the same way but there are a lot of ubsan complaints like this

07:38:09 ../lib/Support/MemoryBuffer.cpp:138:33: runtime error: null pointer passed as argument 2, which is declared to never be null
07:38:09 /usr/include/string.h:43:28: note: nonnull attribute specified here
07:38:09 SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior ../lib/Support/MemoryBuffer.cpp:138:33 in

[kadircet]

thanks for the ping! see the discussion in https://reviews.llvm.org/D142890#4149679

clang-tools-extra/clangd/ParsedAST.cpp

ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   ASTDiags.EndSourceFile();
 
   std::optional<std::vector<Diag>> Diags;
   // FIXME: Also skip generation of diagnostics alltogether to speed up ast
   // builds when we are patching a stale preamble.
   if (PreserveDiags) {
     Diags = CompilerInvocationDiags;
     // Add diagnostics from the preamble, if any.
-    if (Preamble)
-      Diags->insert(Diags->end(), Preamble->Diags.begin(),
-                    Preamble->Diags.end());
+    if (Preamble) {
+      auto PDiags = Patch->patchedDiags();
+      Diags->insert(Diags->end(), PDiags.begin(), PDiags.end());
+    }
     // Finally, add diagnostics coming from the AST.
     {
       std::vector<Diag> D = ASTDiags.take(&*CTContext);
       Diags->insert(Diags->end(), D.begin(), D.end());
     }
   }
   ParsedAST Result(Filename, Inputs.Version, std::move(Preamble),
                    std::move(Clang), std::move(Action), std::move(Tokens),

clang-tools-extra/clangd/Preamble.h

public:
   PreambleBounds modifiedBounds() const { return ModifiedBounds; }
 
   /// Returns textual patch contents.
   llvm::StringRef text() const { return PatchContents; }
 
   /// Whether diagnostics generated using this patch are trustable.
   bool preserveDiagnostics() const;
 
+  /// Returns diag locations for Modified contents, only contains diags attached
+  /// to an #include or #define directive.
+  std::vector<Diag> patchedDiags() const { return PatchedDiags; }
   static constexpr llvm::StringLiteral HeaderName = "__preamble_patch__.h";
 
 private:
   static PreamblePatch create(llvm::StringRef FileName,
                               const ParseInputs &Modified,
                               const PreambleData &Baseline,
                               PatchType PatchType);
 
   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.
+  // Includes that are present in both \p Baseline and \p Modified. Used for
+  // patching includes of baseline preamble.
   std::vector<Inclusion> PreambleIncludes;
+  // Diags that were attached to an #include or a #define directive.
+  std::vector<Diag> PatchedDiags;
   PreambleBounds ModifiedBounds = {0, false};
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

 //===--- Preamble.cpp - Reusing expensive parts of the AST ----------------===//
 //
 // 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 "Preamble.h"
 #include "Compiler.h"
 #include "Config.h"
+#include "Diagnostics.h"
 #include "Headers.h"
 #include "SourceCode.h"
 #include "clang-include-cleaner/Record.h"
 #include "support/Logger.h"
 #include "support/ThreadsafeFS.h"
 #include "support/Trace.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/Diagnostic.h"

 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/raw_ostream.h"
-#include <iterator>
+#include <map>
 #include <memory>
 #include <string>
 #include <system_error>
 #include <utility>
 #include <vector>
 
 namespace clang {
 namespace clangd {
scanPreamble(llvm::StringRef Contents, const tooling::CompileCommand &Cmd) {
   ScannedPreamble SP;
   SP.Bounds = Bounds;
   PP.addPPCallbacks(
       std::make_unique<DirectiveCollector>(PP, SP.TextualDirectives));
   if (llvm::Error Err = Action.Execute())
     return std::move(Err);
   Action.EndSourceFile();
   SP.Includes = std::move(Includes.MainFileIncludes);
   return SP;
 }
 
 const char *spellingForIncDirective(tok::PPKeywordKind IncludeDirective) {
   switch (IncludeDirective) {
   case tok::pp_include:
     return "include";
   case tok::pp_import:
     return "import";
public:
   }
 
   double getTime() { return Timer.getTime(); }
 
 private:
   WallTimer Timer;
 };
 
+// Represents an include as spelled in the source code, without taking resolved
+// path or location into account.
+struct IncludeKey {
+  tok::PPKeywordKind Directive;
+  llvm::StringRef Written;
+
+  bool operator<(const IncludeKey &RHS) const {
+    return std::tie(Directive, Written) < std::tie(RHS.Directive, RHS.Written);
+  }
+  bool operator==(const IncludeKey &RHS) const {
+    return std::tie(Directive, Written) == std::tie(RHS.Directive, RHS.Written);
+  }
+};
 } // namespace
 
 std::shared_ptr<const PreambleData>
 buildPreamble(PathRef FileName, CompilerInvocation CI,
               const ParseInputs &Inputs, bool StoreInMemory,
               PreambleParsedCallback PreambleCallback,
               PreambleBuildStats *Stats) {
   // Note that we don't need to copy the input contents, preamble can live
   // without those.
   auto ContentsBuffer =
       llvm::MemoryBuffer::getMemBuffer(Inputs.Contents, FileName);
   auto Bounds = ComputePreambleBounds(*CI.getLangOpts(), *ContentsBuffer, 0);
 
   trace::Span Tracer("BuildPreamble");
   SPAN_ATTACH(Tracer, "File", FileName);
   std::vector<std::unique_ptr<FeatureModule::ASTListener>> ASTListeners;
   if (Inputs.FeatureModules) {
     for (auto &M : *Inputs.FeatureModules) {
       if (auto Listener = M.astListeners())
         ASTListeners.emplace_back(std::move(Listener));
     }
   }
   StoreDiags PreambleDiagnostics;
   PreambleDiagnostics.setDiagCallback(
       [&ASTListeners](const clang::Diagnostic &D, clangd::Diag &Diag) {
         llvm::for_each(ASTListeners,
                        [&](const auto &L) { L->sawDiagnostic(D, Diag); });
       });
   llvm::IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
       CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(),
                                           &PreambleDiagnostics, false);
   const Config &Cfg = Config::current();
   PreambleDiagnostics.setLevelAdjuster([&](DiagnosticsEngine::Level DiagLevel,
                                            const clang::Diagnostic &Info) {
     if (Cfg.Diagnostics.SuppressAll ||
         isBuiltinDiagnosticSuppressed(Info.getID(), Cfg.Diagnostics.Suppress,
buildPreamble(PathRef FileName, CompilerInvocation CI,
   auto StatCacheFS = StatCache->getProducingFS(VFS);
   llvm::IntrusiveRefCntPtr<TimerFS> TimedFS(new TimerFS(StatCacheFS));
 
   WallTimer PreambleTimer;
   PreambleTimer.startTimer();
   auto BuiltPreamble = PrecompiledPreamble::Build(
       CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine,
       Stats ? TimedFS : StatCacheFS, std::make_shared<PCHContainerOperations>(),
       StoreInMemory, CapturedInfo);
   PreambleTimer.stopTimer();
 
   // When building the AST for the main file, we do want the function
   // bodies.
   CI.getFrontendOpts().SkipFunctionBodies = false;
 
   if (Stats != nullptr) {
     Stats->TotalBuildTime = PreambleTimer.getTime();
case '"':
       break;
     default:
       break;
     }
     OS << C;
   }
 }
 
+// Checks if all pointers in \p D are for the same line of the main file.
+static bool diagReferencesMultipleLines(const Diag &D) {
+  int Line = D.Range.start.line;
+  if (D.Range.end.line != Line)
+    return true;
+  bool NotePointsToOutside = llvm::any_of(D.Notes, [&](const Note &N) {
+    return N.File == D.File &&
+           (N.Range.start.line != Line || N.Range.end.line != Line);
+  });
+  if (NotePointsToOutside)
+    return true;
+  bool FixPointsToOutside = llvm::any_of(D.Fixes, [&](const Fix &F) {
+    for (auto &E : F.Edits)
+      if (E.range.start.line != Line || E.range.end.line != Line)
+        return true;
+    return false;
+  });
+  if (FixPointsToOutside)
+    return true;
+  return false;
+}
+
+// Move all references in \p D to \p NewLine. This assumes none of the
+// references to the main file are multi-line.
+static Diag translateDiag(const Diag &D, int NewLine) {
+  assert(!diagReferencesMultipleLines(D));
+  Diag Translated = D;
+  Translated.Range.start.line = Translated.Range.end.line = NewLine;
+  for (auto &N : Translated.Notes) {
+    if (N.File != D.File)
+      continue;
+    N.Range.start.line = N.Range.end.line = NewLine;
+  }
+  for (auto &F : Translated.Fixes) {
+    for (auto &E : F.Edits)
+      E.range.start.line = E.range.end.line = NewLine;
+  }
+  return Translated;
+}
+
+static llvm::DenseMap<int, llvm::SmallVector<const Diag *>>
+mapDiagsToLines(llvm::ArrayRef<Diag> Diags) {
+  llvm::DenseMap<int, llvm::SmallVector<const Diag *>> LineToDiags;
+  for (auto &D : Diags) {
+    if (diagReferencesMultipleLines(D))
+      continue;
+    LineToDiags[D.Range.start.line].emplace_back(&D);
+  }
+  return LineToDiags;
+}
+
+static std::map<IncludeKey, const Inclusion *>
+getExistingIncludes(const ScannedPreamble &BaselineScan,
+                    llvm::ArrayRef<Inclusion> MFI) {
+  std::map<IncludeKey, const Inclusion *> ExistingIncludes;
+  // There might be includes coming from disabled regions, record these for
+  // exclusion. note that we don't have resolved paths for those.
+  for (const auto &Inc : BaselineScan.Includes)
+    ExistingIncludes[{Inc.Directive, Inc.Written}] = nullptr;
+  for (const auto &Inc : MFI)
+    ExistingIncludes[{Inc.Directive, Inc.Written}] = &Inc;
+  return ExistingIncludes;
+}
+
 PreamblePatch PreamblePatch::create(llvm::StringRef FileName,
                                     const ParseInputs &Modified,
                                     const PreambleData &Baseline,
                                     PatchType PatchType) {
   trace::Span Tracer("CreatePreamblePatch");
   SPAN_ATTACH(Tracer, "File", FileName);
   assert(llvm::sys::path::is_absolute(FileName) && "relative FileName!");
   // First scan preprocessor directives in Baseline and Modified. These will be
   // used to figure out newly added directives in Modified. Scanning can fail,
   // the code just bails out and creates an empty patch in such cases, as:
   // - If scanning for Baseline fails, no knowledge of existing includes hence
   //   patch will contain all the includes in Modified. Leading to rebuild of
   //   whole preamble, which is terribly slow.
   // - If scanning for Modified fails, cannot figure out newly added ones so
   //   there's nothing to do but generate an empty patch.
   auto BaselineScan = scanPreamble(
       // Contents needs to be null-terminated.
       Baseline.Preamble.getContents().str(), Modified.CompileCommand);
   if (!BaselineScan) {
     elog("Failed to scan baseline of {0}: {1}", FileName,
          BaselineScan.takeError());
     return PreamblePatch::unmodified(Baseline);
   }
   auto ModifiedScan = scanPreamble(Modified.Contents, Modified.CompileCommand);
   if (!ModifiedScan) {
     elog("Failed to scan modified contents of {0}: {1}", FileName,
          ModifiedScan.takeError());
     return PreamblePatch::unmodified(Baseline);
PreamblePatch PreamblePatch::create(llvm::StringRef FileName,
   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);
   Patch << "\"\n";
 
+  // List of diagnostics associated to a particular line in the baseline.
+  auto LineToDiags = mapDiagsToLines(Baseline.Diags);
+  auto MoveDiagsBetweenLines = [&LineToDiags, &PP](int OldLine, int NewLine) {
+    auto DiagsForInclude = LineToDiags.find(OldLine);
+    if (DiagsForInclude == LineToDiags.end())
+      return;
+    for (auto *D : DiagsForInclude->second)
+      PP.PatchedDiags.emplace_back(translateDiag(*D, NewLine));
+  };
+  // Map from an include to its line in the Modified contents.
+  std::map<IncludeKey, int> IncludeToPatchedLine;
   if (IncludesChanged && PatchType == PatchType::All) {
     // We are only interested in newly added includes, record the ones in
     // Baseline for exclusion.
-    llvm::DenseMap<std::pair<tok::PPKeywordKind, llvm::StringRef>,
-                   /*Resolved=*/llvm::StringRef>
-        ExistingIncludes;
-    for (const auto &Inc : Baseline.Includes.MainFileIncludes)
-      ExistingIncludes[{Inc.Directive, Inc.Written}] = Inc.Resolved;
-    // There might be includes coming from disabled regions, record these for
-    // exclusion too. note that we don't have resolved paths for those.
-    for (const auto &Inc : BaselineScan->Includes)
-      ExistingIncludes.try_emplace({Inc.Directive, Inc.Written});
+    auto ExistingIncludes =
+        getExistingIncludes(*BaselineScan, Baseline.Includes.MainFileIncludes);
     // Calculate extra includes that needs to be inserted.
     for (auto &Inc : ModifiedScan->Includes) {
       auto It = ExistingIncludes.find({Inc.Directive, Inc.Written});
       // Include already present in the baseline preamble. Set resolved path and
       // put into preamble includes.
       if (It != ExistingIncludes.end()) {
-        Inc.Resolved = It->second.str();
+        if (It->second) {
+          // Copy everything from existing include, apart from the location.
+          Inc.Resolved = It->second->Resolved;
+          Inc.HeaderID = It->second->HeaderID;
+          Inc.BehindPragmaKeep = It->second->BehindPragmaKeep;
+          Inc.FileKind = It->second->FileKind;
+        }
         PP.PreambleIncludes.push_back(Inc);
+        IncludeToPatchedLine[It->first] = Inc.HashLine;
         continue;
       }
       // Include is new in the modified preamble. Inject it into the patch and
       // use #line to set the presumed location to where it is spelled.
       auto LineCol = offsetToClangLineColumn(Modified.Contents, Inc.HashOffset);
       Patch << llvm::formatv("#line {0}\n", LineCol.first);
       Patch << llvm::formatv(
           "#{0} {1}\n", spellingForIncDirective(Inc.Directive), Inc.Written);
     }
+  } else {
+    // Copy include information from baseline preamble if nothing has changed.
+    PP.PreambleIncludes = Baseline.Includes.MainFileIncludes;
+  }
+
+  // Patch locations for diagnostics on baseline includes.
+  for (const auto &Inc : Baseline.Includes.MainFileIncludes) {
+    // Use existing locations if includes haven't changed.
+    size_t NewLine = Inc.HashLine;
+    if (IncludesChanged) {
+      auto PatchedLine =
+          IncludeToPatchedLine.find({Inc.Directive, Inc.Written});
+      // Ignore diags for deleted includes.
+      if (PatchedLine == IncludeToPatchedLine.end())
+        continue;
+      NewLine = PatchedLine->second;
+    }
+    MoveDiagsBetweenLines(Inc.HashLine, NewLine);
   }
 
+  // Maps a directive from its spelling to its location in the Modified
+  // contents.
+  llvm::StringMap<int> DirectiveToPatchedLine;
   if (DirectivesChanged) {
     // We need to patch all the directives, since they are order dependent. e.g:
     // #define BAR(X) NEW(X) // Newly introduced in Modified
     // #define BAR(X) OLD(X) // Exists in the Baseline
     //
     // 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) {
       // Introduce an #undef directive before #defines to suppress any
       // re-definition warnings.
       if (TD.Directive == tok::pp_define)
         Patch << "#undef " << TD.MacroName << '\n';
       Patch << "#line " << TD.DirectiveLine << '\n';
       Patch << TD.Text << '\n';
+      DirectiveToPatchedLine[TD.Text] = TD.DirectiveLine;
     }
+  } else {
+    // Take existing directives as-is, if they're the same.
+    for (const auto &TD : BaselineScan->TextualDirectives)
+      DirectiveToPatchedLine[TD.Text] = TD.DirectiveLine;
   }
+
+  // Patch locations for diagnositcs on baseline macros.
+  for (const auto &TD : BaselineScan->TextualDirectives) {
+    // Textual directive lines are 1-based.
+    auto NewLine = DirectiveToPatchedLine.find(TD.Text);
+    if (NewLine == DirectiveToPatchedLine.end())
+      continue;
+    MoveDiagsBetweenLines(TD.DirectiveLine - 1, NewLine->second - 1);
+  }
+
   dlog("Created preamble patch: {0}", Patch.str());
   Patch.flush();
   return PP;
 }
 
 PreamblePatch PreamblePatch::createFullPatch(llvm::StringRef FileName,
                                              const ParseInputs &Modified,
                                              const PreambleData &Baseline) {

 std::vector<Inclusion> PreamblePatch::preambleIncludes() const {
   return PreambleIncludes;
 }
 
 PreamblePatch PreamblePatch::unmodified(const PreambleData &Preamble) {
   PreamblePatch PP;
   PP.PreambleIncludes = Preamble.Includes.MainFileIncludes;
   PP.ModifiedBounds = Preamble.Preamble.getBounds();
+  PP.PatchedDiags = Preamble.Diags;
   return PP;
 }
 
 bool PreamblePatch::preserveDiagnostics() const {
   return PatchContents.empty() ||
          Config::current().Diagnostics.Mode != Config::DiagnosticsMode::Strict;
 }
 } // namespace clangd
 } // namespace clang

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

TEST(PreamblePatch, DiagnosticsToPreamble) {
   AdditionalFiles["foo.h"] = "#pragma once";
   AdditionalFiles["barx.h"] = "#pragma once";
   llvm::StringLiteral BaselinePreamble("#define FOO 1\n[[#include \"foo.h\"]]");
   {
     // Check with removals from preamble.
     Annotations NewCode("[[#  include \"foo.h\"]]");
     auto AST = createPatchedAST(Annotations(BaselinePreamble).code(),
                                 NewCode.code(), AdditionalFiles);
-    // FIXME: Should be pointing at the range inside the Newcode.
-    EXPECT_THAT(mainFileDiagRanges(*AST), IsEmpty());
+    EXPECT_THAT(mainFileDiagRanges(*AST),
+                UnorderedElementsAreArray(NewCode.ranges()));
   }
   {
     // Check with additions to preamble.
     Annotations Code(BaselinePreamble);
     Annotations NewCode(("[[#include \"barx.h\"]]\n" + BaselinePreamble).str());
     auto AST = createPatchedAST(Code.code(), NewCode.code(), AdditionalFiles);
-    // FIXME: Should preserve existing diagnostics.
     EXPECT_THAT(mainFileDiagRanges(*AST),
-                UnorderedElementsAre(NewCode.ranges()[0]));
+                UnorderedElementsAreArray(NewCode.ranges()));
   }
   {
     llvm::StringLiteral BaselinePreamble = "#define [[FOO]] 1\n";
     Annotations Code(BaselinePreamble);
     // Check ranges for notes.
     Annotations NewCode(("#define BARXYZ 1\n" + BaselinePreamble +
                          "void foo();\n#define [[FOO]] 2")
                             .str());
TEST(PreamblePatch, TranslatesIncludeDiagnosticsInPreamble) {
   WithContextValue WithCfg(Config::Key, std::move(Cfg));
 
   llvm::StringLiteral BaselinePreamble("#include [[<foo>]]\n");
   {
     // Check with additions to preamble.
     Annotations Code(BaselinePreamble);
     Annotations NewCode(("#define BAR\n" + BaselinePreamble).str());
     auto AST = createPatchedAST(Code.code(), NewCode.code());
-    // FIXME: We should point at the NewCode.
-    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAre(Code.range()));
+    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAreArray(NewCode.ranges()));
   }
   {
     // Check with removals from preamble.
     Annotations Code(("#define BAR\n" + BaselinePreamble).str());
     Annotations NewCode(BaselinePreamble);
     auto AST = createPatchedAST(Code.code(), NewCode.code());
-    // FIXME: We should point at the NewCode.
-    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAre(Code.range()));
+    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAreArray(NewCode.ranges()));
   }
   {
     // Drop line with diags.
     Annotations Code(BaselinePreamble);
     Annotations NewCode("#define BAR\n#include [[<bar>]]\n");
     auto AST = createPatchedAST(Code.code(), NewCode.code());
-    // FIXME: We should only have ranges from the NewCode.
-    EXPECT_THAT(mainFileDiagRanges(*AST),
-                ElementsAreArray({Code.range(), NewCode.range()}));
+    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAreArray(NewCode.ranges()));
   }
 }
 
 TEST(PreamblePatch, TranslatesNonIncludeDiagnosticsInPreamble) {
   Config Cfg;
   Cfg.Diagnostics.Mode = Config::DiagnosticsMode::Fast;
   WithContextValue WithCfg(Config::Key, std::move(Cfg));
 
   llvm::StringLiteral BaselinePreamble("#define [[__LINE__]]\n");
   ASSERT_THAT(TestTU::withCode(Annotations(BaselinePreamble).code())
                   .build()
                   .getDiagnostics(),
               Not(llvm::ValueIs(IsEmpty())));
   {
     // Check with additions to preamble.
     Annotations Code(BaselinePreamble);
     Annotations NewCode(("#define BAR\n" + BaselinePreamble).str());
     auto AST = createPatchedAST(Code.code(), NewCode.code());
-    // FIXME: We should point at the NewCode.
-    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAre(Code.range()));
+    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAreArray(NewCode.ranges()));
   }
   {
     // Check with removals from preamble.
     Annotations Code(("#define BAR\n" + BaselinePreamble).str());
     Annotations NewCode(BaselinePreamble);
     auto AST = createPatchedAST(Code.code(), NewCode.code());
-    // FIXME: We should point at the NewCode.
-    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAre(Code.range()));
+    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAreArray(NewCode.ranges()));
   }
   {
     // Drop line with diags.
     Annotations Code(BaselinePreamble);
     Annotations NewCode("#define BAR\n#include [[<bar>]]\n#include <baz>\n");
     auto AST = createPatchedAST(Code.code(), NewCode.code());
-    // FIXME: We should only have the range from NewCode.
-    EXPECT_THAT(mainFileDiagRanges(*AST),
-                ElementsAreArray({Code.range(), NewCode.range()}));
+    EXPECT_THAT(mainFileDiagRanges(*AST), ElementsAreArray(NewCode.ranges()));
   }
 }
 } // namespace
 } // namespace clangd
 } // namespace clang