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

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.

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: klimek, ilya-biryukov, jkorous-apple, cfe-commits

Differential Revision: https://reviews.llvm.org/D43510

llvm-svn: 326070

clang-tools-extra/clangd/ClangdLSPServer.cpp

@@ -194,12 +194,20 @@ void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
                          ExecuteCommandParams::CLANGD_INSERT_HEADER_INCLUDE +
                          " called on non-added file " + FileURI.file())
                             .str());
-    auto Replaces = Server.insertInclude(FileURI.file(), *Code,
+    llvm::StringRef DeclaringHeader = Params.includeInsertion->declaringHeader;
-                                         Params.includeInsertion->header);
+    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);
@@ -209,7 +217,8 @@ void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
     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

clang-tools-extra/clangd/ClangdServer.cpp

@@ -313,31 +313,42 @@ void ClangdServer::rename(
       Bind(Action, File.str(), NewName.str(), std::move(Callback)));
 }
 
+/// 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 Resolved = URI::resolve(*U, HintPath);
+  if (!Resolved)
+    return Resolved.takeError();
+  return HeaderFile{std::move(*Resolved), /*Verbatim=*/false};
+};
+
 Expected<tooling::Replacements>
 ClangdServer::insertInclude(PathRef File, StringRef Code,
-                            llvm::StringRef Header) {
+                            StringRef DeclaringHeader,
+                            StringRef InsertedHeader) {
+  assert(!DeclaringHeader.empty() && !InsertedHeader.empty());
   std::string ToInclude;
-  if (Header.startswith("<") || Header.startswith("\"")) {
+  auto ResolvedOrginal = toHeaderFile(DeclaringHeader, File);
-    ToInclude = Header;
+  if (!ResolvedOrginal)
-  } else {
+    return ResolvedOrginal.takeError();
-    auto U = URI::parse(Header);
+  auto ResolvedPreferred = toHeaderFile(InsertedHeader, File);
-    if (!U)
+  if (!ResolvedPreferred)
-      return U.takeError();
+    return ResolvedPreferred.takeError();
-    auto Resolved = URI::resolve(*U, /*HintPath=*/File);
+  tooling::CompileCommand CompileCommand = CompileArgs.getCompileCommand(File);
-    if (!Resolved)
+  auto Include = calculateIncludePath(
-      return Resolved.takeError();
+      File, Code, *ResolvedOrginal, *ResolvedPreferred, CompileCommand,
-
+      FSProvider.getTaggedFileSystem(File).Value);
-    tooling::CompileCommand CompileCommand =
+  if (!Include)
-        CompileArgs.getCompileCommand(File);
+    return Include.takeError();
-    auto Include =
+  if (Include->empty())
-        calculateIncludePath(File, Code, *Resolved, CompileCommand,
+    return tooling::Replacements();
-                             FSProvider.getTaggedFileSystem(File).Value);
+  ToInclude = std::move(*Include);
-    if (!Include)
-      return Include.takeError();
-    if (Include->empty())
-      return tooling::Replacements();
-    ToInclude = std::move(*Include);
-  }
 
   auto Style = format::getStyle("file", File, "llvm");
   if (!Style) {

clang-tools-extra/clangd/ClangdServer.h

@@ -241,12 +241,21 @@ public:
               UniqueFunction<void(Expected<std::vector<tooling::Replacement>>)>
                   Callback);
 
-  /// Inserts a new #include of \p Header into \p File, if it's not present.
+  /// Inserts a new #include into \p File, if it's not present in \p Code.
-  /// \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
+  /// \p DeclaringHeader The original header corresponding to this insertion
-  /// can be #included directly.
+  /// 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.

clang-tools-extra/clangd/CodeComplete.cpp

@@ -297,7 +297,12 @@ struct CompletionCandidate {
           // 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);

clang-tools-extra/clangd/Headers.cpp

@@ -24,36 +24,50 @@ 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) &&
+  assert(llvm::sys::path::is_absolute(File));
-         llvm::sys::path::is_absolute(Header));
+  assert(DeclaringHeader.valid() && InsertedHeader.valid());
-
+  if (File == DeclaringHeader.File || File == InsertedHeader.File)
-  if (File == Header)
     return "";
   FS->setCurrentWorkingDirectory(CompileCommand.Directory);
 
@@ -95,29 +109,35 @@ calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
     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()) {
+  auto Included = [&](llvm::StringRef Header) {
-    return llvm::make_error<llvm::StringError>(
+    return WrittenHeaders.find(Header) != WrittenHeaders.end() ||
-        Header + " is already included in " + File,
+           ResolvedHeaders.find(Header) != ResolvedHeaders.end();
-        llvm::inconvertibleErrorCode());
+  };
-  }
+  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;
 }
 

clang-tools-extra/clangd/Headers.h

@@ -19,18 +19,36 @@
 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.
+/// \param DeclaringHeader is the original header corresponding to \p
-/// \return A quoted "path" or <path>. This returns an empty string if:
+/// InsertedHeader e.g. the header that declares a symbol.
-///   - \p Header is already (directly) included in the file (including those
+/// \param InsertedHeader The preferred header to be inserted. This could be the
-///   included via different paths).
+/// same as DeclaringHeader but must be provided.
-///   - \p Header is the same as \p File.
+//  \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);
 

clang-tools-extra/clangd/Protocol.cpp

@@ -371,10 +371,13 @@ json::Expr toJSON(const WorkspaceEdit &WE) {
 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) {

clang-tools-extra/clangd/Protocol.h

@@ -453,11 +453,20 @@ 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
+  /// The declaring header corresponding to this insertion e.g. the header that
-  /// quoted with <> or "" that can be #included directly.
+  /// declares a symbol. This could be either a URI or a literal string quoted
-  std::string header;
+  /// 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);

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

@@ -4,10 +4,10 @@
 ---
 {"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": {

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

@@ -966,6 +966,8 @@ 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);
@@ -981,8 +983,10 @@ void f() {}
   FS.Files[FooCpp] = Code;
   Server.addDocument(FooCpp, Code);
 
-  auto Inserted = [&](llvm::StringRef Header, llvm::StringRef Expected) {
+  auto Inserted = [&](llvm::StringRef Original, llvm::StringRef Preferred,
-    auto Replaces = Server.insertInclude(FooCpp, Code, Header);
+                      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));
@@ -990,8 +994,19 @@ void f() {}
         (llvm::Twine("#include ") + Expected + "").str());
   };
 
-  EXPECT_TRUE(Inserted("\"y.h\"", "\"y.h\""));
+  EXPECT_TRUE(Inserted("\"y.h\"", /*Preferred=*/"","\"y.h\""));
-  EXPECT_TRUE(Inserted("<string>", "<string>"));
+  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

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

@@ -24,17 +24,28 @@ public:
   }
 
 protected:
-  // Calculates the include path for \p Header, or returns "" on error.
+  // Calculates the include path, or returns "" on error.
-  std::string calculate(PathRef Header) {
+  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 =
+    auto ToHeaderFile = [](llvm::StringRef Header) {
-        calculateIncludePath(MainFile, FS.Files[MainFile], Header, *Cmd, VFS);
+      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);
   }
@@ -66,7 +77,21 @@ TEST_F(HeadersTest, DontInsertDuplicateDifferentName) {
   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) {
@@ -88,6 +113,17 @@ TEST_F(HeadersTest, DoNotInsertIfInSameFile) {
   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