[clangd] Fix clangd crash when including a header

[qdelacru]

Summary

Fixes https://github.com/clangd/clangd/issues/819

SourceLocation of macros change when a header file is included above it. This is not checked when creating a PreamblePatch, resulting in reusing previously built preamble with an incorrect source location for the macro in the example test case.
This patch stores the SourceLocation in the struct TextualPPDirective so that it gets checked when comparing old vs new preambles.

Also creates a preamble patch for code completion parsing so that clangd does not crash when following the example test case with a large file.

[kadircet]

Thanks, I can see the problem now (sorry for the late reply, i was on leave last week).

It is amazing that this hasn't bitten us yet during code complete flow (it probably did, but clangd would recover after restart so probably people didn't notice).
I still don't think that we'd like to parse all the new includes during code completion flow hence we should just create a preamble patch for the macro directives that moved around. (I don't think we should make the parser more resilient instead).
Do you mind adding a new factory function to PreamblePatch for creating a macro-directive only patch that can be used in CodeComplete flow to prevent crashes (while paying for some extra latency :/ )?

As for storing source locations, i think we can just store the offsets (we already have them in spellDirective when we decompose sourcelocation.

Does that make sense?

clang-tools-extra/clangd/Preamble.cpp

512

this is the condition i am suggesting to drop.

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

554

we actually have a more problematic case, you can just insert whitespace to the first line rather than the include directive. now all the offsets are wrong but there's nothing to patch.

[qdelacru]

Create macro directive only preamble patch for code completion

[kadircet]

clang-tools-extra/clangd/Preamble.cpp

148–149

nit: i'd put offset before text, since it's faster to compare.

534

sorry for not being explicit about this on the first round but rather than duplicating the logic, can we have a shared private factory function that takes an extra parameter to control whether includes/macros should be included in the patch? It can probably look something like this:

class PreamblePatch {
public:
  enum class PatchType {
     MacroDirectives,
     All,
  };
  // these two will just call the create with relevant patchtype
  static PreamblePatch createFullPatch(FileName, Modified, Baseline);
  static PreamblePatch createMacroPatch(FileName, Modified, Baseline);
private:
  static PreamblePatch create(filename, Modified, Baseline, PatchType);
}

[kadircet]

Thanks a lot, LGTM!

Do you want me to land this for you?

[qdelacru]

Yes please land this. Thanks!

clang-tools-extra/clangd/CodeComplete.cpp

CodeCompleteResult codeComplete(PathRef FileName, Position Pos,
   }
   auto Flow = CodeCompleteFlow(
       FileName, Preamble ? Preamble->Includes : IncludeStructure(),
       SpecFuzzyFind, Opts);
   return (!Preamble || Opts.RunParser == CodeCompleteOptions::NeverParse)
              ? std::move(Flow).runWithoutSema(ParseInput.Contents, *Offset,
                                               *ParseInput.TFS)
              : std::move(Flow).run({FileName, *Offset, *Preamble,
-                                    // We want to serve code completions with
-                                    // low latency, so don't bother patching.
-                                    /*PreamblePatch=*/llvm::None, ParseInput});
+                                    /*PreamblePatch=*/
+                                    PreamblePatch::createMacroPatch(
+                                        FileName, ParseInput, *Preamble),
+                                    ParseInput});
 }
 
 SignatureHelp signatureHelp(PathRef FileName, Position Pos,
                             const PreambleData &Preamble,
                             const ParseInputs &ParseInput) {
   auto Offset = positionToOffset(ParseInput.Contents, Pos);
   if (!Offset) {
     elog("Signature help position was invalid {0}", Offset.takeError());
     return SignatureHelp();
   }
   SignatureHelp Result;
   clang::CodeCompleteOptions Options;
   Options.IncludeGlobals = false;
   Options.IncludeMacros = false;
   Options.IncludeCodePatterns = false;
   Options.IncludeBriefComments = false;
   semaCodeComplete(
       std::make_unique<SignatureHelpCollector>(Options, ParseInput.Index,
                                                Result),
       Options,
       {FileName, *Offset, Preamble,
-       PreamblePatch::create(FileName, ParseInput, Preamble), ParseInput});
+       PreamblePatch::createFullPatch(FileName, ParseInput, Preamble),
+       ParseInput});
   return Result;
 }
 
 bool isIndexedForCodeCompletion(const NamedDecl &ND, ASTContext &ASTCtx) {
   auto InTopLevelScope = [](const NamedDecl &ND) {
     switch (ND.getDeclContext()->getDeclKind()) {
     case Decl::TranslationUnit:
     case Decl::Namespace:

clang-tools-extra/clangd/ParsedAST.cpp

ASTDiags.setDiagCallback(
       [&ASTListeners](const clang::Diagnostic &D, clangd::Diag &Diag) {
         llvm::for_each(ASTListeners,
                        [&](const auto &L) { L->sawDiagnostic(D, Diag); });
       });
 
   llvm::Optional<PreamblePatch> Patch;
   bool PreserveDiags = true;
   if (Preamble) {
-    Patch = PreamblePatch::create(Filename, Inputs, *Preamble);
+    Patch = PreamblePatch::createFullPatch(Filename, Inputs, *Preamble);
     Patch->apply(*CI);
     PreserveDiags = Patch->preserveDiagnostics();
   }
   auto Clang = prepareCompilerInstance(
       std::move(CI), PreamblePCH,
       llvm::MemoryBuffer::getMemBufferCopy(Inputs.Contents, Filename), VFS,
       ASTDiags);
   if (!Clang) {

clang-tools-extra/clangd/Preamble.h

 
 /// Stores information required to parse a TU using a (possibly stale) Baseline
 /// preamble. Later on this information can be injected into the main file by
 /// updating compiler invocation with \c apply. This injected section
 /// approximately reflects additions to the preamble in Modified contents, e.g.
 /// new include directives.
 class PreamblePatch {
 public:
+  enum class PatchType { MacroDirectives, All };
   /// \p Preamble is used verbatim.
   static PreamblePatch unmodified(const PreambleData &Preamble);
   /// Builds a patch that contains new PP directives introduced to the preamble
   /// section of \p Modified compared to \p Baseline.
   /// FIXME: This only handles include directives, we should at least handle
   /// define/undef.
-  static PreamblePatch create(llvm::StringRef FileName,
+  static PreamblePatch createFullPatch(llvm::StringRef FileName,
+                                       const ParseInputs &Modified,
+                                       const PreambleData &Baseline);
+  static PreamblePatch createMacroPatch(llvm::StringRef FileName,
                                         const ParseInputs &Modified,
                                         const PreambleData &Baseline);
+
   /// Adjusts CI (which compiles the modified inputs) to be used with the
   /// baseline preamble. This is done by inserting an artifical include to the
   /// \p CI that contains new directives calculated in create.
   void apply(CompilerInvocation &CI) const;
 
   /// 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; }
 
   /// Whether diagnostics generated using this patch are trustable.
   bool preserveDiagnostics() const { return PatchContents.empty(); }
 
 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.
   std::vector<Inclusion> PreambleIncludes;
   PreambleBounds ModifiedBounds = {0, false};
 };

clang-tools-extra/clangd/Preamble.cpp

 };
 
 // Represents directives other than includes, where basic textual information is
 // enough.
 struct TextualPPDirective {
   unsigned DirectiveLine;
   // Full text that's representing the directive, including the `#`.
   std::string Text;
+  unsigned Offset;
 
   bool operator==(const TextualPPDirective &RHS) const {
-    return std::tie(DirectiveLine, Text) ==
-           std::tie(RHS.DirectiveLine, RHS.Text);
+    return std::tie(DirectiveLine, Offset, Text) ==
+           std::tie(RHS.DirectiveLine, RHS.Offset, 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) {
+                           unsigned &DirectiveLine, unsigned &Offset) {
   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);
+  Offset = 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.
void MacroDefined(const Token &MacroNameTok,
     TextualDirectives.emplace_back();
     TextualPPDirective &TD = TextualDirectives.back();
 
     const auto *MI = MD->getMacroInfo();
     TD.Text =
         spellDirective("#define ",
                        CharSourceRange::getTokenRange(
                            MI->getDefinitionLoc(), MI->getDefinitionEndLoc()),
-                       LangOpts, SM, TD.DirectiveLine);
+                       LangOpts, SM, TD.DirectiveLine, TD.Offset);
   }
 
 private:
   bool InMainFile = true;
   const LangOptions &LangOpts;
   const SourceManager &SM;
   std::vector<TextualPPDirective> &TextualDirectives;
 };
default:
       break;
     }
     OS << C;
   }
 }
 
 PreamblePatch PreamblePatch::create(llvm::StringRef FileName,
                                     const ParseInputs &Modified,
-                                    const PreambleData &Baseline) {
+                                    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
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";
 
-  if (IncludesChanged) {
+  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
for (auto &Inc : ModifiedScan->Includes) {
       // 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);
     }
   }
 
   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) {
       Patch << "#line " << TD.DirectiveLine << '\n';
       Patch << TD.Text << '\n';
     }
   }
   dlog("Created preamble patch: {0}", Patch.str());
   Patch.flush();
   return PP;
 }
 
+PreamblePatch PreamblePatch::createFullPatch(llvm::StringRef FileName,
+                                             const ParseInputs &Modified,
+                                             const PreambleData &Baseline) {
+  return create(FileName, Modified, Baseline, PatchType::All);
+}
+
+PreamblePatch PreamblePatch::createMacroPatch(llvm::StringRef FileName,
+                                              const ParseInputs &Modified,
+                                              const PreambleData &Baseline) {
+  return create(FileName, Modified, Baseline, PatchType::MacroDirectives);
+}
+
 void PreamblePatch::apply(CompilerInvocation &CI) const {
   // No need to map an empty file.
   if (PatchContents.empty())
     return;
   auto &PPOpts = CI.getPreprocessorOpts();
   auto PatchBuffer =
       // we copy here to ensure contents are still valid if CI outlives the
       // PreamblePatch.

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

TEST(CompletionTest, ObjCCategoryDecls) {
   @interface Foo (^)
   @end
   )objc");
     TU.Code = Test.code().str();
     auto Results = completions(TU, Test.point());
     EXPECT_THAT(Results.Completions, UnorderedElementsAre(Labeled("BarExt")));
   }
 }
+
+TEST(CompletionTest, PreambleCodeComplete) {
+  llvm::StringLiteral Baseline = "\n#define MACRO 12\nint num = MACRO;";
+  llvm::StringLiteral ModifiedCC =
+      "#include \"header.h\"\n#define MACRO 12\nint num = MACRO; int num2 = M^";
+
+  Annotations Test(ModifiedCC);
+  auto BaselineTU = TestTU::withCode(Baseline);
+  auto ModifiedTU = TestTU::withCode(Test.code());
+
+  MockFS FS;
+  auto Inputs = ModifiedTU.inputs(FS);
+  auto Result = codeComplete(testPath(ModifiedTU.Filename), Test.point(),
+                             BaselineTU.preamble().get(), Inputs, {});
+  EXPECT_THAT(Result.Completions, Not(testing::IsEmpty()));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

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

collectPatchedIncludes(llvm::StringRef ModifiedContents,
   auto TU = TestTU::withCode(BaselineContents);
   TU.Filename = MainFileName.str();
   // ms-compatibility changes meaning of #import, make sure it is turned off.
   TU.ExtraArgs = {"-fno-ms-compatibility"};
   auto BaselinePreamble = TU.preamble();
   // Create the patch.
   TU.Code = ModifiedContents.str();
   auto PI = TU.inputs(FS);
-  auto PP = PreamblePatch::create(testPath(TU.Filename), PI, *BaselinePreamble);
+  auto PP = PreamblePatch::createFullPatch(testPath(TU.Filename), PI,
+                                           *BaselinePreamble);
   // Collect patch contents.
   IgnoreDiagnostics Diags;
   auto CI = buildCompilerInvocation(PI, Diags);
   PP.apply(*CI);
   // Run preprocessor over the modified contents with patched Invocation. We
   // provide a preamble and trim contents to ensure only the implicit header
   // introduced by the patch is parsed and nothing else.
   // We don't run PP directly over the patch cotents to test production
TEST(PreamblePatchTest, PatchesPreambleIncludes) {
   auto BaselinePreamble = buildPreamble(
       TU.Filename, *buildCompilerInvocation(PI, Diags), PI, true, nullptr);
   // We drop c.h from modified and add a new header. Since the latter is patched
   // we should only get a.h in preamble includes.
   TU.Code = R"cpp(
     #include "a.h"
     #include "b.h"
   )cpp";
-  auto PP = PreamblePatch::create(testPath(TU.Filename), TU.inputs(FS),
+  auto PP = PreamblePatch::createFullPatch(testPath(TU.Filename), TU.inputs(FS),
                                            *BaselinePreamble);
   // Only a.h should exists in the preamble, as c.h has been dropped and b.h was
   // newly introduced.
   EXPECT_THAT(PP.preambleIncludes(),
               ElementsAre(AllOf(Field(&Inclusion::Written, "\"a.h\""),
                                 Field(&Inclusion::Resolved, testPath("a.h")))));
 }
 
 llvm::Optional<ParsedAST> createPatchedAST(llvm::StringRef Baseline,
std::string getPreamblePatch(llvm::StringRef Baseline,
                              llvm::StringRef Modified) {
   auto BaselinePreamble = TestTU::withCode(Baseline).preamble();
   if (!BaselinePreamble) {
     ADD_FAILURE() << "Failed to build baseline preamble";
     return "";
   }
   MockFS FS;
   auto TU = TestTU::withCode(Modified);
-  return PreamblePatch::create(testPath("main.cpp"), TU.inputs(FS),
+  return PreamblePatch::createFullPatch(testPath("main.cpp"), TU.inputs(FS),
                                         *BaselinePreamble)
       .text()
       .str();
 }
 
 TEST(PreamblePatchTest, Define) {
   // BAR should be defined while parsing the AST.
   struct {
     const char *const Contents;
TEST(PreamblePatch, ModifiedBounds) {
   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();
     MockFS FS;
-    auto PP = PreamblePatch::create(testPath(TU.Filename), TU.inputs(FS),
-                                    *BaselinePreamble);
+    auto PP = PreamblePatch::createFullPatch(testPath(TU.Filename),
+                                             TU.inputs(FS), *BaselinePreamble);
 
     IgnoreDiagnostics Diags;
     auto CI = buildCompilerInvocation(TU.inputs(FS), 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);
   }
 }
 
 TEST(PreamblePatch, DropsDiagnostics) {
   llvm::StringLiteral Code = "#define FOO\nx;/* error-ok */";
   // First check that this code generates diagnostics.
   EXPECT_THAT(*TestTU::withCode(Code).build().getDiagnostics(),
               testing::Not(testing::IsEmpty()));
   // Ensure they are dropeed when a patched preamble is used.
   EXPECT_FALSE(createPatchedAST("", Code)->getDiagnostics());
 }
+
+TEST(PreamblePatch, MacroLoc) {
+  llvm::StringLiteral Baseline = "\n#define MACRO 12\nint num = MACRO;";
+  llvm::StringLiteral Modified = " \n#define MACRO 12\nint num = MACRO;";
+  auto AST = createPatchedAST(Baseline, Modified);
+  ASSERT_TRUE(AST);
+}
 } // namespace
 } // namespace clangd
 } // namespace clang