[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

511

this is the condition i am suggesting to drop.

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

553

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

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

533

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

clang-tools-extra/clangd/Preamble.h

public:
   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,
                               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

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, Text, Offset) ==
+           std::tie(RHS.DirectiveLine, RHS.Text, RHS.Offset);
   }
 };
 
 // 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;
 };
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::createMacroPatch(llvm::StringRef FileName,
+                                              const ParseInputs &Modified,
+                                              const PreambleData &Baseline) {
+
+  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 PP;
+  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);
+  Patch << "#line 0 \"";
+
+  escapeBackslashAndQuotes(FileName, Patch);
+  Patch << "\"\n";
+
+  if (BaselineScan->TextualDirectives != ModifiedScan->TextualDirectives) {
+    for (const auto &TD : ModifiedScan->TextualDirectives) {
+      Patch << "#line " << TD.DirectiveLine << '\n';
+      Patch << TD.Text << '\n';
+    }
+    dlog("Created macro only preamble patch: {0}", Patch.str());
+    Patch.flush();
+    return PP;
+  }
+  return PreamblePatch::unmodified(Baseline);
+}
+
 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

 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