[clangd] Make use of preamble bounds from the patch inside ReplayPreamble

[kadircet]

Summary

Clangd was using bounds from the stale preamble, which might result in
crashes. For example:

#include "a.h"
#include "b.h" // this line is newly inserted
#include "c.h"

PreambleBounds for the baseline only contains first two lines, but
ReplayPreamble logic contains an include from the third line. This would
result in a crash as we only lex preamble part of the current file
during ReplayPreamble.

This patch adds a preambleBounds method to PreamblePatch, which can be
used to figure out preamble bounds for the current version of the file.
Then uses it when attaching ReplayPreamble, so that it can lex the
up-to-date preamble region.

[sammccall]

clang-tools-extra/clangd/Preamble.cpp

220

nit: = {} again

288

(we know these are trivial structs, could drop move)

clang-tools-extra/clangd/Preamble.h

123

current -> Modified

I'd suggest calling this modifiedBounds()

136

nit: = {} seems less confusing.

Call this ModifiedBounds (or ModifiedPreambleBounds) as well? Not obvious to me what it is at the moment.

[kadircet]

• Address comments

clang-tools-extra/clangd/Preamble.h

136

unfortunately PreambleBounds doesn't have a default constructor :/

clang-tools-extra/clangd/ParsedAST.cpp

ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
 
   IncludeStructure Includes;
   // If we are using a preamble, copy existing includes.
   if (Preamble) {
     Includes = Preamble->Includes;
     Includes.MainFileIncludes = Patch->preambleIncludes();
     // Replay the preamble includes so that clang-tidy checks can see them.
     ReplayPreamble::attach(Patch->preambleIncludes(), *Clang,
-                           Preamble->Preamble.getBounds());
+                           Patch->modifiedBounds());
   }
   // Important: collectIncludeStructure is registered *after* ReplayPreamble!
   // Otherwise we would collect the replayed includes again...
   // (We can't *just* use the replayed includes, they don't have Resolved path).
   Clang->getPreprocessor().addPPCallbacks(
       collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
   // Copy over the macros in the preamble region of the main file, and combine
   // with non-preamble macros below.

clang-tools-extra/clangd/Preamble.h

 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "FS.h"
 #include "Headers.h"
 #include "index/CanonicalIncludes.h"
 #include "support/Path.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
+#include "clang/Lex/Lexer.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringRef.h"
 
 #include <memory>
 #include <string>
 #include <vector>
 
 namespace clang {
public:
   /// Returns #include directives from the \c Modified preamble that were
   /// resolved using the \c Baseline preamble. This covers the new locations of
   /// inclusions that were moved around, but not inclusions of new files. Those
   /// will be recorded when parsing the main file: the includes in the injected
   /// section will be resolved back to their spelled positions in the main file
   /// using the presumed-location mechanism.
   std::vector<Inclusion> preambleIncludes() const;
 
+  /// Returns preamble bounds for the Modified.
+  PreambleBounds modifiedBounds() const { return ModifiedBounds; }
+
   /// Returns textual patch contents.
   llvm::StringRef text() const { return PatchContents; }
 
 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;
+  PreambleBounds ModifiedBounds = {0, false};
 };
 
 /// 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

private:
   const LangOptions &LangOpts;
   const SourceManager &SM;
   std::vector<TextualPPDirective> &TextualDirectives;
 };
 
 struct ScannedPreamble {
   std::vector<Inclusion> Includes;
   std::vector<TextualPPDirective> TextualDirectives;
+  PreambleBounds Bounds = {0, false};
 };
 
 /// 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 VFS and \p Cmd is used to build the compiler invocation,
 /// which might stat/read files.
 llvm::Expected<ScannedPreamble>
 scanPreamble(llvm::StringRef Contents,
scanPreamble(llvm::StringRef Contents,
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "failed BeginSourceFile");
   const auto &SM = Clang->getSourceManager();
   Preprocessor &PP = Clang->getPreprocessor();
   IncludeStructure Includes;
   PP.addPPCallbacks(collectIncludeStructureCallback(SM, &Includes));
   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;
 }
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),
                           PreamblePatchHeaderName);
   PP.PatchFileName = PatchName.str().str();
+  PP.ModifiedBounds = ModifiedScan->Bounds;
 
   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";

 
 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();
   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);

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

for (size_t I = 0; I < Includes.size(); ++I) {
     // This also contains quotes/angles so increment the range by one from both
     // sides.
     EXPECT_EQ(
         SkippedFiles[I].text(SM),
         Code.substr(FileRange.Begin - 1, FileRange.End - FileRange.Begin + 2));
     EXPECT_EQ(SkippedFiles[I].kind(), tok::header_name);
   }
 
+  TU.AdditionalFiles["a.h"] = "";
+  TU.AdditionalFiles["b.h"] = "";
+  TU.AdditionalFiles["c.h"] = "";
   // Make sure replay logic works with patched preambles.
-  TU.Code = "";
-  StoreDiags Diags;
+  llvm::StringLiteral Baseline = R"cpp(
+    #include "a.h"
+    #include "c.h")cpp";
   MockFSProvider FS;
+  TU.Code = Baseline.str();
   auto Inputs = TU.inputs(FS);
-  auto CI = buildCompilerInvocation(Inputs, Diags);
-  auto EmptyPreamble =
-      buildPreamble(testPath(TU.Filename), *CI, Inputs, true, nullptr);
-  ASSERT_TRUE(EmptyPreamble);
-  TU.Code = "#include <a.h>";
+  auto BaselinePreamble = TU.preamble();
+  ASSERT_TRUE(BaselinePreamble);
+
+  // First make sure we don't crash on various modifications to the preamble.
+  llvm::StringLiteral Cases[] = {
+      // clang-format off
+      // New include in middle.
+      R"cpp(
+        #include "a.h"
+        #include "b.h"
+        #include "c.h")cpp",
+      // New include at top.
+      R"cpp(
+        #include "b.h"
+        #include "a.h"
+        #include "c.h")cpp",
+      // New include at bottom.
+      R"cpp(
+        #include "a.h"
+        #include "c.h"
+        #include "b.h")cpp",
+      // Same size with a missing include.
+      R"cpp(
+        #include "a.h"
+        #include "b.h")cpp",
+      // Smaller with no new includes.
+      R"cpp(
+        #include "a.h")cpp",
+      // Smaller with a new includes.
+      R"cpp(
+        #include "b.h")cpp",
+      // clang-format on
+  };
+  for (llvm::StringLiteral Case : Cases) {
+    TU.Code = Case.str();
+
+    IgnoreDiagnostics Diags;
+    auto CI = buildCompilerInvocation(TU.inputs(FS), Diags);
+    auto PatchedAST = ParsedAST::build(testPath(TU.Filename), TU.inputs(FS),
+                                       std::move(CI), {}, BaselinePreamble);
+    ASSERT_TRUE(PatchedAST);
+    EXPECT_TRUE(PatchedAST->getDiagnostics().empty());
+  }
+
+  // Then ensure correctness by making sure includes were seen only once.
+  // Note that we first see the includes from the patch, as preamble includes
+  // are replayed after exiting the built-in file.
   Includes.clear();
+  TU.Code = R"cpp(
+    #include "a.h"
+    #include "b.h")cpp";
+  IgnoreDiagnostics Diags;
+  auto CI = buildCompilerInvocation(TU.inputs(FS), Diags);
   auto PatchedAST = ParsedAST::build(testPath(TU.Filename), TU.inputs(FS),
-                                     std::move(CI), {}, EmptyPreamble);
+                                     std::move(CI), {}, BaselinePreamble);
   ASSERT_TRUE(PatchedAST);
-  // Make sure includes were seen only once.
+  EXPECT_TRUE(PatchedAST->getDiagnostics().empty());
   EXPECT_THAT(Includes,
               ElementsAre(WithFileName(testPath("__preamble_patch__.h")),
-                          WithFileName("a.h")));
+                          WithFileName("b.h"), WithFileName("a.h")));
 }
 
 TEST(ParsedASTTest, PatchesAdditionalIncludes) {
   llvm::StringLiteral ModifiedContents = R"cpp(
     #include "baz.h"
     #include "foo.h"
     #include "sub/aux.h"
     void bar() {

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

TEST(TranslatePreamblePatchLocation, Simple) {
   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);
 }
+
+TEST(PreamblePatch, ModifiedBounds) {
+  struct {
+    llvm::StringLiteral Baseline;
+    llvm::StringLiteral Modified;
+  } Cases[] = {
+      // Size increased
+      {
+          "",
+          R"cpp(
+            #define FOO
+            FOO)cpp",
+      },
+      // Stayed same
+      {"#define FOO", "#define BAR"},
+      // Got smaller
+      {
+          R"cpp(
+            #define FOO
+            #undef FOO)cpp",
+          "#define FOO"},
+  };
+
+  for (const auto &Case : Cases) {
+    auto TU = TestTU::withCode(Case.Baseline);
+    auto BaselinePreamble = TU.preamble();
+    ASSERT_TRUE(BaselinePreamble);
+
+    Annotations Modified(Case.Modified);
+    TU.Code = Modified.code().str();
+    MockFSProvider FSProvider;
+    auto PP = PreamblePatch::create(testPath(TU.Filename),
+                                    TU.inputs(FSProvider), *BaselinePreamble);
+
+    IgnoreDiagnostics Diags;
+    auto CI = buildCompilerInvocation(TU.inputs(FSProvider), Diags);
+    ASSERT_TRUE(CI);
+
+    const auto ExpectedBounds =
+        Lexer::ComputePreamble(Case.Modified, *CI->getLangOpts());
+    EXPECT_EQ(PP.modifiedBounds().Size, ExpectedBounds.Size);
+    EXPECT_EQ(PP.modifiedBounds().PreambleEndsAtStartOfLine,
+              ExpectedBounds.PreambleEndsAtStartOfLine);
+  }
+}
 } // namespace
 } // namespace clangd
 } // namespace clang