[clangd] Patch main file macros in preamble

[kadircet]

Summary

Depends on D146026
Fixes https://github.com/clangd/clangd/issues/1537.

[hokein]

clang-tools-extra/clangd/Preamble.cpp

921

nit: maybe just return a const MainFileMacros&.

[kadircet]

• Rebase
• Prevent unnecessary copies

clang-tools-extra/clangd/Preamble.h

public:
   bool preserveDiagnostics() const;
 
   /// Returns diag locations for Modified contents.
   llvm::ArrayRef<Diag> patchedDiags() const { return PatchedDiags; }
 
   static constexpr llvm::StringLiteral HeaderName = "__preamble_patch__.h";
 
   llvm::ArrayRef<PragmaMark> marks() const;
-  MainFileMacros mainFileMacros() const;
+  const MainFileMacros &mainFileMacros() const;
 
 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 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;
+  MainFileMacros PatchedMacros;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

 
 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;
+  MainFileMacros Macros;
 };
 
 /// 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 (Clang->getFrontendOpts().Inputs.empty())
     return error("compiler instance had no inputs");
   // We are only interested in main file includes.
   Clang->getPreprocessorOpts().SingleFileParseMode = true;
   Clang->getPreprocessorOpts().UsePredefines = false;
   PreprocessOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
     return error("failed BeginSourceFile");
   Preprocessor &PP = Clang->getPreprocessor();
+  const auto &SM = PP.getSourceManager();
   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));
+  PP.addPPCallbacks(collectPragmaMarksCallback(SM, SP.Marks));
+  PP.addPPCallbacks(std::make_unique<CollectMainFileMacros>(SM, SP.Macros));
   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;
 }
 
for (const auto &TD : ModifiedScan->TextualDirectives) {
         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);
+  PP.PatchedMacros = std::move(ModifiedScan->Macros);
   dlog("Created preamble patch: {0}", Patch.str());
   Patch.flush();
   return PP;
 }
 
 PreamblePatch PreamblePatch::createFullPatch(llvm::StringRef FileName,
                                              const ParseInputs &Modified,
                                              const PreambleData &Baseline) {
return PatchContents.empty() ||
          Config::current().Diagnostics.AllowStalePreamble;
 }
 llvm::ArrayRef<PragmaMark> PreamblePatch::marks() const {
   if (PatchContents.empty())
     return Baseline->Marks;
   return PatchedMarks;
 }
 
-MainFileMacros PreamblePatch::mainFileMacros() const {
+const MainFileMacros &PreamblePatch::mainFileMacros() const {
   if (PatchContents.empty())
     return Baseline->Macros;
-  // FIXME: Patch main file macros.
-  return MainFileMacros();
+  return PatchedMacros;
 }
 } // namespace clangd
 } // namespace clang

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

 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/VirtualFileSystem.h"
-#include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest-matchers.h"
 #include "gtest/gtest.h"
 #include <memory>
 #include <optional>
 #include <string>
 #include <utility>
 #include <vector>
#endif)cpp");
     Annotations NewCode(R"cpp(
 #ifndef FOO
 #define FOO
 #define BAR
 #pragma $x[[mark XX
 ]]
 #pragma $y[[mark YY
 ]]
+#define BAZ
 
 #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->getMacros().Names.keys(),
+                UnorderedElementsAreArray({"FOO", "BAR", "BAZ"}));
     EXPECT_THAT(AST->getMarks(),
                 UnorderedElementsAre(Mark(NewCode.range("x"), " XX"),
                                      Mark(NewCode.range("y"), " YY")));
   }
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang