[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

571

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

2747–2750

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

405

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

321

again, i'd inline the marshalling

clang-tools-extra/clangd/ParsedAST.cpp

567

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–89

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–89

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

123

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–89

yes I meant exactly that signature.

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

622

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

[dgoldman]

Update comment + test

[kadircet]

clang-tools-extra/clangd/ParsedAST.cpp

573

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

573

Done, LMK if I did that properly.

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/AST.h

 //===--- AST.h - Utility AST functions  -------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Various code that examines C++ source code using AST.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_AST_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_AST_H
 
+#include "Headers.h"
+#include "index/Symbol.h"
 #include "index/SymbolID.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/TypeLoc.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Lex/MacroInfo.h"
 #include "llvm/ADT/StringRef.h"

 ///
 /// - class extension declaration: "@interface MyClass () ... @end"
 /// - which really links to class definition: "@implementation MyClass ... @end"
 ///
 /// For Objective-C protocols, e.g. "@protocol MyProtocol ... @end" this will
 /// return nullptr as protocols don't have an implementation.
 const ObjCImplDecl *getCorrespondingObjCImpl(const ObjCContainerDecl *D);
 
+/// 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.
+Symbol::IncludeDirective
+preferredIncludeDirective(llvm::StringRef FileName, const LangOptions &LangOpts,
+                          ArrayRef<Inclusion> MainFileIncludes,
+                          ArrayRef<const Decl *> TopLevelDecls);
+
 /// Returns a QualType as string. The result doesn't contain unwritten scopes
 /// like anonymous/inline namespace.
 std::string printType(const QualType QT, const DeclContext &CurContext,
                       llvm::StringRef Placeholder = "");
 
 /// Indicates if \p D is a template instantiation implicitly generated by the
 /// compiler, e.g.
 ///     template <class T> struct vector {};

clang-tools-extra/clangd/AST.cpp

if (CD->IsClassExtension()) {
         return ID->getImplementation();
       return nullptr;
     }
     return CD->getImplementation();
   }
   return nullptr;
 }
 
+Symbol::IncludeDirective
+preferredIncludeDirective(llvm::StringRef FileName, const LangOptions &LangOpts,
+                          ArrayRef<Inclusion> MainFileIncludes,
+                          ArrayRef<const 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 : 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 (const Decl *D : TopLevelDecls)
+    if (isa<ObjCContainerDecl, ObjCIvarDecl, ObjCMethodDecl, ObjCPropertyDecl>(
+            D))
+      return Symbol::IncludeDirective::Import;
+
+  return Symbol::IncludeDirective::Include;
+}
+
 std::string printType(const QualType QT, const DeclContext &CurContext,
                       const llvm::StringRef Placeholder) {
   std::string Result;
   llvm::raw_string_ostream OS(Result);
   PrintingPolicy PP(CurContext.getParentASTContext().getPrintingPolicy());
   PP.SuppressTagKeyword = true;
   PP.SuppressUnwrittenScope = true;
 

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);
 };
 
 } // 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().MainFileIncludes, 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)

clang-tools-extra/clangd/ClangdServer.h

struct Options {
 
     FeatureModuleSet *FeatureModules = nullptr;
     /// If true, use the dirty buffer contents when building Preambles.
     bool UseDirtyHeaders = false;
 
     // If true, parse emplace-like functions in the preamble.
     bool PreambleParseForwardingFunctions = false;
 
+    /// Whether include fixer insertions for Objective-C code should use #import
+    /// instead of #include.
+    bool ImportInsertions = false;
+
     explicit operator TUScheduler::Options() const;
   };
   // Sensible default options for use in tests.
   // Features like indexing must be enabled if desired.
   static Options optsForTest();
 
   /// Creates a new ClangdServer instance.
   ///
private:
 
   bool UseDirtyHeaders = false;
 
   // Whether the client supports folding only complete lines.
   bool LineFoldingOnly = false;
 
   bool PreambleParseForwardingFunctions = false;
 
+  bool ImportInsertions = false;
+
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<std::optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   std::optional<std::string> WorkspaceRoot;
   std::optional<AsyncTaskRunner> IndexTasks; // for stdlib indexing.
   std::optional<TUScheduler> WorkScheduler;

clang-tools-extra/clangd/ClangdServer.cpp

ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const ThreadsafeFS &TFS, const Options &Opts,
                            Callbacks *Callbacks)
     : FeatureModules(Opts.FeatureModules), CDB(CDB), TFS(TFS),
       DynamicIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
       ClangTidyProvider(Opts.ClangTidyProvider),
       UseDirtyHeaders(Opts.UseDirtyHeaders),
       LineFoldingOnly(Opts.LineFoldingOnly),
       PreambleParseForwardingFunctions(Opts.PreambleParseForwardingFunctions),
+      ImportInsertions(Opts.ImportInsertions),
       WorkspaceRoot(Opts.WorkspaceRoot),
       Transient(Opts.ImplicitCancellation ? TUScheduler::InvalidateOnUpdate
                                           : TUScheduler::NoInvalidation),
       DirtyFS(std::make_unique<DraftStoreFS>(TFS, DraftMgr)) {
   if (Opts.AsyncThreadsCount != 0)
     IndexTasks.emplace();
   // Pass a callback into `WorkScheduler` to extract symbols from a newly
   // parsed file and rebuild the file index synchronously each time an AST

 }
 
 void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
                                llvm::StringRef Version,
                                WantDiagnostics WantDiags, bool ForceRebuild) {
   std::string ActualVersion = DraftMgr.addDraft(File, Version, Contents);
   ParseOptions Opts;
   Opts.PreambleParseForwardingFunctions = PreambleParseForwardingFunctions;
+  Opts.ImportInsertions = ImportInsertions;
 
   // Compile command is set asynchronously during update, as it can be slow.
   ParseInputs Inputs;
   Inputs.TFS = &getHeaderFS();
   Inputs.Contents = std::string(Contents);
   Inputs.Version = std::move(ActualVersion);
   Inputs.ForceRebuild = ForceRebuild;
   Inputs.Opts = std::move(Opts);

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 std::nullopt;
   }
 
   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);
     // 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);
+              ToInclude->first, Directive == Symbol::Import
+                                    ? tooling::IncludeDirective::Import
+                                    : tooling::IncludeDirective::Include);
         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/Compiler.h

public:
 
   void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
                         const clang::Diagnostic &Info) override;
 };
 
 // Options to run clang e.g. when parsing AST.
 struct ParseOptions {
   bool PreambleParseForwardingFunctions = false;
+
+  bool ImportInsertions = false;
 };
 
 /// Information required to run clang, e.g. to parse AST or do code completion.
 struct ParseInputs {
   tooling::CompileCommand CompileCommand;
   const ThreadsafeFS *TFS;
   std::string Contents;
   // Version identifier for Contents, provided by the client and opaque to us.

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.
   std::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;
   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))
+                               Directive == Symbol::Import
+                                  ? tooling::IncludeDirective::Import
+                                  : tooling::IncludeDirective::Include))
             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<const Decl *> getLocalTopLevelDecls() const;
 
   const std::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());
+      ArrayRef<Inclusion> MainFileIncludes;
       if (Preamble) {
+        MainFileIncludes = Preamble->Includes.MainFileIncludes;
         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 =
+          Inputs.Opts.ImportInsertions
+              ? preferredIncludeDirective(Filename, Clang->getLangOpts(),
+                                          MainFileIncludes, {})
+              : Symbol::Include;
       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());
     }
   }
 

 const Preprocessor &ParsedAST::getPreprocessor() const {
   return Clang->getPreprocessor();
 }
 
 llvm::ArrayRef<Decl *> ParsedAST::getLocalTopLevelDecls() {
   return LocalTopLevelDecls;
 }
 
+llvm::ArrayRef<const 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();
   // FIXME(ibiryukov): we do not account for the dynamically allocated part of
   // Message and Fixes inside each diagnostic.
   std::size_t Total = clangd::getUsedBytes(LocalTopLevelDecls) +

clang-tools-extra/clangd/tool/ClangdMain.cpp

if (EnableConfig)
       Providers.push_back(provideClangdConfig());
     Providers.push_back(provideDefaultChecks());
     Providers.push_back(disableUnusableChecks());
     ClangTidyOptProvider = combine(std::move(Providers));
     Opts.ClangTidyProvider = ClangTidyOptProvider;
   }
   Opts.UseDirtyHeaders = UseDirtyHeaders;
   Opts.PreambleParseForwardingFunctions = PreambleParseForwardingFunctions;
+  Opts.ImportInsertions = ImportInsertions;
   Opts.QueryDriverGlobs = std::move(QueryDriverGlobs);
   Opts.TweakFilter = [&](const Tweak &T) {
     if (T.hidden() && !HiddenFeatures)
       return false;
     if (TweakList.getNumOccurrences())
       return llvm::is_contained(TweakList, T.id());
     return true;
   };

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);
 }
 } // namespace
 } // namespace clangd
 } // namespace clang

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

 //===-- ASTTests.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 "AST.h"
 
 #include "Annotations.h"
 #include "ParsedAST.h"
 #include "TestTU.h"
+#include "index/Symbol.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/Basic/AttrKinds.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
TEST(ClangdAST, HasReservedName) {
   EXPECT_FALSE(hasReservedScope(
       *findUnqualifiedDecl(AST, "error_code").getDeclContext()));
 
   EXPECT_FALSE(hasReservedName(findUnqualifiedDecl(AST, "secret")));
   EXPECT_TRUE(
       hasReservedScope(*findUnqualifiedDecl(AST, "secret").getDeclContext()));
 }
 
+TEST(ClangdAST, PreferredIncludeDirective) {
+  auto ComputePreferredDirective = [](TestTU &TU) {
+    auto AST = TU.build();
+    return preferredIncludeDirective(AST.tuPath(), AST.getLangOpts(),
+                                     AST.getIncludeStructure().MainFileIncludes,
+                                     AST.getLocalTopLevelDecls());
+  };
+  TestTU ObjCTU = TestTU::withCode(R"cpp(
+  int main() {}
+  )cpp");
+  ObjCTU.Filename = "TestTU.m";
+  EXPECT_EQ(ComputePreferredDirective(ObjCTU),
+            Symbol::IncludeDirective::Import);
+
+  TestTU HeaderTU = TestTU::withCode(R"cpp(
+  #import "TestTU.h"
+  )cpp");
+  HeaderTU.Filename = "TestTUHeader.h";
+  HeaderTU.ExtraArgs = {"-xobjective-c++-header"};
+  EXPECT_EQ(ComputePreferredDirective(HeaderTU),
+            Symbol::IncludeDirective::Import);
+
+  // ObjC language option is not enough for headers.
+  HeaderTU.Code = R"cpp(
+  #include "TestTU.h"
+  )cpp";
+  EXPECT_EQ(ComputePreferredDirective(HeaderTU),
+            Symbol::IncludeDirective::Include);
+
+  HeaderTU.Code = R"cpp(
+  @interface Foo
+  @end
+
+  Foo * getFoo();
+  )cpp";
+  EXPECT_EQ(ComputePreferredDirective(HeaderTU),
+            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) {