[clang][Tooling] Add support for generating #import edits

And make use of this from clangd's CodeComplete and IncludeFixer, although currently they are both restricted only to #include symbols.

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

clang-tools-extra/clangd/CodeComplete.cpp

@@ -395,7 +395,8 @@ struct CodeCompletionBuilder {
         CodeCompletion::IncludeCandidate Include;
         Include.Header = ToInclude->first;
         if (ToInclude->second && ShouldInsert)
-          Include.Insertion = Includes.insert(ToInclude->first);
+          Include.Insertion = Includes.insert(
+              ToInclude->first, tooling::IncludeDirective::Include);
         Completion.Includes.push_back(std::move(Include));
       } else
         log("Failed to generate include insertion edits for adding header "

clang-tools-extra/clangd/Headers.cpp

@@ -345,10 +345,12 @@ IncludeInserter::calculateIncludePath(const HeaderFile &InsertedHeader,
 }
 
 llvm::Optional<TextEdit>
-IncludeInserter::insert(llvm::StringRef VerbatimHeader) const {
+IncludeInserter::insert(llvm::StringRef VerbatimHeader,
+                        tooling::IncludeDirective Directive) const {
   llvm::Optional<TextEdit> Edit;
-  if (auto Insertion = Inserter.insert(VerbatimHeader.trim("\"<>"),
-                                       VerbatimHeader.startswith("<")))
+  if (auto Insertion =
+          Inserter.insert(VerbatimHeader.trim("\"<>"),
+                          VerbatimHeader.startswith("<"), Directive))
     Edit = replacementToEdit(Code, *Insertion);
   return Edit;
 }

clang-tools-extra/clangd/Headers.h

@@ -50,7 +50,7 @@ struct SymbolInclude {
   /// The header to include. This is either a URI or a verbatim include which is
   /// quoted with <> or "".
   llvm::StringRef Header;
-  /// The include directive to use, e.g. #import or #include.
+  /// The include directive(s) that can be used, e.g. #import and/or #include.
   Symbol::IncludeDirective Directive;
 };
 
@@ -248,7 +248,8 @@ class IncludeInserter {
 
   /// Calculates an edit that inserts \p VerbatimHeader into code. If the header
   /// is already included, this returns std::nullopt.
-  llvm::Optional<TextEdit> insert(llvm::StringRef VerbatimHeader) const;
+  llvm::Optional<TextEdit> insert(llvm::StringRef VerbatimHeader,
+                                  tooling::IncludeDirective Directive) const;
 
 private:
   StringRef FileName;

clang-tools-extra/clangd/IncludeFixer.cpp

@@ -249,18 +249,22 @@ std::vector<Fix> IncludeFixer::fix(DiagnosticsEngine::Level DiagLevel,
 }
 
 llvm::Optional<Fix> IncludeFixer::insertHeader(llvm::StringRef Spelled,
-                                               llvm::StringRef Symbol) const {
+                                               llvm::StringRef Symbol,
+                                               tooling::IncludeDirective Directive) const {
   Fix F;
 
-  if (auto Edit = Inserter->insert(Spelled))
+  if (auto Edit = Inserter->insert(Spelled, Directive))
     F.Edits.push_back(std::move(*Edit));
   else
     return std::nullopt;
 
+  llvm::StringRef DirectiveSpelling =
+      Directive == tooling::IncludeDirective::Include ? "Include" : "Import";
   if (Symbol.empty())
-    F.Message = llvm::formatv("Include {0}", Spelled);
+    F.Message = llvm::formatv("{0} {1}", DirectiveSpelling, Spelled);
   else
-    F.Message = llvm::formatv("Include {0} for symbol {1}", Spelled, Symbol);
+    F.Message = llvm::formatv("{0} {1} for symbol {2}",
+        DirectiveSpelling, Spelled, Symbol);
 
   return F;
 }
@@ -325,7 +329,8 @@ std::vector<Fix> IncludeFixer::fixesForSymbols(const SymbolSlab &Syms) const {
           if (!InsertedHeaders.try_emplace(ToInclude->first).second)
             continue;
           if (auto Fix =
-                  insertHeader(ToInclude->first, (Sym.Scope + Sym.Name).str()))
+                  insertHeader(ToInclude->first, (Sym.Scope + Sym.Name).str(),
+                               tooling::IncludeDirective::Include))
             Fixes.push_back(std::move(*Fix));
         }
       } else {

clang-tools-extra/clangd/IncludeFixer.h

@@ -17,6 +17,7 @@
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Sema/ExternalSemaSource.h"
+#include "clang/Tooling/Inclusions/HeaderIncludes.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
@@ -54,8 +55,10 @@ class IncludeFixer {
   /// Generates header insertion fixes for all symbols. Fixes are deduplicated.
   std::vector<Fix> fixesForSymbols(const SymbolSlab &Syms) const;
 
-  llvm::Optional<Fix> insertHeader(llvm::StringRef Name,
-                                   llvm::StringRef Symbol = "") const;
+  llvm::Optional<Fix>
+  insertHeader(llvm::StringRef Name, llvm::StringRef Symbol = "",
+               tooling::IncludeDirective Directive =
+                   tooling::IncludeDirective::Include) const;
 
   struct UnresolvedName {
     std::string Name;   // E.g. "X" in foo::X.

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

@@ -16,6 +16,7 @@
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
+#include "clang/Tooling/Inclusions/HeaderIncludes.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -117,7 +118,8 @@ class HeadersTest : public ::testing::Test {
     return Path.value_or("");
   }
 
-  llvm::Optional<TextEdit> insert(llvm::StringRef VerbatimHeader) {
+  llvm::Optional<TextEdit> insert(llvm::StringRef VerbatimHeader,
+                                  tooling::IncludeDirective Directive) {
     Clang = setupClang();
     PreprocessOnlyAction Action;
     EXPECT_TRUE(
@@ -126,7 +128,7 @@ class HeadersTest : public ::testing::Test {
     IncludeInserter Inserter(MainFile, /*Code=*/"", format::getLLVMStyle(),
                              CDB.getCompileCommand(MainFile)->Directory,
                              &Clang->getPreprocessor().getHeaderSearchInfo());
-    auto Edit = Inserter.insert(VerbatimHeader);
+    auto Edit = Inserter.insert(VerbatimHeader, Directive);
     Action.EndSourceFile();
     return Edit;
   }
@@ -330,9 +332,13 @@ TEST_F(HeadersTest, DontInsertDuplicateResolved) {
 }
 
 TEST_F(HeadersTest, PreferInserted) {
-  auto Edit = insert("<y>");
-  EXPECT_TRUE(Edit);
-  EXPECT_TRUE(StringRef(Edit->newText).contains("<y>"));
+  auto Edit = insert("<y>", tooling::IncludeDirective::Include);
+  ASSERT_TRUE(Edit);
+  EXPECT_EQ(Edit->newText, "#include <y>\n");
+
+  Edit = insert("\"header.h\"", tooling::IncludeDirective::Import);
+  ASSERT_TRUE(Edit);
+  EXPECT_EQ(Edit->newText, "#import \"header.h\"\n");
 }
 
 TEST(Headers, NoHeaderSearchInfo) {

clang/include/clang/Tooling/Inclusions/HeaderIncludes.h

@@ -44,16 +44,18 @@ class IncludeCategoryManager {
   SmallVector<llvm::Regex, 4> CategoryRegexs;
 };
 
+enum class IncludeDirective { Include, Import };
+
 /// Generates replacements for inserting or deleting #include directives in a
 /// file.
 class HeaderIncludes {
 public:
   HeaderIncludes(llvm::StringRef FileName, llvm::StringRef Code,
                  const IncludeStyle &Style);
 
-  /// Inserts an #include directive of \p Header into the code. If \p IsAngled
-  /// is true, \p Header will be quoted with <> in the directive; otherwise, it
-  /// will be quoted with "".
+  /// Inserts an #include or #import directive of \p Header into the code.
+  /// If \p IsAngled is true, \p Header will be quoted with <> in the directive;
+  /// otherwise, it will be quoted with "".
   ///
   /// When searching for points to insert new header, this ignores #include's
   /// after the #include block(s) in the beginning of a file to avoid inserting
@@ -71,26 +73,30 @@ class HeaderIncludes {
   /// \p IncludeName already exists (with exactly the same spelling), this
   /// returns std::nullopt.
   llvm::Optional<tooling::Replacement> insert(llvm::StringRef Header,
-                                              bool IsAngled) const;
+                                              bool IsAngled,
+                                              IncludeDirective Directive) const;
 
-  /// Removes all existing #includes of \p Header quoted with <> if \p IsAngled
-  /// is true or "" if \p IsAngled is false.
-  /// This doesn't resolve the header file path; it only deletes #includes with
-  /// exactly the same spelling.
+  /// Removes all existing #includes and #imports of \p Header quoted with <> if
+  /// \p IsAngled is true or "" if \p IsAngled is false.
+  /// This doesn't resolve the header file path; it only deletes #includes and
+  /// #imports with exactly the same spelling.
   tooling::Replacements remove(llvm::StringRef Header, bool IsAngled) const;
 
   // Matches a whole #include directive.
   static const llvm::Regex IncludeRegex;
 
 private:
   struct Include {
-    Include(StringRef Name, tooling::Range R) : Name(Name), R(R) {}
+    Include(StringRef Name, tooling::Range R, IncludeDirective D)
+        : Name(Name), R(R), Directive(D) {}
 
     // An include header quoted with either <> or "".
     std::string Name;
     // The range of the whole line of include directive including any leading
     // whitespaces and trailing comment.
     tooling::Range R;
+    // Either #include or #import.
+    IncludeDirective Directive;
   };
 
   void addExistingInclude(Include IncludeToAdd, unsigned NextLineOffset);

clang/lib/Format/Format.cpp

@@ -3302,7 +3302,8 @@ fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
     (void)Matched;
     auto IncludeName = Matches[2];
     auto Replace =
-        Includes.insert(IncludeName.trim("\"<>"), IncludeName.startswith("<"));
+        Includes.insert(IncludeName.trim("\"<>"), IncludeName.startswith("<"),
+                        tooling::IncludeDirective::Include);
     if (Replace) {
       auto Err = Result.add(*Replace);
       if (Err) {

clang/lib/Tooling/Inclusions/HeaderIncludes.cpp

@@ -296,7 +296,9 @@ HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
       addExistingInclude(
           Include(Matches[2],
                   tooling::Range(
-                      Offset, std::min(Line.size() + 1, Code.size() - Offset))),
+                      Offset, std::min(Line.size() + 1, Code.size() - Offset)),
+                  Matches[1] == "import" ? tooling::IncludeDirective::Import
+                                         : tooling::IncludeDirective::Include),
           NextLineOffset);
     }
     Offset = NextLineOffset;
@@ -342,17 +344,20 @@ void HeaderIncludes::addExistingInclude(Include IncludeToAdd,
 }
 
 llvm::Optional<tooling::Replacement>
-HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled) const {
+HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled,
+                       IncludeDirective Directive) const {
   assert(IncludeName == trimInclude(IncludeName));
   // If a <header> ("header") already exists in code, "header" (<header>) with
-  // different quotation will still be inserted.
+  // different quotation and/or directive will still be inserted.
   // FIXME: figure out if this is the best behavior.
   auto It = ExistingIncludes.find(IncludeName);
-  if (It != ExistingIncludes.end())
+  if (It != ExistingIncludes.end()) {
     for (const auto &Inc : It->second)
-      if ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
-          (!IsAngled && StringRef(Inc.Name).startswith("\"")))
+      if (Inc.Directive == Directive &&
+          ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
+           (!IsAngled && StringRef(Inc.Name).startswith("\""))))
         return std::nullopt;
+  }
   std::string Quoted =
       std::string(llvm::formatv(IsAngled ? "<{0}>" : "\"{0}\"", IncludeName));
   StringRef QuotedName = Quoted;
@@ -371,8 +376,10 @@ HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled) const {
     }
   }
   assert(InsertOffset <= Code.size());
+  llvm::StringRef DirectiveSpelling =
+      Directive == IncludeDirective::Include ? "include" : "import";
   std::string NewInclude =
-      std::string(llvm::formatv("#include {0}\n", QuotedName));
+      llvm::formatv("#{0} {1}\n", DirectiveSpelling, QuotedName);
   // When inserting headers at end of the code, also append '\n' to the code
   // if it does not end with '\n'.
   // FIXME: when inserting multiple #includes at the end of code, only one

clang/unittests/Tooling/HeaderIncludesTest.cpp

@@ -20,10 +20,12 @@ namespace {
 
 class HeaderIncludesTest : public ::testing::Test {
 protected:
-  std::string insert(llvm::StringRef Code, llvm::StringRef Header) {
+  std::string insert(llvm::StringRef Code, llvm::StringRef Header,
+                     IncludeDirective Directive = IncludeDirective::Include) {
     HeaderIncludes Includes(FileName, Code, Style);
     assert(Header.startswith("\"") || Header.startswith("<"));
-    auto R = Includes.insert(Header.trim("\"<>"), Header.startswith("<"));
+    auto R =
+        Includes.insert(Header.trim("\"<>"), Header.startswith("<"), Directive);
     if (!R)
       return std::string(Code);
     auto Result = applyAllReplacements(Code, Replacements(*R));
@@ -60,6 +62,25 @@ TEST_F(HeaderIncludesTest, RepeatedIncludes) {
   EXPECT_EQ(Expected, insert(Code, "\"a2.h\""));
 }
 
+TEST_F(HeaderIncludesTest, InsertImportWithSameInclude) {
+  std::string Code = "#include \"a.h\"\n";
+  std::string Expected = Code + "#import \"a.h\"\n";
+  EXPECT_EQ(Expected, insert(Code, "\"a.h\"", IncludeDirective::Import));
+}
+
+TEST_F(HeaderIncludesTest, DontInsertAlreadyImported) {
+  std::string Code = "#import \"a.h\"\n";
+  EXPECT_EQ(Code, insert(Code, "\"a.h\"", IncludeDirective::Import));
+}
+
+TEST_F(HeaderIncludesTest, DeleteImportAndSameInclude) {
+  std::string Code = R"cpp(
+#include <abc.h>
+#import <abc.h>
+int x;)cpp";
+  EXPECT_EQ("\nint x;", remove(Code, "<abc.h>"));
+}
+
 TEST_F(HeaderIncludesTest, NoExistingIncludeWithDefine) {
   std::string Code = "#ifndef A_H\n"
                      "#define A_H\n"