[clangd] don't insert new includes if either original header or canonical header is already included.

[ioeric]

Summary

Changes:
o Store both the original header and the canonical header in LSP command.
o Also check that both original and canonical headers are not already included
by comparing both resolved header path and written literal includes.

This addresses the use case where private IWYU pragma is defined in a private
header while it would still be preferrable to include the private header, in the
internal implementation file. If we have seen that the priviate header is already
included, we don't try to insert the canonical include.

[sammccall]

clangd/ClangdServer.h

239 ↗

(On Diff #135045)

I really do want to know what this parameter does, but I don't understand it from this comment.

243 ↗

(On Diff #135045)

"this has the same semantic as OriginalHeader" contradicts "[OriginalHeader] may be different from preferredHeader" :-)

[sammccall]

Mostly nits around naming/doc: you've convinced me that the two headers that are paths or uris or <headers> at different times is the right thing, but we need to be really clear about it.

clangd/ClangdServer.cpp

320 ↗

(On Diff #135045)

For readability, I'd suggest pulling this out into a real function, and returning something like
struct HeaderFile { std::string File; bool Verbatim; } which could be passed to the Headers.h functions.

The fold of "literal or uri" -> "literal or file" all typed as 'string' is a bit too much I think.
But maybe more explicit names and comments would moke this clearer.

clangd/ClangdServer.h

239 ↗

(On Diff #135045)

Discussed offline - this is the header file (either URI or <header>) that directly declares the symbol.
I'd suggest calling this DeclaringHeader.

240 ↗

(On Diff #135045)

the "this may be different" probably belongs after the doc of PreferredHeader (backreferences are easier to follow than forward ones)

243 ↗

(On Diff #135045)

T think InsertedHeader might be clearer about its role.

243 ↗

(On Diff #135045)

I'd suggest being explicit, since this is different to above: "if this is a URI, insertInclude will determine how to spell it. If this is a quoted string, it will be inserted verbatim".

244 ↗

(On Diff #135045)

Drop the "empty" special case, the caller should always pass this in.

246 ↗

(On Diff #135045)

Worth pointing out at the end "Both OriginalHeader and InsertedHeader will be considered to determine whether an include needs to be added".

clangd/Headers.cpp

72 ↗

(On Diff #135045)

how could this happen?

126 ↗

(On Diff #135045)

oops, missed this in the previous commit - this isn't an error condition, but a "return empty string" condition.

Can we add a test for this?

clangd/Headers.h

23 ↗

(On Diff #135045)

only used in headers.cpp, make static?

29 ↗

(On Diff #135045)

again, this special case isn't needed as copying a stringref is cheap

33 ↗

(On Diff #135045)

when would originalheader be empty?

37 ↗

(On Diff #135045)

again - not quite true.

clangd/Protocol.h

434 ↗

(On Diff #135045)

(also update comments here)

[ioeric]

• Addressed review comments.

[ioeric]

clangd/ClangdServer.h

243 ↗

(On Diff #135045)

Bad wording... sorry about that!

clangd/Headers.cpp

126 ↗

(On Diff #135045)

There are tests for this in HeadersTests.cpp. It's just that in the tests, we treat both errors and empty insertion as no insertion.

clangd/Headers.h

23 ↗

(On Diff #135045)

Oops, the same logic appeared in ClangdServer. I wanted to share this but forgot to...

[sammccall]

Thanks, and sorry for the delay!

clangd/Headers.cpp

45 ↗

(On Diff #135074)

nit: StringSet? std::set is pretty bad unless we need ordering

[ioeric]

• Merge with origin/master
• Use llvm::StringSet

clang-tools-extra/trunk/clangd/ClangdLSPServer.cpp

if (Params.command == ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND &&
     auto &FileURI = Params.includeInsertion->textDocument.uri;
     auto Code = Server.getDocument(FileURI.file());
     if (!Code)
       return replyError(ErrorCode::InvalidParams,
                         ("command " +
                          ExecuteCommandParams::CLANGD_INSERT_HEADER_INCLUDE +
                          " called on non-added file " + FileURI.file())
                             .str());
-    auto Replaces = Server.insertInclude(FileURI.file(), *Code,
-                                         Params.includeInsertion->header);
+    llvm::StringRef DeclaringHeader = Params.includeInsertion->declaringHeader;
+    if (DeclaringHeader.empty())
+      return replyError(
+          ErrorCode::InvalidParams,
+          "declaringHeader must be provided for include insertion.");
+    llvm::StringRef PreferredHeader = Params.includeInsertion->preferredHeader;
+    auto Replaces = Server.insertInclude(
+        FileURI.file(), *Code, DeclaringHeader,
+        PreferredHeader.empty() ? DeclaringHeader : PreferredHeader);
     if (!Replaces) {
       std::string ErrMsg =
           ("Failed to generate include insertion edits for adding " +
-           Params.includeInsertion->header + " into " + FileURI.file())
+           DeclaringHeader + " (" + PreferredHeader + ") into " +
+           FileURI.file())
               .str();
       log(ErrMsg + ":" + llvm::toString(Replaces.takeError()));
       replyError(ErrorCode::InternalError, ErrMsg);
       return;
     }
     auto Edits = replacementsToEdits(*Code, *Replaces);
     WorkspaceEdit WE;
     WE.changes = {{FileURI.uri(), Edits}};
 
-    reply("Inserted header " + Params.includeInsertion->header);
+    reply(("Inserted header " + DeclaringHeader + " (" + PreferredHeader + ")")
+              .str());
     ApplyEdit(std::move(WE));
   } else {
     // We should not get here because ExecuteCommandParams would not have
     // parsed in the first place and this handler should not be called. But if
     // more commands are added, this will be here has a safe guard.
     replyError(
         ErrorCode::InvalidParams,
         llvm::formatv("Unsupported command \"{0}\".", Params.command).str());

clang-tools-extra/trunk/clangd/ClangdServer.h

public:
   formatOnType(StringRef Code, PathRef File, Position Pos);
 
   /// Rename all occurrences of the symbol at the \p Pos in \p File to
   /// \p NewName.
   void rename(PathRef File, Position Pos, llvm::StringRef NewName,
               UniqueFunction<void(Expected<std::vector<tooling::Replacement>>)>
                   Callback);
 
-  /// Inserts a new #include of \p Header into \p File, if it's not present.
-  /// \p Header is either an URI that can be resolved to an #include path that
-  /// is suitable to be inserted or a literal string quoted with <> or "" that
-  /// can be #included directly.
+  /// Inserts a new #include into \p File, if it's not present in \p Code.
+  ///
+  /// \p DeclaringHeader The original header corresponding to this insertion
+  /// e.g. the header that declared a symbol. This can be either a URI or a
+  /// literal string quoted with <> or "" that can be #included directly.
+  /// \p InsertedHeader The preferred header to be inserted. This may be
+  /// different from \p DeclaringHeader as a header file can have a different
+  /// canonical include. This can be either a URI or a literal string quoted
+  /// with <> or "" that can be #included directly.
+  ///
+  /// Both OriginalHeader and InsertedHeader will be considered to determine
+  /// whether an include needs to be added.
   Expected<tooling::Replacements> insertInclude(PathRef File, StringRef Code,
-                                                StringRef Header);
+                                                StringRef DeclaringHeader,
+                                                StringRef InsertedHeader);
 
   /// Gets current document contents for \p File. Returns None if \p File is not
   /// currently tracked.
   /// FIXME(ibiryukov): This function is here to allow offset-to-Position
   /// conversions in outside code, maybe there's a way to get rid of it.
   llvm::Optional<std::string> getDocument(PathRef File);
 
   /// Only for testing purposes.

clang-tools-extra/trunk/clangd/ClangdServer.cpp

auto Action = [Pos](Path File, std::string NewName,
     return Callback(Replacements);
   };
 
   WorkScheduler.runWithAST(
       "Rename", File,
       Bind(Action, File.str(), NewName.str(), std::move(Callback)));
 }
 
-Expected<tooling::Replacements>
-ClangdServer::insertInclude(PathRef File, StringRef Code,
-                            llvm::StringRef Header) {
-  std::string ToInclude;
-  if (Header.startswith("<") || Header.startswith("\"")) {
-    ToInclude = Header;
-  } else {
-    auto U = URI::parse(Header);
-    if (!U)
-      return U.takeError();
-    auto Resolved = URI::resolve(*U, /*HintPath=*/File);
-    if (!Resolved)
-      return Resolved.takeError();
+/// 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 Resolved = URI::resolve(*U, HintPath);
+  return HeaderFile{std::move(*Resolved), /*Verbatim=*/false};
+};
 
-    tooling::CompileCommand CompileCommand =
-        CompileArgs.getCompileCommand(File);
-    auto Include =
-        calculateIncludePath(File, Code, *Resolved, CompileCommand,
-                             FSProvider.getTaggedFileSystem(File).Value);
-    if (!Include)
-      return Include.takeError();
-    if (Include->empty())
-      return tooling::Replacements();
-    ToInclude = std::move(*Include);
-  }
+Expected<tooling::Replacements>
+ClangdServer::insertInclude(PathRef File, StringRef Code,
+                            StringRef DeclaringHeader,
+                            StringRef InsertedHeader) {
+  assert(!DeclaringHeader.empty() && !InsertedHeader.empty());
+  std::string ToInclude;
+  auto ResolvedOrginal = toHeaderFile(DeclaringHeader, File);
+  if (!ResolvedOrginal)
+    return ResolvedOrginal.takeError();
+  auto ResolvedPreferred = toHeaderFile(InsertedHeader, File);
+  if (!ResolvedPreferred)
+    return ResolvedPreferred.takeError();
+  tooling::CompileCommand CompileCommand = CompileArgs.getCompileCommand(File);
+  auto Include = calculateIncludePath(
+      File, Code, *ResolvedOrginal, *ResolvedPreferred, CompileCommand,
 
   auto Style = format::getStyle("file", File, "llvm");
   if (!Style) {
     llvm::consumeError(Style.takeError());
     // FIXME(ioeric): needs more consistent style support in clangd server.
     Style = format::getLLVMStyle();
   }
   // Replacement with offset UINT_MAX and length 0 will be treated as include

clang-tools-extra/trunk/clangd/CodeComplete.cpp

if (IndexResult) {
           // LSP favors additionalTextEdits over command. But we are still using
           // command here because it would be expensive to calculate #include
           // insertion edits for all candidates, and the include insertion edit
           // is unlikely to conflict with the code completion edits.
           Command Cmd;
           // Command title is not added since this is not a user-facing command.
           Cmd.command = ExecuteCommandParams::CLANGD_INSERT_HEADER_INCLUDE;
           IncludeInsertion Insertion;
-          Insertion.header = D->IncludeHeader;
+          // Fallback to canonical header if declaration location is invalid.
+          Insertion.declaringHeader =
+              IndexResult->CanonicalDeclaration.FileURI.empty()
+                  ? D->IncludeHeader
+                  : IndexResult->CanonicalDeclaration.FileURI;
+          Insertion.preferredHeader = D->IncludeHeader;
           Insertion.textDocument.uri = URIForFile(FileName);
           Cmd.includeInsertion = std::move(Insertion);
           I.command = std::move(Cmd);
         }
       }
     }
     I.scoreInfo = Scores;
     I.sortText = sortText(Scores.finalScore, Name);

clang-tools-extra/trunk/clangd/Headers.h

 #include "clang/Basic/VirtualFileSystem.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringRef.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;
+};
+
 /// 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 Header is an absolute file path.
-/// \return A quoted "path" or <path>. This returns an empty string if:
-///   - \p Header is already (directly) included in the file (including those
-///   included via different paths).
-///   - \p Header is the same as \p 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).
+///   - \p DeclaringHeader or \p InsertedHeader is the same as \p File.
 llvm::Expected<std::string>
-calculateIncludePath(PathRef File, llvm::StringRef Code, llvm::StringRef Header,
+calculateIncludePath(PathRef File, llvm::StringRef Code,
+                     const HeaderFile &DeclaringHeader,
+                     const HeaderFile &InsertedHeader,
                      const tooling::CompileCommand &CompileCommand,
                      IntrusiveRefCntPtr<vfs::FileSystem> FS);
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H

clang-tools-extra/trunk/clangd/Headers.cpp

 #include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 class RecordHeaders : public PPCallbacks {
 public:
-  RecordHeaders(std::set<std::string> &Headers) : Headers(Headers) {}
+  RecordHeaders(llvm::StringSet<> &WrittenHeaders,
+                llvm::StringSet<> &ResolvedHeaders)
+      : WrittenHeaders(WrittenHeaders), ResolvedHeaders(ResolvedHeaders) {}
 
   void InclusionDirective(SourceLocation /*HashLoc*/,
                           const Token & /*IncludeTok*/,
-                          llvm::StringRef /*FileName*/, bool /*IsAngled*/,
+                          llvm::StringRef FileName, bool IsAngled,
                           CharSourceRange /*FilenameRange*/,
                           const FileEntry *File, llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
                           const Module * /*Imported*/) override {
+    WrittenHeaders.insert(
+        (IsAngled ? "<" + FileName + ">" : "\"" + FileName + "\"").str());
     if (File != nullptr && !File->tryGetRealPathName().empty())
-      Headers.insert(File->tryGetRealPathName());
+      ResolvedHeaders.insert(File->tryGetRealPathName());
   }
 
 private:
-  std::set<std::string> &Headers;
+  llvm::StringSet<> &WrittenHeaders;
+  llvm::StringSet<> &ResolvedHeaders;
 };
 
 } // namespace
 
+bool isLiteralInclude(llvm::StringRef Include) {
+  return Include.startswith("<") || Include.startswith("\"");
+}
+
+bool HeaderFile::valid() const {
+  return (Verbatim && isLiteralInclude(File)) ||
+         (!Verbatim && llvm::sys::path::is_absolute(File));
+}
+
 /// 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,
-                     llvm::StringRef Header,
+                     const HeaderFile &DeclaringHeader,
+                     const HeaderFile &InsertedHeader,
                      const tooling::CompileCommand &CompileCommand,
                      IntrusiveRefCntPtr<vfs::FileSystem> FS) {
-  assert(llvm::sys::path::is_absolute(File) &&
-         llvm::sys::path::is_absolute(Header));
-
-  if (File == Header)
+  assert(llvm::sys::path::is_absolute(File));
+  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)
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::set<std::string> ExistingHeaders;
+  llvm::StringSet<> WrittenHeaders;
+  llvm::StringSet<> ResolvedHeaders;
   Clang->getPreprocessor().addPPCallbacks(
-      llvm::make_unique<RecordHeaders>(ExistingHeaders));
+      llvm::make_unique<RecordHeaders>(WrittenHeaders, ResolvedHeaders));
   if (!Action.Execute())
     return llvm::make_error<llvm::StringError>(
         "Failed to execute preprocessor only action for file " + File,
         llvm::inconvertibleErrorCode());
-  if (ExistingHeaders.find(Header) != ExistingHeaders.end()) {
-    return llvm::make_error<llvm::StringError>(
-        Header + " is already included in " + File,
-        llvm::inconvertibleErrorCode());
-  }
+  auto Included = [&](llvm::StringRef Header) {
+    return WrittenHeaders.find(Header) != WrittenHeaders.end() ||
+           ResolvedHeaders.find(Header) != ResolvedHeaders.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(
-      Header, CompileCommand.Directory, &IsSystem);
+      InsertedHeader.File, CompileCommand.Directory, &IsSystem);
   if (IsSystem)
     Suggested = "<" + Suggested + ">";
   else
     Suggested = "\"" + Suggested + "\"";
 
-  log("Suggested #include for " + Header + " is: " + Suggested);
+  log("Suggested #include for " + InsertedHeader.File + " is: " + Suggested);
   return Suggested;
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/clangd/Protocol.h

struct WorkspaceEdit {
   /// Note: "documentChanges" is not currently used because currently there is
   /// no support for versioned edits.
 };
 bool fromJSON(const json::Expr &, WorkspaceEdit &);
 json::Expr toJSON(const WorkspaceEdit &WE);
 
 struct IncludeInsertion {
   /// The document in which the command was invoked.
+  /// If either originalHeader or preferredHeader has been (directly) included
+  /// in the current file, no new include will be inserted.
   TextDocumentIdentifier textDocument;
 
-  /// The header to be inserted. This could be either a URI ir a literal string
-  /// quoted with <> or "" that can be #included directly.
-  std::string header;
+  /// The declaring header corresponding to this insertion e.g. the header that
+  /// declares a symbol. This could be either a URI or a literal string quoted
+  /// with <> or "" that can be #included directly.
+  std::string declaringHeader;
+  /// The preferred header to be inserted. This may be different from
+  /// originalHeader as a header file can have a different canonical include.
+  /// This could be either a URI or a literal string quoted with <> or "" that
+  /// can be #included directly. If empty, declaringHeader is used to calculate
+  /// the #include path.
+  std::string preferredHeader;
 };
 bool fromJSON(const json::Expr &, IncludeInsertion &);
 json::Expr toJSON(const IncludeInsertion &II);
 
 /// Exact commands are not specified in the protocol so we define the
 /// ones supported by Clangd here. The protocol specifies the command arguments
 /// to be "any[]" but to make this safer and more manageable, each command we
 /// handle maps to a certain llvm::Optional of some struct to contain its

clang-tools-extra/trunk/clangd/Protocol.cpp

json::Expr toJSON(const WorkspaceEdit &WE) {
   for (auto &Change : *WE.changes)
     FileChanges[Change.first] = json::ary(Change.second);
   return json::obj{{"changes", std::move(FileChanges)}};
 }
 
 bool fromJSON(const json::Expr &II, IncludeInsertion &R) {
   json::ObjectMapper O(II);
   return O && O.map("textDocument", R.textDocument) &&
-         O.map("header", R.header);
+         O.map("declaringHeader", R.declaringHeader) &&
+         O.map("preferredHeader", R.preferredHeader);
 }
 json::Expr toJSON(const IncludeInsertion &II) {
-  return json::obj{{"textDocument", II.textDocument}, {"header", II.header}};
+  return json::obj{{"textDocument", II.textDocument},
+                   {"declaringHeader", II.declaringHeader},
+                   {"preferredHeader", II.preferredHeader}};
 }
 
 json::Expr toJSON(const ApplyWorkspaceEditParams &Params) {
   return json::obj{{"edit", Params.edit}};
 }
 
 bool fromJSON(const json::Expr &Params, TextDocumentPositionParams &R) {
   json::ObjectMapper O(Params);

clang-tools-extra/trunk/test/clangd/insert-include.test

 # RUN: clangd -lit-test < %s | FileCheck -strict-whitespace %s
 # RUN: clangd -lit-test -pch-storage=memory < %s | FileCheck -strict-whitespace %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 ---
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///main.cpp","languageId":"cpp","version":1,"text":"void f() {}"}}}
 ---
-{"jsonrpc":"2.0","id":3,"method":"workspace/executeCommand","params":{"arguments":[{"header":"<vector>","textDocument":{"uri":"test:///main.cpp"}}],"command":"clangd.insertInclude"}}
+{"jsonrpc":"2.0","id":3,"method":"workspace/executeCommand","params":{"arguments":[{"declaringHeader":"\"/usr/include/bits/vector\"", "preferredHeader":"<vector>","textDocument":{"uri":"test:///main.cpp"}}],"command":"clangd.insertInclude"}}
 #      CHECK:    "id": 3,
 # CHECK-NEXT:    "jsonrpc": "2.0",
-# CHECK-NEXT:    "result": "Inserted header <vector>"
+# CHECK-NEXT:    "result": "Inserted header \"/usr/include/bits/vector\" (<vector>)"
 # CHECK-NEXT:  }
 #      CHECK:    "method": "workspace/applyEdit",
 # CHECK-NEXT:    "params": {
 # CHECK-NEXT:      "edit": {
 # CHECK-NEXT:        "changes": {
 # CHECK-NEXT:          "file:///clangd-test/main.cpp": [
 # CHECK-NEXT:            {
 # CHECK-NEXT:              "newText": "#include <vector>\n",

clang-tools-extra/trunk/unittests/clangd/ClangdTests.cpp

)cpp";
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
   ASSERT_EQ(DiagConsumer.Count, 2); // Sanity check - we actually ran both?
 }
 
 TEST_F(ClangdVFSTest, InsertIncludes) {
   MockFSProvider FS;
   ErrorCheckingDiagConsumer DiagConsumer;
   MockCompilationDatabase CDB;
+  std::string SearchDirArg = (llvm::Twine("-I") + testRoot()).str();
+  CDB.ExtraClangFlags.insert(CDB.ExtraClangFlags.end(), {SearchDirArg.c_str()});
   ClangdServer Server(CDB, DiagConsumer, FS,
                       /*AsyncThreadsCount=*/0,
                       /*StorePreamblesInMemory=*/true);
 
   // No need to sync reparses, because reparses are performed on the calling
   // thread.
   auto FooCpp = testPath("foo.cpp");
   const auto Code = R"cpp(
 #include "x.h"
 
 void f() {}
 )cpp";
   FS.Files[FooCpp] = Code;
   Server.addDocument(FooCpp, Code);
 
-  auto Inserted = [&](llvm::StringRef Header, llvm::StringRef Expected) {
-    auto Replaces = Server.insertInclude(FooCpp, Code, Header);
+  auto Inserted = [&](llvm::StringRef Original, llvm::StringRef Preferred,
+                      llvm::StringRef Expected) {
+    auto Replaces = Server.insertInclude(
+        FooCpp, Code, Original, Preferred.empty() ? Original : Preferred);
     EXPECT_TRUE(static_cast<bool>(Replaces));
     auto Changed = tooling::applyAllReplacements(Code, *Replaces);
     EXPECT_TRUE(static_cast<bool>(Changed));
     return llvm::StringRef(*Changed).contains(
         (llvm::Twine("#include ") + Expected + "").str());
   };
 
-  EXPECT_TRUE(Inserted("\"y.h\"", "\"y.h\""));
-  EXPECT_TRUE(Inserted("<string>", "<string>"));
+  EXPECT_TRUE(Inserted("\"y.h\"", /*Preferred=*/"","\"y.h\""));
+  EXPECT_TRUE(Inserted("\"y.h\"", /*Preferred=*/"\"Y.h\"", "\"Y.h\""));
+  EXPECT_TRUE(Inserted("<string>", /*Preferred=*/"", "<string>"));
+  EXPECT_TRUE(Inserted("<string>", /*Preferred=*/"", "<string>"));
+
+  std::string OriginalHeader = URI::createFile(testPath("y.h")).toString();
+  std::string PreferredHeader = URI::createFile(testPath("Y.h")).toString();
+  EXPECT_TRUE(Inserted(OriginalHeader,
+                       /*Preferred=*/"", "\"y.h\""));
+  EXPECT_TRUE(Inserted(OriginalHeader,
+                       /*Preferred=*/"<Y.h>", "<Y.h>"));
+  EXPECT_TRUE(Inserted(OriginalHeader, PreferredHeader, "\"Y.h\""));
+  EXPECT_TRUE(Inserted("<y.h>", PreferredHeader, "\"Y.h\""));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/unittests/clangd/HeadersTests.cpp

 class HeadersTest : public ::testing::Test {
 public:
   HeadersTest() {
     CDB.ExtraClangFlags = {SearchDirArg.c_str()};
     FS.Files[MainFile] = "";
   }
 
 protected:
-  // Calculates the include path for \p Header, or returns "" on error.
-  std::string calculate(PathRef Header) {
+  // Calculates the include path, or returns "" on error.
+  std::string calculate(PathRef Original, PathRef Preferred = "",
+                        bool ExpectError = false) {
+    if (Preferred.empty())
+      Preferred = Original;
     auto VFS = FS.getTaggedFileSystem(MainFile).Value;
     auto Cmd = CDB.getCompileCommand(MainFile);
     assert(static_cast<bool>(Cmd));
     VFS->setCurrentWorkingDirectory(Cmd->Directory);
-    auto Path =
-        calculateIncludePath(MainFile, FS.Files[MainFile], Header, *Cmd, VFS);
+    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);
     if (!Path) {
       llvm::consumeError(Path.takeError());
+      EXPECT_TRUE(ExpectError);
       return std::string();
+    } else {
+      EXPECT_FALSE(ExpectError);
     }
     return std::move(*Path);
   }
   MockFSProvider FS;
   MockCompilationDatabase CDB;
   std::string MainFile = testPath("main.cpp");
   std::string Subdir = testPath("sub");
   std::string SearchDirArg = (llvm::Twine("-I") + Subdir).str();

 }
 
 TEST_F(HeadersTest, DontInsertDuplicateDifferentName) {
   std::string BarHeader = testPath("sub/bar.h");
   FS.Files[BarHeader] = "";
   FS.Files[MainFile] = R"cpp(
 #include "sub/bar.h"  // not shortest
 )cpp";
-  EXPECT_EQ(calculate(BarHeader), "");
+  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), "");
 }
 
 TEST_F(HeadersTest, StillInsertIfTrasitivelyIncluded) {
   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\"");
 }
 
 TEST_F(HeadersTest, DoNotInsertIfInSameFile) {
   MainFile = testPath("main.h");
   FS.Files[MainFile] = "";
   EXPECT_EQ(calculate(MainFile), "");
 }
 
+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\"");
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang