[clangd] Populate #include insertions as additional edits in completion items.

[ioeric]

Summary

o Remove IncludeInsertion LSP command.
o Populate include insertion edits synchromously in completion items.
o Share the code completion compiler instance and precompiled preamble to get existing inclusions in main file.
o Include insertion logic lives only in CodeComplete now.
o Use tooling::HeaderIncludes for inserting new includes.
o Refactored tests.

[sammccall]

I'm concerned about the scope of this patch - it does too many things and touches too many files for me to feel comfortable that I understand it well enough to review.
Is it possible to split it up? (You mention 6 distinct things in the description).

Have tried to focus comments on what seems to be the core parts.

clangd/CodeComplete.cpp

247

I can't parse this name as it relates to this class. This inserts includes, it isn't itself an insertion.
Maybe IncludeInserter?

clangd/CodeComplete.h

79

IncLocations isn't an appropriate name for the variable - it's coupled to the preamble

clangd/Headers.h

1

The interface to this header looks a bit complex and entangled after this change.

• The construction/lifecycle of IncludePaths is confusing (comments below)
• here we provide all the information needed to compute the edits (style, code...) but don't actually do so

it would be nice to achieve a simpler interface with the restricted set of functionality, or expand the functionality...

28

I had some trouble working out the role of this struct (e.g. that it might have been a *proposed* inclusion).
Strictly it's not just the location, but also identifies the header being included.

So this could maybe be

// An #include directive that we found in the main file.
struct Inclusion

but it does also seem a little strange that this is part of the public interface at all (other than for testing)

39

This doesn't seem to add any semantics over the vector, so it just seems to obscure a bit.
In fact it does have semantics: it's the stored list of inclusion locations for the preamble. But those semantics don't belong in this layer. I'd suggest moving it back to the preamble code, or avoiding the typedef altogether.

67–102

the main value-add of this class over HeaderSearch is that it determines whether to insert, so that should probably be mentioned.

68

can you describe the design a little bit here?
e.g. "We simply keep track of which files were already inserted, and delegate to clang's HeaderSearch if a header is new"

69

This name suggests a set or sequence of include path entries (-I flags).
Can we merge this with IncludeInsertion and call it IncludeInserter or similar?

72

hmm, if we require the InclusionLocations to be prebuilt for the preamble, why can't we do the same for non-preamble includes? i.e. can't we split up the "collect existing includes" and "calculate new includes" parts to simplify the interface here? Among other things, testing would be a ton easier.

75

The contract around PP is really confusing: we need to be able to mutate it (to add PP hooks), we need someone to use it in the middle, we need to use its HeaderSearch structure later, and to add more confusion we have shared ownership.

Maybe instead something like:

• this class takes an IncludeSearch& (sadly not const-correct) to use for lookups
• it exposes a function like std::unique_ptr<PPCallbacks> record() which returns a callback hook which writes into this includeptahs object.

Then usage looks like:

IncludePaths Paths(SM, PreambleLocs, PP.getHeaderSearch());
PP.addPPCallbacks(Paths.record());
// do something with PP
Paths.calculate(...);

and AFAICT, all the lifetime issues/interactions would be communicated by the types.

76

consider adding a addExistingInclude(InclusionLocation) instead of the InclusionLocations parameter here. This gives you somewhere to document what this does and what it's for (constructors are always crowded), and makes the callsites a little more self-documenting.

It also lets you take InclusionLocations out of the public interface here, which I think would be a win.

95

is this used?

97

and this?

107

is this unused? this seems like it could be a natural part of this class...

clangd/tool/ClangdMain.cpp

103 ↗

(On Diff #145385)

why this change?

unittests/clangd/HeadersTests.cpp

97

Is there any way we can avoid creating another copy of this code?

[ioeric]

• Merged with origin/master
• Address review comments.
• Revert unintended change.

[ioeric]

In D46497#1089755, @sammccall wrote:

I'm concerned about the scope of this patch - it does too many things and touches too many files for me to feel comfortable that I understand it well enough to review.
Is it possible to split it up? (You mention 6 distinct things in the description).

Sorry about the big change! And thanks for the early feedback on headers library!

I think there are a few things going on here:

1. Removing the existing way of include insertion (additional command)
2. Populating #include insertion in code completion workflow.
3. Refactoring #include path calculation to make 2 work
4. Move some helpers (e.g. replacementToEdit) to public header.
5. Share InclusionLocation between preamble code and header insertion code.

I was trying to keep 1 and 2 in the same patch so that we don't regress features, but it's probably fine if no one is actually using it... I'm happy to split that out if having them in the same patch makes review hard.

I think 2 and 3 could stay in the same patch as it would be hard to see where the design in 3 came from without 2. But if you got the idea, I guess they could still be split.

4 and 5 should be independent enough to be 2 separate patches. I thought the changes were small and could be included in this patch. I'll split them as well now that you mention.

clangd/Headers.h

1

Refactored code according to your suggestions. PTAL

here we provide all the information needed to compute the edits (style, code...) but don't actually do so

Sorry... forgot to delete some leftover after an earlier revision.

28

Thanks! Inclusion sounds better.

Apart from header insertion and testing, this is also shared with the preamble code.

69

I think there are three pretty independent parts of include insertion:

1. URI/Verbatim to HeaderFile conversion (which lives in HeaderInsertion in CodeComplete.cpp right now).
2. Calculate include paths based on existing inclusions and HeaderSearch.
3. Calculate include insertion edits using tooling::HeaderIncludes.

Initially, I grouped all these in one class (something like HeaderInsertion) which returns an edit given the URIs in the symbol info, but it turned out to be hard to test. For example, among all 3 components, 2 was the most interesting piece to test, while an integration test is probably enough for 1 and 3, but we would need to create URIs (due to 1) and retrieve include name from edits (due to 3) in order to test 2. Since each individual piece seemed independent enough, I decided to keep them separated. I also kept URI logic out of the header library so that it could potentially be shared with other features that are not index-based include insertion, or be moved out of clangd as a tooling library.

In the new revision, I refactored the code according to your suggestions and got rid of this class.

72

Good idea!

107

Thanks for the catches! Missed these after iterating from an earlier version.,,

clangd/tool/ClangdMain.cpp

103 ↗

(On Diff #145385)

Oops, this was not intended. I increased the limit for profiling. Reverted...

unittests/clangd/HeadersTests.cpp

97

This was moved here. Headers library no longer needs compiler instance. I'm not sure if there is an easier way to test HeaderSearch and PPCallbacks without a compiler instance...

[ioeric]

• Minor cleanup

[ioeric]

• Rebase on D46670, D46675, D46676

[ioeric]

This has been split and now depends on D46670, D46675, D46676.

[sammccall]

clangd/CodeComplete.cpp

247

Do you think it would make sense to move this class to Headers.h? This would reduce the non-completion related code in this large file.

261

If we wanted to avoid URI dependencies in Headers.h but move this class there, the Declaring/Inserted headers could be HeaderFile (possibly invalid). toHeaderFile would stay here.

This might be a neater separation of concerns anyway.

317

This strikes me as overly defensive. We should never have an include header but no declaration location, and if the declaration header is invalid then the include inserter should take care of that, not this code (I think?)

700

PreambleInclusions?

969

this callback is a bit confusing, particularly sequencing/lifetimes.
semaCodeComplete is local, could we simplify it by coupling a bit, passing an IncludeInserter* ("will be populated by the preprocessor") and moving this code (and the preamble-includes handling) inside semaCodeComplete?

This would be a bit more natural if we had IncludeInsertion::addExisting(Insertion) instead of IncludeInsertion::setExisting(vector<Insertion>), I think.

1130

nit: /*Inclusions=*/

clangd/CodeComplete.h

73

should this be PreambleInclusions, or is it broader than that?

clangd/Headers.h

69

OK, SGTM

unittests/clangd/HeadersTests.cpp

97

Moved from where?

You could consider adding inclusions to TestTU (sorry it's in D46524 which I thought had already landed, but should today)

[ioeric]

• Addressed review comments.
• Merged with origin/master

[sammccall]

Nice, ship it!

clangd/CodeComplete.cpp

317

is this comment still relevant here?

791

nit: FormatStyle

802

nit: includes is a pointer, capture by value?

961

This leaks. reset (or = nullptr)

[ioeric]

• Addressed review comments.

[ioeric]

• Rebase on D46676

[ioeric]

• Rebase...

[ioeric]

• Merged with origin/master

clangd/ClangdServer.cpp

if (!IP)
       return CB(IP.takeError());
 
     auto PreambleData = IP->Preamble;
 
     // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
     // both the old and the new version in case only one of them matches.
     CompletionList Result = clangd::codeComplete(
         File, IP->Command, PreambleData ? &PreambleData->Preamble : nullptr,
+        PreambleData ? PreambleData->Inclusions : std::vector<Inclusion>(),
         IP->Contents, Pos, FS, PCHs, CodeCompleteOpts);
     CB(std::move(Result));
   };
 
   WorkScheduler.runWithPreamble("CodeComplete", File,
                                 Bind(Task, File.str(), std::move(CB)));
 }
 

clangd/CodeComplete.h

 // Code completion provides suggestions for what the user might type next.
 // After "std::string S; S." we might suggest members of std::string.
 // Signature help describes the parameters of a function as you type them.
 //
 //===---------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CODECOMPLETE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CODECOMPLETE_H
 
+#include "Headers.h"
 #include "Logger.h"
 #include "Path.h"
 #include "Protocol.h"
 #include "index/Index.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Sema/CodeCompleteOptions.h"
 #include "clang/Tooling/CompilationDatabase.h"
 
struct CodeCompleteOptions {
   /// clangd.
   const SymbolIndex *Index = nullptr;
 };
 
 /// Get code completions at a specified \p Pos in \p FileName.
 CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
+                            const std::vector<Inclusion> &PreambleInclusions,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs,
                             CodeCompleteOptions Opts);
 
 /// Get signature help at a specified \p Pos in \p FileName.
 SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/CodeComplete.cpp

 // as there are typically fewer candidates.
 //
 //===---------------------------------------------------------------------===//
 
 #include "CodeComplete.h"
 #include "CodeCompletionStrings.h"
 #include "Compiler.h"
 #include "FuzzyMatch.h"
+#include "Headers.h"
 #include "Logger.h"
 #include "SourceCode.h"
 #include "Trace.h"
+#include "URI.h"
 #include "index/Index.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Index/USRGeneration.h"
 #include "clang/Sema/CodeCompleteConsumer.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Tooling/Core/Replacement.h"
std::string sortText(float Score, llvm::StringRef Name) {
   llvm::raw_string_ostream OS(S);
   write_hex(OS, encodeFloat(-Score), llvm::HexPrintStyle::Lower,
             /*Width=*/2 * sizeof(Score));
   OS << Name;
   OS.flush();
   return S;
 }
 
+/// Creates a `HeaderFile` from \p Header which can be either a URI or a literal
+/// include.
+static llvm::Expected<HeaderFile> toHeaderFile(StringRef Header,
+                                               llvm::StringRef HintPath) {
+  if (isLiteralInclude(Header))
+    return HeaderFile{Header.str(), /*Verbatim=*/true};
+  auto U = URI::parse(Header);
+  if (!U)
+    return U.takeError();
+
+  auto IncludePath = URI::includeSpelling(*U);
+  if (!IncludePath)
+    return IncludePath.takeError();
+  if (!IncludePath->empty())
+    return HeaderFile{std::move(*IncludePath), /*Verbatim=*/true};
+
+  auto Resolved = URI::resolve(*U, HintPath);
+  if (!Resolved)
+    return Resolved.takeError();
+  return HeaderFile{std::move(*Resolved), /*Verbatim=*/false};
+}
+
 /// A code completion result, in clang-native form.
 /// It may be promoted to a CompletionItem if it's among the top-ranked results.
 struct CompletionCandidate {
   llvm::StringRef Name; // Used for filtering and sorting.
   // We may have a result from Sema, from the index, or both.
   const CodeCompletionResult *SemaResult = nullptr;
   const Symbol *IndexResult = nullptr;
 
   // Computes the "symbol quality" score for this completion. Higher is better.
   float score() const {
     float Score = 1;
     if (IndexResult)
       Score *= quality(*IndexResult);
     if (SemaResult) {
       // For now we just use the Sema priority, mapping it onto a 0-2 interval.
       // That makes 1 neutral-ish, so we don't reward/penalize non-Sema results.
       // Priority 80 is a really bad score.
       Score *= 2 - std::min<float>(80, SemaResult->Priority) / 40;
 
       switch (static_cast<CXAvailabilityKind>(SemaResult->Availability)) {
       case CXAvailability_Available:
         // No penalty.
         break;
       case CXAvailability_Deprecated:
         Score *= 0.1f;
         break;
       case CXAvailability_NotAccessible:
       case CXAvailability_NotAvailable:
         Score = 0;
         break;
       }
     }
     return Score;
   }
 
   // Builds an LSP completion item.
-  CompletionItem build(llvm::StringRef FileName,
-                       const CompletionItemScores &Scores,
+  CompletionItem build(StringRef FileName, const CompletionItemScores &Scores,
                        const CodeCompleteOptions &Opts,
-                       CodeCompletionString *SemaCCS) const {
+                       CodeCompletionString *SemaCCS,
+                       const IncludeInserter *Includes) const {
     assert(bool(SemaResult) == bool(SemaCCS));
     CompletionItem I;
     if (SemaResult) {
       I.kind = toCompletionItemKind(SemaResult->Kind, SemaResult->CursorKind);
       getLabelAndInsertText(*SemaCCS, &I.label, &I.insertText,
                             Opts.EnableSnippets);
       I.filterText = getFilterText(*SemaCCS);
       I.documentation = getDocumentation(*SemaCCS);
if (IndexResult) {
                            ? IndexResult->CompletionSnippetInsertText
                            : IndexResult->CompletionPlainInsertText;
 
       if (auto *D = IndexResult->Detail) {
         if (I.documentation.empty())
           I.documentation = D->Documentation;
         if (I.detail.empty())
           I.detail = D->CompletionDetail;
+        if (Includes && !D->IncludeHeader.empty()) {
+          auto Edit = [&]() -> Expected<Optional<TextEdit>> {
+            auto ResolvedDeclaring = toHeaderFile(
+                IndexResult->CanonicalDeclaration.FileURI, FileName);
+            if (!ResolvedDeclaring)
+              return ResolvedDeclaring.takeError();
+            auto ResolvedInserted = toHeaderFile(D->IncludeHeader, FileName);
+            if (!ResolvedInserted)
+              return ResolvedInserted.takeError();
+            return Includes->insert(*ResolvedDeclaring, *ResolvedInserted);
+          }();
+          if (!Edit) {
+            std::string ErrMsg =
+                ("Failed to generate include insertion edits for adding header "
+                 "(FileURI=\"" +
+                 IndexResult->CanonicalDeclaration.FileURI +
+                 "\", IncludeHeader=\"" + D->IncludeHeader + "\") into " +
+                 FileName)
+                    .str();
+            log(ErrMsg + ":" + llvm::toString(Edit.takeError()));
+          } else if (*Edit) {
+            I.additionalTextEdits = {std::move(**Edit)};
+          }
+        }
       }
     }
     I.scoreInfo = Scores;
     I.sortText = sortText(Scores.finalScore, Name);
     I.insertTextFormat = Opts.EnableSnippets ? InsertTextFormat::Snippet
                                              : InsertTextFormat::PlainText;
     return I;
   }
private:
   CodeCompletionTUInfo CCTUInfo;
 
 }; // SignatureHelpCollector
 
 struct SemaCompleteInput {
   PathRef FileName;
   const tooling::CompileCommand &Command;
   PrecompiledPreamble const *Preamble;
+  const std::vector<Inclusion> &PreambleInclusions;
   StringRef Contents;
   Position Pos;
   IntrusiveRefCntPtr<vfs::FileSystem> VFS;
   std::shared_ptr<PCHContainerOperations> PCHs;
 };
 
 // Invokes Sema code completion on a file.
+// If \p Includes is set, it will be initialized after a compiler instance has
+// been set up.
 bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
                       const clang::CodeCompleteOptions &Options,
-                      const SemaCompleteInput &Input) {
+                      const SemaCompleteInput &Input,
+                      std::unique_ptr<IncludeInserter> *Includes = nullptr) {
   trace::Span Tracer("Sema completion");
   std::vector<const char *> ArgStrs;
   for (const auto &S : Input.Command.CommandLine)
     ArgStrs.push_back(S.c_str());
 
   if (Input.VFS->setCurrentWorkingDirectory(Input.Command.Directory)) {
     log("Couldn't set working directory");
     // We run parsing anyway, our lit-tests rely on results for non-existing
bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
   Clang->setCodeCompletionConsumer(Consumer.release());
 
   SyntaxOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
     log("BeginSourceFile() failed when running codeComplete for " +
         Input.FileName);
     return false;
   }
+  if (Includes) {
+    // Initialize Includes if provided.
+
+    // FIXME(ioeric): needs more consistent style support in clangd server.
+    auto Style = format::getStyle("file", Input.FileName, "LLVM",
+                                  Input.Contents, Input.VFS.get());
+    if (!Style) {
+      log("Failed to get FormatStyle for file" + Input.FileName +
+          ". Fall back to use LLVM style. Error: " +
+          llvm::toString(Style.takeError()));
+      Style = format::getLLVMStyle();
+    }
+    *Includes = llvm::make_unique<IncludeInserter>(
+        Input.FileName, Input.Contents, *Style, Input.Command.Directory,
+        Clang->getPreprocessor().getHeaderSearchInfo());
+    for (const auto &Inc : Input.PreambleInclusions)
+      Includes->get()->addExisting(Inc);
+    Clang->getPreprocessor().addPPCallbacks(collectInclusionsInMainFileCallback(
+        Clang->getSourceManager(), [Includes](Inclusion Inc) {
+          Includes->get()->addExisting(std::move(Inc));
+        }));
+  }
   if (!Action.Execute()) {
     log("Execute() failed when running codeComplete for " + Input.FileName);
     return false;
   }
   Action.EndSourceFile();
 
   return true;
 }

 //   - TopN determines the results with the best score.
 class CodeCompleteFlow {
   PathRef FileName;
   const CodeCompleteOptions &Opts;
   // Sema takes ownership of Recorder. Recorder is valid until Sema cleanup.
   CompletionRecorder *Recorder = nullptr;
   int NSema = 0, NIndex = 0, NBoth = 0; // Counters for logging.
   bool Incomplete = false; // Would more be available with a higher limit?
-  llvm::Optional<FuzzyMatcher> Filter; // Initialized once Sema runs.
+  llvm::Optional<FuzzyMatcher> Filter;       // Initialized once Sema runs.
+  std::unique_ptr<IncludeInserter> Includes; // Initialized once compiler runs.
 
 public:
   // A CodeCompleteFlow object is only useful for calling run() exactly once.
   CodeCompleteFlow(PathRef FileName, const CodeCompleteOptions &Opts)
       : FileName(FileName), Opts(Opts) {}
 
   CompletionList run(const SemaCompleteInput &SemaCCInput) && {
     trace::Span Tracer("CodeCompleteFlow");
+
     // We run Sema code completion first. It builds an AST and calculates:
     //   - completion results based on the AST.
     //   - partial identifier and context. We need these for the index query.
     CompletionList Output;
     auto RecorderOwner = llvm::make_unique<CompletionRecorder>(Opts, [&]() {
       assert(Recorder && "Recorder is not set");
+      assert(Includes && "Includes is not set");
+      // If preprocessor was run, inclusions from preprocessor callback should
+      // already be added to Inclusions.
       Output = runWithSema();
+      Includes.reset(); // Make sure this doesn't out-live Clang.
       SPAN_ATTACH(Tracer, "sema_completion_kind",
                   getCompletionKindString(Recorder->CCContext.getKind()));
     });
 
     Recorder = RecorderOwner.get();
     semaCodeComplete(std::move(RecorderOwner), Opts.getClangCompleteOpts(),
-                     SemaCCInput);
+                     SemaCCInput, &Includes);
 
     SPAN_ATTACH(Tracer, "sema_results", NSema);
     SPAN_ATTACH(Tracer, "index_results", NIndex);
     SPAN_ATTACH(Tracer, "merged_results", NBoth);
     SPAN_ATTACH(Tracer, "returned_results", Output.items.size());
     SPAN_ATTACH(Tracer, "incomplete", Output.isIncomplete);
     log(llvm::formatv("Code complete: {0} results from Sema, {1} from Index, "
                       "{2} matched, {3} returned{4}.",
                       NSema, NIndex, NBoth, Output.items.size(),
if (Candidates.push({C, Scores}))
       Incomplete = true;
   }
 
   CompletionItem toCompletionItem(const CompletionCandidate &Candidate,
                                   const CompletionItemScores &Scores) {
     CodeCompletionString *SemaCCS = nullptr;
     if (auto *SR = Candidate.SemaResult)
       SemaCCS = Recorder->codeCompletionString(*SR, Opts.IncludeBriefComments);
-    return Candidate.build(FileName, Scores, Opts, SemaCCS);
+    return Candidate.build(FileName, Scores, Opts, SemaCCS, Includes.get());
   }
 };
 
 CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
+                            const std::vector<Inclusion> &PreambleInclusions,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs,
                             CodeCompleteOptions Opts) {
   return CodeCompleteFlow(FileName, Opts)
-      .run({FileName, Command, Preamble, Contents, Pos, VFS, PCHs});
+      .run({FileName, Command, Preamble, PreambleInclusions, Contents, Pos, VFS,
+            PCHs});
 }
 
 SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs) {
   SignatureHelp Result;
   clang::CodeCompleteOptions Options;
   Options.IncludeGlobals = false;
   Options.IncludeMacros = false;
   Options.IncludeCodePatterns = false;
   Options.IncludeBriefComments = true;
   semaCodeComplete(llvm::make_unique<SignatureHelpCollector>(Options, Result),
                    Options,
-                   {FileName, Command, Preamble, Contents, Pos, std::move(VFS),
+                   {FileName, Command, Preamble,
+                    /*PreambleInclusions=*/{}, Contents, Pos, std::move(VFS),
                     std::move(PCHs)});
   return Result;
 }
 
 } // namespace clangd
 } // namespace clang

clangd/Headers.h

 //===--- Headers.h - Include headers -----------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H
 
 #include "Path.h"
 #include "Protocol.h"
+#include "SourceCode.h"
 #include "clang/Basic/VirtualFileSystem.h"
+#include "clang/Format/Format.h"
+#include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PPCallbacks.h"
-#include "clang/Tooling/CompilationDatabase.h"
+#include "clang/Tooling/Core/HeaderIncludes.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 
 namespace clang {
 namespace clangd {
 
 /// Returns true if \p Include is literal include like "path" or <path>.
 bool isLiteralInclude(llvm::StringRef Include);
 
 /// Represents a header file to be #include'd.
 struct HeaderFile {
   std::string File;
   /// If this is true, `File` is a literal string quoted with <> or "" that
   /// can be #included directly; otherwise, `File` is an absolute file path.
   bool Verbatim;
 
   bool valid() const;
 };
 
 // An #include directive that we found in the main file.
 struct Inclusion {
   Range R;             // Inclusion range.
   std::string Written; // Inclusion name as written e.g. <vector>.
   Path Resolved;       // Resolved path of included file. Empty if not resolved.
 };
 
 /// Returns a PPCallback that visits all inclusions in the main file.
 std::unique_ptr<PPCallbacks>
 collectInclusionsInMainFileCallback(const SourceManager &SM,
                                     std::function<void(Inclusion)> Callback);
 
 /// Determines the preferred way to #include a file, taking into account the
 /// search path. Usually this will prefer a shorter representation like
 /// 'Foo/Bar.h' over a longer one like 'Baz/include/Foo/Bar.h'.
 ///
 /// \param File is an absolute file path.
+/// \param Inclusions Existing inclusions in the main file.
 /// \param DeclaringHeader is the original header corresponding to \p
 /// InsertedHeader e.g. the header that declares a symbol.
 /// \param InsertedHeader The preferred header to be inserted. This could be the
 /// same as DeclaringHeader but must be provided.
 //  \return A quoted "path" or <path>. This returns an empty string if:
 ///   - Either \p DeclaringHeader or \p InsertedHeader is already (directly)
-///   included in the file (including those included via different paths).
+///   in \p Inclusions (including those included via different paths).
 ///   - \p DeclaringHeader or \p InsertedHeader is the same as \p File.
-llvm::Expected<std::string>
-calculateIncludePath(PathRef File, llvm::StringRef Code,
-                     const HeaderFile &DeclaringHeader,
-                     const HeaderFile &InsertedHeader,
-                     const tooling::CompileCommand &CompileCommand,
-                     IntrusiveRefCntPtr<vfs::FileSystem> FS);
+llvm::Expected<std::string> calculateIncludePath(
+    PathRef File, StringRef BuildDir, HeaderSearch &HeaderSearchInfo,
+    const std::vector<Inclusion> &Inclusions, const HeaderFile &DeclaringHeader,
+    const HeaderFile &InsertedHeader);
+
+// Calculates insertion edit for including a new header in a file.
+class IncludeInserter {
+public:
+  IncludeInserter(StringRef FileName, StringRef Code,
+                  const format::FormatStyle &Style, StringRef BuildDir,
+                  HeaderSearch &HeaderSearchInfo)
+      : FileName(FileName), Code(Code), BuildDir(BuildDir),
+        HeaderSearchInfo(HeaderSearchInfo),
+        Inserter(FileName, Code, Style.IncludeStyle) {}
+
+  void addExisting(Inclusion Inc) { Inclusions.push_back(std::move(Inc)); }
+
+  /// Returns a TextEdit that inserts a new header; if the header is not
+  /// inserted e.g. it's an existing header, this returns None. If any header is
+  /// invalid, this returns error.
+  ///
+  /// \param DeclaringHeader is the original header corresponding to \p
+  /// InsertedHeader e.g. the header that declares a symbol.
+  /// \param InsertedHeader The preferred header to be inserted. This could be
+  /// the same as DeclaringHeader but must be provided.
+  Expected<Optional<TextEdit>> insert(const HeaderFile &DeclaringHeader,
+                                      const HeaderFile &InsertedHeader) const;
+
+private:
+  StringRef FileName;
+  StringRef Code;
+  StringRef BuildDir;
+  HeaderSearch &HeaderSearchInfo;
+  std::vector<Inclusion> Inclusions;
+  tooling::HeaderIncludes Inserter; // Computers insertion replacement.
+};
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H

clangd/Headers.cpp

 #include "Headers.h"
 #include "Compiler.h"
 #include "Logger.h"
 #include "SourceCode.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/HeaderSearch.h"
-#include "clang/Lex/PreprocessorOptions.h"
-#include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 class RecordHeaders : public PPCallbacks {
 public:

 std::unique_ptr<PPCallbacks>
 collectInclusionsInMainFileCallback(const SourceManager &SM,
                                     std::function<void(Inclusion)> Callback) {
   return llvm::make_unique<RecordHeaders>(SM, std::move(Callback));
 }
 
 /// FIXME(ioeric): we might not want to insert an absolute include path if the
 /// path is not shortened.
-llvm::Expected<std::string>
-calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
-                     const HeaderFile &DeclaringHeader,
-                     const HeaderFile &InsertedHeader,
-                     const tooling::CompileCommand &CompileCommand,
-                     IntrusiveRefCntPtr<vfs::FileSystem> FS) {
-  assert(llvm::sys::path::is_absolute(File));
+llvm::Expected<std::string> calculateIncludePath(
+    PathRef File, StringRef BuildDir, HeaderSearch &HeaderSearchInfo,
+    const std::vector<Inclusion> &Inclusions, const HeaderFile &DeclaringHeader,
+    const HeaderFile &InsertedHeader) {
   assert(DeclaringHeader.valid() && InsertedHeader.valid());
   if (File == DeclaringHeader.File || File == InsertedHeader.File)
     return "";
-  FS->setCurrentWorkingDirectory(CompileCommand.Directory);
-
-  // Set up a CompilerInstance and create a preprocessor to collect existing
-  // #include headers in \p Code. Preprocesor also provides HeaderSearch with
-  // which we can calculate the shortest include path for \p Header.
-  std::vector<const char *> Argv;
-  for (const auto &S : CompileCommand.CommandLine)
-    Argv.push_back(S.c_str());
-  IgnoringDiagConsumer IgnoreDiags;
-  auto CI = clang::createInvocationFromCommandLine(
-      Argv,
-      CompilerInstance::createDiagnostics(new DiagnosticOptions(), &IgnoreDiags,
-                                          false),
-      FS);
-  if (!CI)
-    return llvm::make_error<llvm::StringError>(
-        "Failed to create a compiler instance for " + File,
-        llvm::inconvertibleErrorCode());
-  CI->getFrontendOpts().DisableFree = false;
-  // Parse the main file to get all existing #includes in the file, and then we
-  // can make sure the same header (even with different include path) is not
-  // added more than once.
-  CI->getPreprocessorOpts().SingleFileParseMode = true;
-
-  // The diagnostic options must be set before creating a CompilerInstance.
-  CI->getDiagnosticOpts().IgnoreWarnings = true;
-  auto Clang = prepareCompilerInstance(
-      std::move(CI), /*Preamble=*/nullptr,
-      llvm::MemoryBuffer::getMemBuffer(Code, File),
-      std::make_shared<PCHContainerOperations>(), FS, IgnoreDiags);
-
-  if (Clang->getFrontendOpts().Inputs.empty())
-    return llvm::make_error<llvm::StringError>(
-        "Empty frontend action inputs empty for file " + File,
-        llvm::inconvertibleErrorCode());
-  PreprocessOnlyAction Action;
-  if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
-    return llvm::make_error<llvm::StringError>(
-        "Failed to begin preprocessor only action for file " + File,
-        llvm::inconvertibleErrorCode());
-  std::vector<Inclusion> Inclusions;
-  Clang->getPreprocessor().addPPCallbacks(collectInclusionsInMainFileCallback(
-      Clang->getSourceManager(),
-      [&Inclusions](Inclusion Inc) { Inclusions.push_back(std::move(Inc)); }));
-  if (!Action.Execute())
-    return llvm::make_error<llvm::StringError>(
-        "Failed to execute preprocessor only action for file " + File,
-        llvm::inconvertibleErrorCode());
   llvm::StringSet<> IncludedHeaders;
   for (const auto &Inc : Inclusions) {
     IncludedHeaders.insert(Inc.Written);
     if (!Inc.Resolved.empty())
       IncludedHeaders.insert(Inc.Resolved);
   }
   auto Included = [&](llvm::StringRef Header) {
     return IncludedHeaders.find(Header) != IncludedHeaders.end();
   };
   if (Included(DeclaringHeader.File) || Included(InsertedHeader.File))
     return "";
 
-  auto &HeaderSearchInfo = Clang->getPreprocessor().getHeaderSearchInfo();
   bool IsSystem = false;
 
   if (InsertedHeader.Verbatim)
     return InsertedHeader.File;
 
   std::string Suggested = HeaderSearchInfo.suggestPathToFileForDiagnostics(
-      InsertedHeader.File, CompileCommand.Directory, &IsSystem);
+      InsertedHeader.File, BuildDir, &IsSystem);
   if (IsSystem)
     Suggested = "<" + Suggested + ">";
   else
     Suggested = "\"" + Suggested + "\"";
 
   log("Suggested #include for " + InsertedHeader.File + " is: " + Suggested);
   return Suggested;
 }
 
+Expected<Optional<TextEdit>>
+IncludeInserter::insert(const HeaderFile &DeclaringHeader,
+                        const HeaderFile &InsertedHeader) const {
+  auto Validate = [](const HeaderFile &Header) {
+    return Header.valid()
+               ? llvm::Error::success()
+               : llvm::make_error<llvm::StringError>(
+                     "Invalid HeaderFile: " + Header.File +
+                         " (verbatim=" + std::to_string(Header.Verbatim) + ").",
+                     llvm::inconvertibleErrorCode());
+  };
+  if (auto Err = Validate(DeclaringHeader))
+    return std::move(Err);
+  if (auto Err = Validate(InsertedHeader))
+    return std::move(Err);
+  auto Include =
+      calculateIncludePath(FileName, BuildDir, HeaderSearchInfo, Inclusions,
+                           DeclaringHeader, InsertedHeader);
+  if (!Include)
+    return Include.takeError();
+  if (Include->empty())
+    return llvm::None;
+  StringRef IncludeRef = *Include;
+  auto Insertion =
+      Inserter.insert(IncludeRef.trim("\"<>"), IncludeRef.startswith("<"));
+  if (!Insertion)
+    return llvm::None;
+  return replacementToEdit(Code, *Insertion);
+}
+
 } // namespace clangd
 } // namespace clang

clangd/Protocol.h

struct Position {
   /// WARNING: this is in UTF-16 codepoints, not bytes or characters!
   /// Use the functions in SourceCode.h to construct/interpret Positions.
   int character = 0;
 
   friend bool operator==(const Position &LHS, const Position &RHS) {
     return std::tie(LHS.line, LHS.character) ==
            std::tie(RHS.line, RHS.character);
   }
+  friend bool operator!=(const Position &LHS, const Position &RHS) {
+    return !(LHS == RHS);
+  }
   friend bool operator<(const Position &LHS, const Position &RHS) {
     return std::tie(LHS.line, LHS.character) <
            std::tie(RHS.line, RHS.character);
   }
   friend bool operator<=(const Position &LHS, const Position &RHS) {
     return std::tie(LHS.line, LHS.character) <=
            std::tie(RHS.line, RHS.character);
   }

unittests/clangd/CodeCompleteTests.cpp

 
 // GMock helpers for matching completion items.
 MATCHER_P(Named, Name, "") { return arg.insertText == Name; }
 MATCHER_P(Labeled, Label, "") { return arg.label == Label; }
 MATCHER_P(Kind, K, "") { return arg.kind == K; }
 MATCHER_P(Filter, F, "") { return arg.filterText == F; }
 MATCHER_P(Doc, D, "") { return arg.documentation == D; }
 MATCHER_P(Detail, D, "") { return arg.detail == D; }
+MATCHER_P(InsertInclude, IncludeHeader, "") {
+  if (arg.additionalTextEdits.size() != 1)
+    return false;
+  const auto &Edit = arg.additionalTextEdits[0];
+  if (Edit.range.start != Edit.range.end)
+    return false;
+  SmallVector<StringRef, 2> Matches;
+  llvm::Regex RE(R"(#include[ ]*(["<][^">]*[">]))");
+  return RE.match(Edit.newText, &Matches) && Matches[1] == IncludeHeader;
+}
 MATCHER_P(PlainText, Text, "") {
   return arg.insertTextFormat == clangd::InsertTextFormat::PlainText &&
          arg.insertText == Text;
 }
 MATCHER_P(Snippet, Text, "") {
   return arg.insertTextFormat == clangd::InsertTextFormat::Snippet &&
          arg.insertText == Text;
 }
 MATCHER(NameContainsFilter, "") {
   if (arg.filterText.empty())
     return true;
   return llvm::StringRef(arg.insertText).contains(arg.filterText);
 }
+MATCHER(HasAdditionalEdits, "") { return !arg.additionalTextEdits.empty(); }
+
 // Shorthand for Contains(Named(Name)).
 Matcher<const std::vector<CompletionItem> &> Has(std::string Name) {
   return Contains(Named(std::move(Name)));
 }
 Matcher<const std::vector<CompletionItem> &> Has(std::string Name,
                                                  CompletionItemKind K) {
   return Contains(AllOf(Named(std::move(Name)), Kind(K)));
 }
 MATCHER(IsDocumented, "") { return !arg.documentation.empty(); }
 
 std::unique_ptr<SymbolIndex> memIndex(std::vector<Symbol> Symbols) {
   SymbolSlab::Builder Slab;
   for (const auto &Sym : Symbols)
     Slab.insert(Sym);
   return MemIndex::build(std::move(Slab).build());
 }
 
-// Builds a server and runs code completion.
-// If IndexSymbols is non-empty, an index will be built and passed to opts.
-CompletionList completions(StringRef Text,
+CompletionList completions(ClangdServer &Server, StringRef Text,
                            std::vector<Symbol> IndexSymbols = {},
                            clangd::CodeCompleteOptions Opts = {}) {
   std::unique_ptr<SymbolIndex> OverrideIndex;
   if (!IndexSymbols.empty()) {
     assert(!Opts.Index && "both Index and IndexSymbols given!");
     OverrideIndex = memIndex(std::move(IndexSymbols));
     Opts.Index = OverrideIndex.get();
   }
 
-  MockFSProvider FS;
-  MockCompilationDatabase CDB;
-  IgnoreDiagnostics DiagConsumer;
-  ClangdServer Server(CDB, FS, DiagConsumer, ClangdServer::optsForTest());
   auto File = testPath("foo.cpp");
   Annotations Test(Text);
   runAddDocument(Server, File, Test.code());
   auto CompletionList =
       cantFail(runCodeComplete(Server, File, Test.point(), Opts));
   // Sanity-check that filterText is valid.
   EXPECT_THAT(CompletionList.items, Each(NameContainsFilter()));
   return CompletionList;
 }
 
+// Builds a server and runs code completion.
+// If IndexSymbols is non-empty, an index will be built and passed to opts.
+CompletionList completions(StringRef Text,
+                           std::vector<Symbol> IndexSymbols = {},
+                           clangd::CodeCompleteOptions Opts = {}) {
+  MockFSProvider FS;
+  MockCompilationDatabase CDB;
+  IgnoreDiagnostics DiagConsumer;
+  ClangdServer Server(CDB, FS, DiagConsumer, ClangdServer::optsForTest());
+  return completions(Server, Text, std::move(IndexSymbols), std::move(Opts));
+}
+
 std::string replace(StringRef Haystack, StringRef Needle, StringRef Repl) {
   std::string Result;
   raw_string_ostream OS(Result);
   std::pair<StringRef, StringRef> Split;
   for (Split = Haystack.split(Needle); !Split.second.empty();
        Split = Split.first.split(Needle))
     OS << Split.first << Repl;
   Result += Split.first;
auto Results = completions(
           namespace ns { int local; void both(); }
           void f() { ::ns::^ }
       )cpp",
       {func("ns::both"), cls("ns::Index")}, Opts);
   EXPECT_EQ(Results.items.size(), Opts.Limit);
   EXPECT_TRUE(Results.isIncomplete);
 }
 
+TEST(CompletionTest, IncludeInsertionPreprocessorIntegrationTests) {
+  MockFSProvider FS;
+  MockCompilationDatabase CDB;
+  std::string Subdir = testPath("sub");
+  std::string SearchDirArg = (llvm::Twine("-I") + Subdir).str();
+  CDB.ExtraClangFlags = {SearchDirArg.c_str()};
+  std::string BarHeader = testPath("sub/bar.h");
+  FS.Files[BarHeader] = "";
+
+  IgnoreDiagnostics DiagConsumer;
+  ClangdServer Server(CDB, FS, DiagConsumer, ClangdServer::optsForTest());
+  Symbol::Details Scratch;
+  auto BarURI = URI::createFile(BarHeader).toString();
+  Symbol Sym = cls("ns::X");
+  Sym.CanonicalDeclaration.FileURI = BarURI;
+  Scratch.IncludeHeader = BarURI;
+  Sym.Detail = &Scratch;
+  // Shoten include path based on search dirctory and insert.
+  auto Results = completions(Server,
+                             R"cpp(
+          int main() { ns::^ }
+      )cpp",
+                             {Sym});
+  EXPECT_THAT(Results.items,
+              ElementsAre(AllOf(Named("X"), InsertInclude("\"bar.h\""))));
+  // Duplicate based on inclusions in preamble.
+  Results = completions(Server,
+                        R"cpp(
+          #include "sub/bar.h"  // not shortest, so should only match resolved.
+          int main() { ns::^ }
+      )cpp",
+                        {Sym});
+  EXPECT_THAT(Results.items,
+              ElementsAre(AllOf(Named("X"), Not(HasAdditionalEdits()))));
+}
+
 TEST(CompletionTest, IndexSuppressesPreambleCompletions) {
   MockFSProvider FS;
   MockCompilationDatabase CDB;
   IgnoreDiagnostics DiagConsumer;
   ClangdServer Server(CDB, FS, DiagConsumer, ClangdServer::optsForTest());
 
   FS.Files[testPath("bar.h")] =
       R"cpp(namespace ns { struct preamble { int member; }; })cpp";

unittests/clangd/HeadersTests.cpp

 //===-- HeadersTests.cpp - Include headers unit tests -----------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Headers.h"
+
+#include "Compiler.h"
 #include "TestFS.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/CompilerInvocation.h"
+#include "clang/Frontend/FrontendActions.h"
+#include "clang/Lex/PreprocessorOptions.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
+using ::testing::AllOf;
+using ::testing::UnorderedElementsAre;
+
 class HeadersTest : public ::testing::Test {
 public:
   HeadersTest() {
     CDB.ExtraClangFlags = {SearchDirArg.c_str()};
     FS.Files[MainFile] = "";
+    // Make sure directory sub/ exists.
+    FS.Files[testPath("sub/EMPTY")] = "";
+  }
+
+private:
+  std::unique_ptr<CompilerInstance> setupClang() {
+    auto Cmd = CDB.getCompileCommand(MainFile);
+    assert(static_cast<bool>(Cmd));
+    auto VFS = FS.getFileSystem();
+    VFS->setCurrentWorkingDirectory(Cmd->Directory);
+
+    std::vector<const char *> Argv;
+    for (const auto &S : Cmd->CommandLine)
+      Argv.push_back(S.c_str());
+    auto CI = clang::createInvocationFromCommandLine(
+        Argv,
+        CompilerInstance::createDiagnostics(new DiagnosticOptions(),
+                                            &IgnoreDiags, false),
+        VFS);
+    EXPECT_TRUE(static_cast<bool>(CI));
+    CI->getFrontendOpts().DisableFree = false;
+
+    // The diagnostic options must be set before creating a CompilerInstance.
+    CI->getDiagnosticOpts().IgnoreWarnings = true;
+    auto Clang = prepareCompilerInstance(
+        std::move(CI), /*Preamble=*/nullptr,
+        llvm::MemoryBuffer::getMemBuffer(FS.Files[MainFile], MainFile),
+        std::make_shared<PCHContainerOperations>(), VFS, IgnoreDiags);
+
+    EXPECT_FALSE(Clang->getFrontendOpts().Inputs.empty());
+    return Clang;
   }
 
 protected:
+  std::vector<Inclusion> collectIncludes() {
+    auto Clang = setupClang();
+    PreprocessOnlyAction Action;
+    EXPECT_TRUE(
+        Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
+    std::vector<Inclusion> Inclusions;
+    Clang->getPreprocessor().addPPCallbacks(collectInclusionsInMainFileCallback(
+        Clang->getSourceManager(),
+        [&](Inclusion Inc) { Inclusions.push_back(std::move(Inc)); }));
+    EXPECT_TRUE(Action.Execute());
+    Action.EndSourceFile();
+    return Inclusions;
+  }
+
   // Calculates the include path, or returns "" on error.
   std::string calculate(PathRef Original, PathRef Preferred = "",
+                        const std::vector<Inclusion> &Inclusions = {},
                         bool ExpectError = false) {
+    auto Clang = setupClang();
+    PreprocessOnlyAction Action;
+    EXPECT_TRUE(
+        Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
+
     if (Preferred.empty())
       Preferred = Original;
-    auto VFS = FS.getFileSystem();
-    auto Cmd = CDB.getCompileCommand(MainFile);
-    assert(static_cast<bool>(Cmd));
-    VFS->setCurrentWorkingDirectory(Cmd->Directory);
     auto ToHeaderFile = [](llvm::StringRef Header) {
       return HeaderFile{Header,
                         /*Verbatim=*/!llvm::sys::path::is_absolute(Header)};
     };
-    auto Path = calculateIncludePath(MainFile, FS.Files[MainFile],
-                                     ToHeaderFile(Original),
-                                     ToHeaderFile(Preferred), *Cmd, VFS);
+
+    auto Path = calculateIncludePath(
+        MainFile, CDB.getCompileCommand(MainFile)->Directory,
+        Clang->getPreprocessor().getHeaderSearchInfo(), Inclusions,
+        ToHeaderFile(Original), ToHeaderFile(Preferred));
+    Action.EndSourceFile();
     if (!Path) {
       llvm::consumeError(Path.takeError());
       EXPECT_TRUE(ExpectError);
       return std::string();
     } else {
       EXPECT_FALSE(ExpectError);
     }
     return std::move(*Path);
   }
+
+  Expected<Optional<TextEdit>>
+  insert(const HeaderFile &DeclaringHeader, const HeaderFile &InsertedHeader,
+         const std::vector<Inclusion> &ExistingInclusions = {}) {
+    auto Clang = setupClang();
+    PreprocessOnlyAction Action;
+    EXPECT_TRUE(
+        Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
+
+    IncludeInserter Inserter(MainFile, /*Code=*/"", format::getLLVMStyle(),
+                             CDB.getCompileCommand(MainFile)->Directory,
+                             Clang->getPreprocessor().getHeaderSearchInfo());
+    for (const auto &Inc : ExistingInclusions)
+      Inserter.addExisting(Inc);
+
+    auto Edit = Inserter.insert(DeclaringHeader, InsertedHeader);
+    Action.EndSourceFile();
+    return Edit;
+  }
+
   MockFSProvider FS;
   MockCompilationDatabase CDB;
   std::string MainFile = testPath("main.cpp");
   std::string Subdir = testPath("sub");
   std::string SearchDirArg = (llvm::Twine("-I") + Subdir).str();
+  IgnoringDiagConsumer IgnoreDiags;
 };
 
-TEST_F(HeadersTest, InsertInclude) {
-  std::string Path = testPath("sub/bar.h");
-  FS.Files[Path] = "";
-  EXPECT_EQ(calculate(Path), "\"bar.h\"");
-}
+MATCHER_P(Written, Name, "") { return arg.Written == Name; }
+MATCHER_P(Resolved, Name, "") { return arg.Resolved == Name; }
 
-TEST_F(HeadersTest, DontInsertDuplicateSameName) {
-  FS.Files[MainFile] = R"cpp(
-#include "bar.h"
-)cpp";
-  std::string BarHeader = testPath("sub/bar.h");
-  FS.Files[BarHeader] = "";
-  EXPECT_EQ(calculate(BarHeader), "");
-}
-
-TEST_F(HeadersTest, DontInsertDuplicateDifferentName) {
-  std::string BarHeader = testPath("sub/bar.h");
-  FS.Files[BarHeader] = "";
+TEST_F(HeadersTest, CollectRewrittenAndResolved) {
   FS.Files[MainFile] = R"cpp(
-#include "sub/bar.h"  // not shortest
+#include "sub/bar.h" // not shortest
 )cpp";
-  EXPECT_EQ(calculate("\"sub/bar.h\""), ""); // Duplicate rewritten.
-  EXPECT_EQ(calculate(BarHeader), "");       // Duplicate resolved.
-  EXPECT_EQ(calculate(BarHeader, "\"BAR.h\""), ""); // Do not insert preferred.
-}
-
-TEST_F(HeadersTest, DontInsertDuplicatePreferred) {
   std::string BarHeader = testPath("sub/bar.h");
   FS.Files[BarHeader] = "";
-  FS.Files[MainFile] = R"cpp(
-#include "sub/bar.h"  // not shortest
-)cpp";
-  // Duplicate written.
-  EXPECT_EQ(calculate("\"original.h\"", "\"sub/bar.h\""), "");
-  // Duplicate resolved.
-  EXPECT_EQ(calculate("\"original.h\"", BarHeader), "");
+
+  EXPECT_THAT(collectIncludes(),
+              UnorderedElementsAre(
+                  AllOf(Written("\"sub/bar.h\""), Resolved(BarHeader))));
 }
 
-TEST_F(HeadersTest, StillInsertIfTrasitivelyIncluded) {
+TEST_F(HeadersTest, OnlyCollectInclusionsInMain) {
   std::string BazHeader = testPath("sub/baz.h");
   FS.Files[BazHeader] = "";
   std::string BarHeader = testPath("sub/bar.h");
   FS.Files[BarHeader] = R"cpp(
 #include "baz.h"
 )cpp";
   FS.Files[MainFile] = R"cpp(
 #include "bar.h"
 )cpp";
-  EXPECT_EQ(calculate(BazHeader), "\"baz.h\"");
+  EXPECT_THAT(
+      collectIncludes(),
+      UnorderedElementsAre(AllOf(Written("\"bar.h\""), Resolved(BarHeader))));
+}
+
+TEST_F(HeadersTest, UnResolvedInclusion) {
+  FS.Files[MainFile] = R"cpp(
+#include "foo.h"
+)cpp";
+
+  EXPECT_THAT(collectIncludes(),
+              UnorderedElementsAre(AllOf(Written("\"foo.h\""), Resolved(""))));
+}
+
+TEST_F(HeadersTest, InsertInclude) {
+  std::string Path = testPath("sub/bar.h");
+  FS.Files[Path] = "";
+  EXPECT_EQ(calculate(Path), "\"bar.h\"");
 }
 
 TEST_F(HeadersTest, DoNotInsertIfInSameFile) {
   MainFile = testPath("main.h");
-  FS.Files[MainFile] = "";
   EXPECT_EQ(calculate(MainFile), "");
 }
 
+TEST_F(HeadersTest, ShortenedInclude) {
+  std::string BarHeader = testPath("sub/bar.h");
+  EXPECT_EQ(calculate(BarHeader), "\"bar.h\"");
+}
+
+TEST_F(HeadersTest, NotShortenedInclude) {
+  std::string BarHeader = testPath("sub-2/bar.h");
+  EXPECT_EQ(calculate(BarHeader, ""), "\"" + BarHeader + "\"");
+}
+
 TEST_F(HeadersTest, PreferredHeader) {
-  FS.Files[MainFile] = "";
   std::string BarHeader = testPath("sub/bar.h");
-  FS.Files[BarHeader] = "";
   EXPECT_EQ(calculate(BarHeader, "<bar>"), "<bar>");
 
   std::string BazHeader = testPath("sub/baz.h");
-  FS.Files[BazHeader] = "";
   EXPECT_EQ(calculate(BarHeader, BazHeader), "\"baz.h\"");
 }
 
+TEST_F(HeadersTest, DontInsertDuplicatePreferred) {
+  std::vector<Inclusion> Inclusions = {
+      {Range(), /*Written*/ "\"bar.h\"", /*Resolved*/ ""}};
+  EXPECT_EQ(calculate(testPath("sub/bar.h"), "\"bar.h\"", Inclusions), "");
+  EXPECT_EQ(calculate("\"x.h\"", "\"bar.h\"", Inclusions), "");
+}
+
+TEST_F(HeadersTest, DontInsertDuplicateResolved) {
+  std::string BarHeader = testPath("sub/bar.h");
+  std::vector<Inclusion> Inclusions = {
+      {Range(), /*Written*/ "fake-bar.h", /*Resolved*/ BarHeader}};
+  EXPECT_EQ(calculate(BarHeader, "", Inclusions), "");
+  // Do not insert preferred.
+  EXPECT_EQ(calculate(BarHeader, "\"BAR.h\"", Inclusions), "");
+}
+
+HeaderFile literal(StringRef Include) {
+  HeaderFile H{Include, /*Verbatim=*/true};
+  assert(H.valid());
+  return H;
+}
+
+TEST_F(HeadersTest, PreferInserted) {
+  auto Edit = insert(literal("<x>"), literal("<y>"));
+  EXPECT_TRUE(static_cast<bool>(Edit));
+  EXPECT_TRUE(llvm::StringRef((*Edit)->newText).contains("<y>"));
+}
+
+TEST_F(HeadersTest, ExistingInclusion) {
+  Inclusion Existing{Range(), /*Written=*/"<c.h>", /*Resolved=*/""};
+  auto Edit = insert(literal("<c.h>"), literal("<c.h>"), {Existing});
+  EXPECT_TRUE(static_cast<bool>(Edit));
+  EXPECT_FALSE(static_cast<bool>(*Edit));
+}
+
+TEST_F(HeadersTest, ValidateHeaders) {
+  HeaderFile InvalidHeader{"a.h", /*Verbatim=*/true};
+  assert(!InvalidHeader.valid());
+  auto Edit = insert(InvalidHeader, literal("\"c.h\""));
+  EXPECT_FALSE(static_cast<bool>(Edit));
+  llvm::consumeError(Edit.takeError());
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang
-