[clangd] Full support for #import insertions

[dgoldman]

Summary

These are still disabled by default, but will work in ObjC code if you
enable the -import-insertions flag.

This requires ASTSignals to be available; before ASTSignals are
available, we will always use #include. Once they are available, the
behavior varies as follows:

• For source files, use #import if the ObjC language flag is enabled
• For header files:
  • If the ObjC language flag is disabled, use #include
  • If the header file contains any #imports, use #import
  • If the header file references any ObjC decls, use #import
  • Otherwise, use #include

[dgoldman]

Add tests

[kadircet]

clang-tools-extra/clangd/ASTSignals.cpp

29

can we rewrite this as:

bool preferImports(const ParsedAST &AST) {
  // If we don't have any objc-ness, always prefer #include
  if (!AST.getLangOpts().ObjC)
     return false;
  // If this is not a header file and has objc set as language, prefer #import
  if (!isHeaderFile(...))
     return true;
  // Headers lack proper compile flags most of the time, so we might treat a header as objc accidentally.
  // Perform some extra checks to make sure this works. 

  // Any file with a #import, should keep #import-ing.
  for (auto &Inc : AST.getIncludeStructure().MainFileIncludes) {
    if (Inc.Directive == tok::pp_import)
       return true;
  }
  
  // Any file declaring an objc struct should also #import.
  // No need to look over the references, as the file doesn't have any #imports, it must be declaring interesting Objc-like decls.
  for(Decl *D: AST.getLocalTopLevelDecls()) {
        if (isa<...InterestingDeclTypes...>(D)
           return true;
   });
  return false;
}

It's fine to run findExplicitReferences for an extra time here, as we do it only when this is an ObjC header with no #imports

clang-tools-extra/clangd/ASTSignals.h

33

can we rather say Preferred PP-directive to use for inclusions by the file.

clang-tools-extra/clangd/ParsedAST.cpp

570

we already have access to Clang->getLangOpts() here. we also have the FileName available so we can check for header-ness as well. we also have preamble includes to scan for #import directives (well, we might be building an AST without a preamble sometimes but it's fine if IncludeFixer can't suggest #imports in those cases).

so can we detect import-ness based on these here? we're only lacking information about whether the file declares any objc symbols, e.g. we might fail to suggest #import directives in a objc file without any #import statements, but I'd say it's fine considering all the extra complexity needed to propagate it from previous run. (if this turns out to be an important UX issue we can consider plumbing this information directly, so please leave a FIXME).

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

91

can you put this test first with a comment like objc lang option is not enough for headers?

97

this only works because we actually implicitly #import contents of HeaderTU.HeaderCode in TestTU.

can you instead clear up the HeaderCode and define the interface directly in the header file?

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

2729–2732

can you use this Opts (without Signals) in the previous test as well to make sure we're #includeing the symbol in a non-import context even if the option is set?

[dgoldman]

Fixes for review

[dgoldman]

Don't pass ASTSignals through ParsedAST

[kadircet]

clang-tools-extra/clangd/ASTSignals.h

39

could you rather move this to AST.h, with some comments explaining what it does

41

what about accepting an llvm::ArrayRef<Inclusion> MainFileIncludes instead of the full IncludeStructure, as the former can be constructed a lot more easily.

clang-tools-extra/clangd/CodeComplete.cpp

395

nit: I'd actually inline this into use-side to not confuse the reader unnecessarily about whether this is marshalling/conversion

clang-tools-extra/clangd/IncludeFixer.cpp

322

again, i'd inline the marshalling

clang-tools-extra/clangd/ParsedAST.cpp

566

this is a somewhat heavy struct to copy (contains absolute paths to all the includes in the preamble). can you use a pointer here?

clang-tools-extra/clangd/ParsedAST.h

88

can you rather introduce a new overload that returns ArrayRef<const Decl*>

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

87

since we have a dedicated interface now, could you please test it instead (same for uses below)?

98

no need for the #include here

[dgoldman]

Minor fixes for review, moving stuff around

[dgoldman]

clang-tools-extra/clangd/ParsedAST.h

88

Did you mean an overload with ArrayRef<const Decl *> getLocalTopLevelDecls() const;? ASTSignals gets a const ParsedAST &, unless you want me to const cast instead?

[kadircet]

clang-tools-extra/clangd/AST.h

122 ↗

(On Diff #486111)

Maybe something more concise like:

Infer the include directive to use for the given \p FileName. It aims for #import for ObjC files and #include for the rest.

For source files we use LangOpts directly to infer ObjC-ness.
For header files we also check for symbols declared by the file and existing include directives, as the language can be set to objc as a fallback in the absence of compile flags.

To make sure we're first talking about "what" it's trying to achieve and then talk about "how". That way future travelers can verify their changes to heuristics to see if they're breaking the semantics.

clang-tools-extra/clangd/ParsedAST.h

88

yes I meant exactly that signature.

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

622 ↗

(On Diff #486111)

nit: I'd make this a lambda in the test below.

[dgoldman]

Update comment + test

[kadircet]

clang-tools-extra/clangd/ParsedAST.cpp

572

as discussed offline can you guard this with the import-insertion flag as well? we should pass it through ParseOptions.

[dgoldman]

Respect ImportInsertions flag in IncludeFixer

[dgoldman]

clang-tools-extra/clangd/ParsedAST.cpp

572

Done, LMK if I did that properly.

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/ASTSignals.h

 //===--- ASTSignals.h --------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTSIGNALS_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTSIGNALS_H
 
 #include "ParsedAST.h"
+#include "index/Symbol.h"
 #include "index/SymbolID.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace clang {
 namespace clangd {
 
 /// Signals derived from a valid AST of a file.
 /// Provides information that can only be extracted from the AST to actions that
 /// can't access an AST. The signals are computed and updated asynchronously by
 /// the ASTWorker and thus they are always stale and also can be absent.
 /// Example usage: Information about the declarations used in a file affects
 /// code-completion ranking in that file.
 struct ASTSignals {
   /// Number of occurrences of each symbol present in the file.
   llvm::DenseMap<SymbolID, unsigned> ReferencedSymbols;
   /// Namespaces whose symbols are used in the file, and the number of such
   /// distinct symbols.
   llvm::StringMap<unsigned> RelatedNamespaces;
+  /// Preferred preprocessor directive to use for inclusions by the file.
+  Symbol::IncludeDirective InsertionDirective =
+      Symbol::IncludeDirective::Include;
 
   static ASTSignals derive(const ParsedAST &AST);
+
+  static Symbol::IncludeDirective preferredIncludeDirective(
+      llvm::StringRef Filename, const LangOptions &LangOpts,
+      const IncludeStructure &Includes, ArrayRef<Decl *> TopLevelDecls);
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTSIGNALS_H

clang-tools-extra/clangd/ASTSignals.cpp

 //===--- ASTSignals.cpp ------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ASTSignals.h"
 #include "AST.h"
 #include "FindTarget.h"
+#include "Headers.h"
 #include "support/Trace.h"
+#include "clang/AST/DeclObjC.h"
 
 namespace clang {
 namespace clangd {
 ASTSignals ASTSignals::derive(const ParsedAST &AST) {
   trace::Span Span("ASTSignals::derive");
   ASTSignals Signals;
+  Signals.InsertionDirective = preferredIncludeDirective(
+      AST.tuPath(), AST.getLangOpts(), AST.getIncludeStructure(),
+      AST.getLocalTopLevelDecls());
   const SourceManager &SM = AST.getSourceManager();
   findExplicitReferences(
       AST.getASTContext(),
       [&](ReferenceLoc Ref) {
         for (const NamedDecl *ND : Ref.Targets) {
           if (!isInsideMainFile(Ref.NameLoc, SM))
             continue;
           SymbolID ID = getSymbolID(ND);
           if (!ID)
             continue;
           unsigned &SymbolCount = Signals.ReferencedSymbols[ID];
           SymbolCount++;
           // Process namespace only when we see the symbol for the first time.
           if (SymbolCount != 1)
             continue;
           if (const auto *NSD = dyn_cast<NamespaceDecl>(ND->getDeclContext())) {
             if (NSD->isAnonymousNamespace())
               continue;
             std::string NS = printNamespaceScope(*NSD);
             if (!NS.empty())
               Signals.RelatedNamespaces[NS]++;
           }
         }
       },
       AST.getHeuristicResolver());
   return Signals;
 }
+
+Symbol::IncludeDirective ASTSignals::preferredIncludeDirective(
+    llvm::StringRef Filename, const LangOptions &LangOpts,
+    const IncludeStructure &Includes, ArrayRef<Decl *> TopLevelDecls) {
+  // Always prefer #include for non-ObjC code.
+  if (!LangOpts.ObjC)
+    return Symbol::IncludeDirective::Include;
+  // If this is not a header file and has ObjC set as the language, prefer
+  // #import.
+  if (!isHeaderFile(Filename, LangOpts))
+    return Symbol::IncludeDirective::Import;
+
+  // Headers lack proper compile flags most of the time, so we might treat a
+  // header as ObjC accidentally. Perform some extra checks to make sure this
+  // works.
+
+  // Any file with a #import, should keep #import-ing.
+  for (auto &Inc : Includes.MainFileIncludes)
+    if (Inc.Directive == tok::pp_import)
+      return Symbol::IncludeDirective::Import;
+
+  // Any file declaring an ObjC decl should also be #import-ing.
+  // No need to look over the references, as the file doesn't have any #imports,
+  // it must be declaring interesting ObjC-like decls.
+  for (Decl *D : TopLevelDecls)
+    if (isa<ObjCContainerDecl, ObjCIvarDecl, ObjCMethodDecl, ObjCPropertyDecl>(
+            D))
+      return Symbol::IncludeDirective::Import;
+
+  return Symbol::IncludeDirective::Include;
+}
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/CodeComplete.cpp

case MarkupKind::PlainText:
     break;
   case MarkupKind::Markdown:
     Result.value.append(Doc.asMarkdown());
     break;
   }
   return Result;
 }
 
+Symbol::IncludeDirective insertionDirective(const CodeCompleteOptions &Opts) {
+  if (!Opts.ImportInsertions || !Opts.MainFileSignals)
+    return Symbol::IncludeDirective::Include;
+  return Opts.MainFileSignals->InsertionDirective;
+}
+
 // Identifier code completion result.
 struct RawIdentifier {
   llvm::StringRef Name;
   unsigned References; // # of usages in file.
 };
 
 /// A code completion result, in clang-native form.
 /// It may be promoted to a CompletionItem if it's among the top-ranked results.
headerToInsertIfAllowed(const CodeCompleteOptions &Opts) const {
     if (SemaResult && SemaResult->Declaration) {
       // Avoid inserting new #include if the declaration is found in the current
       // file e.g. the symbol is forward declared.
       auto &SM = SemaResult->Declaration->getASTContext().getSourceManager();
       for (const Decl *RD : SemaResult->Declaration->redecls())
         if (SM.isInMainFile(SM.getExpansionLoc(RD->getBeginLoc())))
           return std::nullopt;
     }
+    Symbol::IncludeDirective Directive = insertionDirective(Opts);
     for (const auto &Inc : RankedIncludeHeaders)
-      // FIXME: We should support #import directives here.
-      if ((Inc.Directive & clang::clangd::Symbol::Include) != 0)
+      if ((Inc.Directive & Directive) != 0)
         return Inc.Header;
     return None;
   }
 
   using Bundle = llvm::SmallVector<CompletionCandidate, 4>;
 };
 using ScoredBundle =
     std::pair<CompletionCandidate::Bundle, CodeCompletion::Scores>;
auto Inserted = [&](llvm::StringRef Header)
       auto Spelled = Includes.calculateIncludePath(*ResolvedInserted, FileName);
       if (!Spelled)
         return error("Header not on include path");
       return std::make_pair(
           std::move(*Spelled),
           Includes.shouldInsertInclude(*ResolvedDeclaring, *ResolvedInserted));
     };
     bool ShouldInsert = C.headerToInsertIfAllowed(Opts).has_value();
+    Symbol::IncludeDirective Directive = insertionDirective(Opts);
+    tooling::IncludeDirective InsertDirective =
+        Directive == Symbol::Import ? tooling::IncludeDirective::Import
+                                    : tooling::IncludeDirective::Include;
     // Calculate include paths and edits for all possible headers.
     for (const auto &Inc : C.RankedIncludeHeaders) {
-      // FIXME: We should support #import directives here.
-      if ((Inc.Directive & clang::clangd::Symbol::Include) == 0)
+      if ((Inc.Directive & Directive) == 0)
         continue;
 
       if (auto ToInclude = Inserted(Inc.Header)) {
         CodeCompletion::IncludeCandidate Include;
         Include.Header = ToInclude->first;
         if (ToInclude->second && ShouldInsert)
-          Include.Insertion = Includes.insert(
-              ToInclude->first, tooling::IncludeDirective::Include);
+          Include.Insertion =
+              Includes.insert(ToInclude->first, InsertDirective);
         Completion.Includes.push_back(std::move(Include));
       } else
         log("Failed to generate include insertion edits for adding header "
             "(FileURI='{0}', IncludeHeader='{1}') into {2}: {3}",
             C.IndexResult->CanonicalDeclaration.FileURI, Inc.Header, FileName,
             ToInclude.takeError());
     }
     // Prefer includes that do not need edits (i.e. already exist).

clang-tools-extra/clangd/IncludeFixer.h

 namespace clangd {
 
 /// Attempts to recover from error diagnostics by suggesting include insertion
 /// fixes. For example, member access into incomplete type can be fixes by
 /// include headers with the definition.
 class IncludeFixer {
 public:
   IncludeFixer(llvm::StringRef File, std::shared_ptr<IncludeInserter> Inserter,
-               const SymbolIndex &Index, unsigned IndexRequestLimit)
+               const SymbolIndex &Index, unsigned IndexRequestLimit,
+               Symbol::IncludeDirective Directive)
       : File(File), Inserter(std::move(Inserter)), Index(Index),
-        IndexRequestLimit(IndexRequestLimit) {}
+        IndexRequestLimit(IndexRequestLimit), Directive(Directive) {}
 
   /// Returns include insertions that can potentially recover the diagnostic.
   /// If Info is a note and fixes are returned, they should *replace* the note.
   std::vector<Fix> fix(DiagnosticsEngine::Level DiagLevel,
                        const clang::Diagnostic &Info) const;
 
   /// Returns an ExternalSemaSource that records failed name lookups in Sema.
   /// This allows IncludeFixer to suggest inserting headers that define those
private:
   /// one we've seen during the Sema run.
   std::vector<Fix> fixUnresolvedName() const;
 
   std::string File;
   std::shared_ptr<IncludeInserter> Inserter;
   const SymbolIndex &Index;
   const unsigned IndexRequestLimit; // Make at most 5 index requests.
   mutable unsigned IndexRequestCount = 0;
+  const Symbol::IncludeDirective Directive;
 
   // These collect the last unresolved name so that we can associate it with the
   // diagnostic.
   llvm::Optional<UnresolvedName> LastUnresolvedName;
 
   // There can be multiple diagnostics that are caused by the same unresolved
   // name or incomplete type in one parse, especially when code is
   // copy-and-pasted without #includes. We cache the index results based on

clang-tools-extra/clangd/IncludeFixer.cpp

return std::make_pair(
         Inserter->shouldInsertInclude(*ResolvedDeclaring, *ResolvedInserted));
   };
 
   std::vector<Fix> Fixes;
   // Deduplicate fixes by include headers. This doesn't distinguish symbols in
   // different scopes from the same header, but this case should be rare and is
   // thus ignored.
   llvm::StringSet<> InsertedHeaders;
+  auto InsertDirective = Directive == Symbol::IncludeDirective::Import ?
+      tooling::IncludeDirective::Import : tooling::IncludeDirective::Include;
   for (const auto &Sym : Syms) {
     for (const auto &Inc : getRankedIncludes(Sym)) {
-      // FIXME: We should support #import directives here.
-      if ((Inc.Directive & clang::clangd::Symbol::Include) == 0)
+      if ((Inc.Directive & Directive) == 0)
         continue;
       if (auto ToInclude = Inserted(Sym, Inc.Header)) {
         if (ToInclude->second) {
           if (!InsertedHeaders.try_emplace(ToInclude->first).second)
             continue;
           if (auto Fix =
                   insertHeader(ToInclude->first, (Sym.Scope + Sym.Name).str(),
-                               tooling::IncludeDirective::Include))
+                               InsertDirective))
             Fixes.push_back(std::move(*Fix));
         }
       } else {
         vlog("Failed to calculate include insertion for {0} into {1}: {2}",
              Inc.Header, File, ToInclude.takeError());
       }
     }
   }

clang-tools-extra/clangd/ParsedAST.h

public:
 
   const LangOptions &getLangOpts() const {
     return getASTContext().getLangOpts();
   }
 
   /// This function returns top-level decls present in the main file of the AST.
   /// The result does not include the decls that come from the preamble.
   /// (These should be const, but RecursiveASTVisitor requires Decl*).
-  ArrayRef<Decl *> getLocalTopLevelDecls();
+  ArrayRef<Decl *> getLocalTopLevelDecls() const;
 
   const llvm::Optional<std::vector<Diag>> &getDiagnostics() const {
     return Diags;
   }
 
   /// Returns the estimated size of the AST and the accessory structures, in
   /// bytes. Does not include the size of the preamble.
   std::size_t getUsedBytes() const;

clang-tools-extra/clangd/ParsedAST.cpp

 //===--- ParsedAST.cpp -------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ParsedAST.h"
 #include "../clang-tidy/ClangTidyCheck.h"
 #include "../clang-tidy/ClangTidyDiagnosticConsumer.h"
 #include "../clang-tidy/ClangTidyModuleRegistry.h"
 #include "AST.h"
+#include "ASTSignals.h"
 #include "Compiler.h"
 #include "Config.h"
 #include "Diagnostics.h"
 #include "Feature.h"
 #include "FeatureModule.h"
 #include "Headers.h"
 #include "HeuristicResolver.h"
 #include "IncludeCleaner.h"
 #include "IncludeFixer.h"
 #include "Preamble.h"
 #include "SourceCode.h"
 #include "TidyProvider.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
+#include "index/Symbol.h"
 #include "support/Logger.h"
 #include "support/Trace.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticSema.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
if (PreserveDiags) {
     // Add IncludeFixer which can recover diagnostics caused by missing includes
     // (e.g. incomplete type) and attach include insertion fixes to diagnostics.
     if (Inputs.Index && !BuildDir.getError()) {
       auto Style =
           getFormatStyleForFile(Filename, Inputs.Contents, *Inputs.TFS);
       auto Inserter = std::make_shared<IncludeInserter>(
           Filename, Inputs.Contents, Style, BuildDir.get(),
           &Clang->getPreprocessor().getHeaderSearchInfo());
+      IncludeStructure Includes;
       if (Preamble) {
+        Includes = Preamble->Includes;
         for (const auto &Inc : Preamble->Includes.MainFileIncludes)
           Inserter->addExisting(Inc);
       }
+      // FIXME: Consider piping through ASTSignals to fetch this to handle the
+      // case where a header file contains ObjC decls but no #imports.
+      Symbol::IncludeDirective Directive =
+          ASTSignals::preferredIncludeDirective(Filename, Clang->getLangOpts(),
+                                                Includes, {});
       FixIncludes.emplace(Filename, Inserter, *Inputs.Index,
-                          /*IndexRequestLimit=*/5);
+                          /*IndexRequestLimit=*/5, Directive);
       ASTDiags.contributeFixes([&FixIncludes](DiagnosticsEngine::Level DiagLevl,
                                               const clang::Diagnostic &Info) {
         return FixIncludes->fix(DiagLevl, Info);
       });
       Clang->setExternalSemaSource(FixIncludes->unresolvedNameRecorder());
     }
   }
 

 std::shared_ptr<Preprocessor> ParsedAST::getPreprocessorPtr() {
   return Clang->getPreprocessorPtr();
 }
 
 const Preprocessor &ParsedAST::getPreprocessor() const {
   return Clang->getPreprocessor();
 }
 
-llvm::ArrayRef<Decl *> ParsedAST::getLocalTopLevelDecls() {
+llvm::ArrayRef<Decl *> ParsedAST::getLocalTopLevelDecls() const {
   return LocalTopLevelDecls;
 }
 
 const MainFileMacros &ParsedAST::getMacros() const { return Macros; }
 const std::vector<PragmaMark> &ParsedAST::getMarks() const { return Marks; }
 
 std::size_t ParsedAST::getUsedBytes() const {
   auto &AST = getASTContext();

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

EXPECT_THAT(
       UnorderedElementsAre(
           Pair(ns("tar").ID, 4), Pair(ns("ns1").ID, 1),
           Pair(ns("ns1::ns2").ID, 1), Pair(_ /*int func();*/, 1),
           Pair(cls("tar::X").ID, 1), Pair(var("tar::kConst").ID, 1),
           Pair(func("tar::foo").ID, 2), Pair(func("ns1::ns2::fooInNS2").ID, 1),
           Pair(sym("Y", index::SymbolKind::Variable, "@N@tar@S@X@FI@\\0").ID,
                2),
           Pair(_ /*a*/, 3)));
+  EXPECT_EQ(Signals.InsertionDirective, Symbol::IncludeDirective::Include);
 }
+
+TEST(ASTSignals, DeriveInsertionDirective) {
+  TestTU ObjCTU = TestTU::withCode(R"cpp(
+  int main() {}
+  )cpp");
+  ObjCTU.Filename = "TestTU.m";
+  ASTSignals Signals = ASTSignals::derive(ObjCTU.build());
+  EXPECT_EQ(Signals.InsertionDirective, Symbol::IncludeDirective::Import);
+
+  TestTU HeaderTU = TestTU::withCode(R"cpp(
+  #import "TestTU.h"
+  )cpp");
+  HeaderTU.Filename = "TestTUHeader.h";
+  HeaderTU.ExtraArgs = {"-xobjective-c++-header"};
+  Signals = ASTSignals::derive(HeaderTU.build());
+  EXPECT_EQ(Signals.InsertionDirective, Symbol::IncludeDirective::Import);
+
+  // ObjC language option is not enough for headers.
+  HeaderTU.Code = R"cpp(
+  #include "TestTU.h"
+  )cpp";
+  Signals = ASTSignals::derive(HeaderTU.build());
+  EXPECT_EQ(Signals.InsertionDirective, Symbol::IncludeDirective::Include);
+
+  HeaderTU.Code = R"cpp(
+  #include "TestTU.h"
+
+  @interface Foo
+  @end
+
+  Foo * getFoo();
+  )cpp";
+  Signals = ASTSignals::derive(HeaderTU.build());
+  EXPECT_EQ(Signals.InsertionDirective, Symbol::IncludeDirective::Import);
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

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

 }
 
 TEST(CompletionTest, InsertIncludeOrImport) {
   std::string DeclFile = URI::create(testPath("foo")).toString();
   Symbol Sym = func("Func");
   Sym.CanonicalDeclaration.FileURI = DeclFile.c_str();
   Sym.IncludeHeaders.emplace_back("\"bar.h\"", 1000,
                                   Symbol::Include | Symbol::Import);
-
-  auto Results = completions("Fun^", {Sym}).Completions;
+  CodeCompleteOptions Opts;
+  // Should only take effect in import contexts.
+  Opts.ImportInsertions = true;
+  auto Results = completions("Fun^", {Sym}, Opts).Completions;
   assert(!Results.empty());
   EXPECT_THAT(Results[0],
               AllOf(named("Func"), insertIncludeText("#include \"bar.h\"\n")));
 
-  Results = completions("Fun^", {Sym}, {}, "Foo.m").Completions;
+  ASTSignals Signals;
+  Signals.InsertionDirective = Symbol::IncludeDirective::Import;
+  Opts.MainFileSignals = &Signals;
+  Results = completions("Fun^", {Sym}, Opts, "Foo.m").Completions;
   assert(!Results.empty());
-  // TODO: Once #import integration support is done this should be #import.
   EXPECT_THAT(Results[0],
-              AllOf(named("Func"), insertIncludeText("#include \"bar.h\"\n")));
+              AllOf(named("Func"), insertIncludeText("#import \"bar.h\"\n")));
 
   Sym.IncludeHeaders[0].SupportedDirectives = Symbol::Import;
   Results = completions("Fun^", {Sym}).Completions;
   assert(!Results.empty());
   EXPECT_THAT(Results[0], AllOf(named("Func"), Not(insertInclude())));
 }
 
 TEST(CompletionTest, NoInsertIncludeIfOnePresent) {