[clangd] Patch PragmaMarks in preamble section of the file

[hokein]

clang-tools-extra/clangd/Preamble.cpp

579

any reason not calling collectPragmaMarksCallback just like the https://reviews.llvm.org/D146028 for macros? is it for simplicity?

The pragma mark might be simple enough to parse ourselves, but it doesn't handle well for some trivial/corner cases like (well we probably don't care about these cases)

#pragma mark \
   abc

[kadircet]

• Use collectPragmaMarksCallback

[kadircet]

clang-tools-extra/clangd/Preamble.cpp

579

yeah it was for not running it unnecessarily, but you're right I guess it's not worth the discrepancy. using collectPragmaMarksCallback instead.

clang-tools-extra/clangd/Preamble.h

private:
   std::string PatchFileName;
   // Includes that are present in both Baseline and Modified. Used for
   // patching includes of baseline preamble.
   std::vector<Inclusion> PreambleIncludes;
   // Diags that were attached to a line preserved in Modified contents.
   std::vector<Diag> PatchedDiags;
   PreambleBounds ModifiedBounds = {0, false};
   const PreambleData *Baseline = nullptr;
+  std::vector<PragmaMark> PatchedMarks;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

 #include "CollectMacros.h"
 #include "Compiler.h"
 #include "Config.h"
 #include "Diagnostics.h"
 #include "Headers.h"
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "clang-include-cleaner/Record.h"
+#include "index/CanonicalIncludes.h"
 #include "support/Logger.h"
+#include "support/Path.h"
 #include "support/ThreadsafeFS.h"
 #include "support/Trace.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticLex.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
+#include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/PPCallbacks.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/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.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 <cstddef>
+#include <functional>
 #include <memory>
 #include <optional>
 #include <string>
 #include <system_error>
+#include <tuple>
 #include <utility>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,

 };
 
 struct ScannedPreamble {
   std::vector<Inclusion> Includes;
   std::vector<TextualPPDirective> TextualDirectives;
   // Literal lines of the preamble contents.
   std::vector<llvm::StringRef> Lines;
   PreambleBounds Bounds = {0, false};
+  std::vector<PragmaMark> Marks;
 };
 
 /// Scans the preprocessor directives in the preamble section of the file by
 /// running preprocessor over \p Contents. Returned includes do not contain
 /// resolved paths. \p Cmd is used to build the compiler invocation, which might
 /// stat/read files.
 llvm::Expected<ScannedPreamble>
 scanPreamble(llvm::StringRef Contents, const tooling::CompileCommand &Cmd) {
if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
     return error("failed BeginSourceFile");
   Preprocessor &PP = Clang->getPreprocessor();
   IncludeStructure Includes;
   Includes.collect(*Clang);
   ScannedPreamble SP;
   SP.Bounds = Bounds;
   PP.addPPCallbacks(
       std::make_unique<DirectiveCollector>(PP, SP.TextualDirectives));
+  PP.addPPCallbacks(
+      collectPragmaMarksCallback(PP.getSourceManager(), SP.Marks));
   if (llvm::Error Err = Action.Execute())
     return std::move(Err);
   Action.EndSourceFile();
   SP.Includes = std::move(Includes.MainFileIncludes);
   llvm::append_range(SP.Lines, llvm::split(Contents, "\n"));
   return SP;
 }
 
std::optional<Diag> translateDiag(const Diag &D) {
     // Translate ranges inside notes and fixes too, dropping the ones that are
     // no longer relevant.
     llvm::erase_if(NewD.Notes, [this](Note &N) { return !translateNote(N); });
     llvm::erase_if(NewD.Fixes, [this](Fix &F) { return !translateFix(F); });
     return NewD;
   }
 };
 } // 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 =
for (const auto &TD : ModifiedScan->TextualDirectives) {
       if (TD.Directive == tok::pp_define)
         Patch << "#undef " << TD.MacroName << '\n';
       Patch << "#line " << TD.DirectiveLine << '\n';
       Patch << TD.Text << '\n';
     }
   }
 
   PP.PatchedDiags = patchDiags(Baseline.Diags, *BaselineScan, *ModifiedScan);
+  PP.PatchedMarks = std::move(ModifiedScan->Marks);
   dlog("Created preamble patch: {0}", Patch.str());
   Patch.flush();
   return PP;
 }
 
 PreamblePatch PreamblePatch::createFullPatch(llvm::StringRef FileName,
                                              const ParseInputs &Modified,
                                              const PreambleData &Baseline) {

 
 bool PreamblePatch::preserveDiagnostics() const {
   return PatchContents.empty() ||
          Config::current().Diagnostics.AllowStalePreamble;
 }
 llvm::ArrayRef<PragmaMark> PreamblePatch::marks() const {
   if (PatchContents.empty())
     return Baseline->Marks;
-  // FIXME: Patch pragma marks.
-  return {};
+  return PatchedMarks;
 }
 
 MainFileMacros PreamblePatch::mainFileMacros() const {
   if (PatchContents.empty())
     return Baseline->Macros;
   // FIXME: Patch main file macros.
   return MainFileMacros();
 }
 } // namespace clangd
 } // namespace clang

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

#define $main[[BAR]] 2)");
     Annotations NewCode(R"(
 #define BAZ 0
 #define BAR 1)");
     auto AST = createPatchedAST(Code.code(), NewCode.code());
     EXPECT_THAT(*AST->getDiagnostics(), IsEmpty());
   }
 }
 
+MATCHER_P2(Mark, Range, Text, "") {
+  return std::tie(arg.Rng, arg.Trivia) == std::tie(Range, Text);
+}
+
 TEST(PreamblePatch, MacroAndMarkHandling) {
   Config Cfg;
   Cfg.Diagnostics.AllowStalePreamble = true;
   WithContextValue WithCfg(Config::Key, std::move(Cfg));
 
   {
     Annotations Code(R"cpp(
 #ifndef FOO
 #define FOO
 // Some comments
 #pragma mark XX
 #define BAR
 
 #endif)cpp");
     Annotations NewCode(R"cpp(
 #ifndef FOO
 #define FOO
 #define BAR
-#pragma mark XX
+#pragma $x[[mark XX
+]]
+#pragma $y[[mark YY
+]]
 
 #endif)cpp");
     auto AST = createPatchedAST(Code.code(), NewCode.code());
     // FIXME: Macros and marks have locations that need to be patched.
     EXPECT_THAT(AST->getMacros().Names, IsEmpty());
-    EXPECT_THAT(AST->getMarks(), IsEmpty());
+    EXPECT_THAT(AST->getMarks(),
+                UnorderedElementsAre(Mark(NewCode.range("x"), " XX"),
+                                     Mark(NewCode.range("y"), " YY")));
   }
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang