[clangd] Add new IncludeType to IncludeHeaderWithReferences

[dgoldman]

Summary

The IncludeType contains both Include (the current behavior) and Import,
which we can use in the future to provide #import suggestions for
Objective-C files/symbols.

[dgoldman]

Change IncludeTypes to be a bitset

[kadircet]

clang-tools-extra/clangd/index/Symbol.h

96

you can use LLVM_MARK_AS_BITMASK_ENUM and relevant magic to implement bitwise operators on the type (see llvm-project/llvm/include/llvm/ADT/BitmaskEnum.h)

119–126

i think we already rely on unsigned being 32 bits here. instead of introducing a new byte, can we make References 30 bits long and use 2 bits for the include type? also change the type from unsigned to uint32_t.

122

the naming here feels a little confusing, can we change the enum name to be IncludeDirective and field name to SupportedDirectives

clang-tools-extra/clangd/index/remote/Index.proto

110

no need for the comment here.

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

59

can we also have an example for both (and none)

[dgoldman]

Fixes for review

• Swap over IncludeType to IncludeDirective and update the protos accordingly

[dgoldman]

Also LMK if you want more in this change (such as a flag to control it, just not sure where it should live + what it should be called).

clang-tools-extra/clangd/index/Symbol.h

122

Done, although for Serialization I left it as a 32 bit and then 8 bit include directives, LMK if I should swap it over to serialize as a 32 bit single value... Guess I would need to use a union or manually bit manipulate it to store + load it?

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

1689–1690 ↗

(On Diff #473347)

Not sure the best way to fix this, should we also scan the -include files for #imports?

[dgoldman]

Fix serialization and isSelfContainedHeader

• Keep serialization as a single var uint32
• Keep old imported logic in isSelfContainedHeader in addition to a new check for import lines only for main files

[dgoldman]

Run clang format

[kadircet]

clang-tools-extra/clangd/SourceCode.cpp

1241 ↗

(On Diff #475180)

can you put a comment here saying any header that contains #imports are supposed to be #import'd so no need to check for anything but the main-file.

1249 ↗

(On Diff #475180)

nit: perform .take_front directly here.

clang-tools-extra/clangd/index/Merge.cpp

251

nit: SI.SupportedDirectives |= OI.SupportedDirectives;

clang-tools-extra/clangd/index/Symbol.h

90

nit: I'd still keep the enum name singular

clang-tools-extra/clangd/index/SymbolCollector.cpp

106

can we rather modify the previous function to look like:

IncludeDirective shouldCollectIncludePath(index::SymbolKind Kind) {
  using SK = index::SymbolKind;
  switch (Kind) {
  case SK::Macro:
  case SK::Enum:
  case SK::Struct:
  case SK::Class:
  case SK::Union:
  case SK::TypeAlias:
  case SK::Using:
  case SK::Function:
  case SK::Variable:
  case SK::EnumConstant:
  case SK::Concept:
    return Include | Import;
  case SK::Protocol:
    return Import;
  default:
    return Invalid;
  }
}

812

can we do this in addDeclaration instead? we already have nameLocation and FileID there, but that's also where we have the decl itself, which might be needed in future if we need more detailed checks.

860

can we instead do

auto It = map.find(key);
if (It == map.end()) {
  It = map.insert(key, calculateValue()).first;
}
bool ContainsImports = It->second;

870

can we perform ContainsImports calculations/caching here instead? that way we can trim a bunch of redundant searches.

clang-tools-extra/clangd/index/remote/marshalling/Marshalling.cpp

431–436

i guess IncludeTypes pieces are leftover?

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

2736

i think this case should still be empty, otherwise we're actually regressing the behaviour by inserting includes for symbols that we previously wouldn't insert.

[dgoldman]

Fixes for review

[kadircet]

clang-tools-extra/clangd/Headers.cpp

225

i don't think this is the right layer to perform such a filtering, we should instead be returning everything here, both the header and the directive to use.

later on inside CodeComplete.cpp, there's headerToInsertIfAllowed for now we can drop headers to be #import'd there, with a fixme saying propagation into other layers.
similarly IncludeFixer should drop these inside fixesForSymbols for now.

clang-tools-extra/clangd/index/SymbolCollector.cpp

809–810

nit: while here combine with the condition above

843–844

nit: feel free to rewrite as const auto &[SID, FID] : IncludeFiles, as you seem to be referring to Entry.second a bunch.

871

nit:

if ((CollectDirectives & Symbol::Import) != 0) {
     auto [It, Inserted] = FileToContainsImportsOrObjC.try_emplace(Entry.second);
     if(Inserted)
       It->second = FilesWithObjCConstructs.contains(Entry.second) || fileContainsImports(Entry.second,ASTCtx->getSourceManager());
     if(It->second) Directives |= Symbol::Import;
}

878

i'd first do this check, as it doesn't require parsing file contents

964

let's drop the Opts.CollectIncludePath check. it doesn't really align with the model of "we're just marking files that contain objc constructs"

[dgoldman]

Fixes for review

[dgoldman]

Rebase

[dgoldman]

Rebase

[kadircet]

it feels like rebasing went wrong. changes from 2 unrelated patches seem to be part of this one now. can you make sure this patch only contains the delta for symbolcollector/clangd pieces?

clang-tools-extra/clangd/Headers.h

48

let's rename this to SymbolInclude as it looks too similar to HeaderFile right now.

also adding a comment like: A header and directives as stored in a Symbol.

50

let's mention that this is either a URI or verbatim include in the comments

clang-tools-extra/clangd/index/SymbolCollector.cpp

870

can we add a comment like Only allow #import for symbols from objc-like files.

clang/lib/Tooling/Inclusions/HeaderAnalysis.cpp

76

SM.getFileEntryForID(SM.getMainFileID()) == FE

[dgoldman]

Fixes + hopefully proper diffbase

[kadircet]

thanks, lgtm!

clang/unittests/Tooling/HeaderAnalysisTest.cpp

69

indentation

69

i think another important positive test case is, having it after some non PP code blocks (e.g. put it after main)

[dgoldman]

Update test

clang-tools-extra/clangd/CodeComplete.cpp

 /// A code completion result, in clang-native form.
 /// It may be promoted to a CompletionItem if it's among the top-ranked results.
 struct CompletionCandidate {
   llvm::StringRef Name; // Used for filtering and sorting.
   // We may have a result from Sema, from the index, or both.
   const CodeCompletionResult *SemaResult = nullptr;
   const Symbol *IndexResult = nullptr;
   const RawIdentifier *IdentifierResult = nullptr;
-  llvm::SmallVector<llvm::StringRef, 1> RankedIncludeHeaders;
+  llvm::SmallVector<SymbolInclude, 1> RankedIncludeHeaders;
 
   // Returns a token identifying the overload set this is part of.
   // 0 indicates it's not part of any overload set.
   size_t overloadSet(const CodeCompleteOptions &Opts, llvm::StringRef FileName,
                      IncludeInserter *Inserter) const {
     if (!Opts.BundleOverloads.value_or(false))
       return 0;
 
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;
     }
-    return RankedIncludeHeaders[0];
+    for (const auto &Inc : RankedIncludeHeaders)
+      // FIXME: We should support #import directives here.
+      if ((Inc.Directive & clang::clangd::Symbol::Include) != 0)
+        return Inc.Header;
+    return None;
   }
 
   using Bundle = llvm::SmallVector<CompletionCandidate, 4>;
 };
 using ScoredBundle =
     std::pair<CompletionCandidate::Bundle, CodeCompletion::Scores>;
 struct ScoredBundleGreater {
   bool operator()(const ScoredBundle &L, const ScoredBundle &R) {
auto Inserted = [&](llvm::StringRef Header)
         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();
     // Calculate include paths and edits for all possible headers.
     for (const auto &Inc : C.RankedIncludeHeaders) {
-      if (auto ToInclude = Inserted(Inc)) {
+      // FIXME: We should support #import directives here.
+      if ((Inc.Directive & clang::clangd::Symbol::Include) == 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);
         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, FileName,
+            C.IndexResult->CanonicalDeclaration.FileURI, Inc.Header, FileName,
             ToInclude.takeError());
     }
     // Prefer includes that do not need edits (i.e. already exist).
     std::stable_partition(Completion.Includes.begin(),
                           Completion.Includes.end(),
                           [](const CodeCompletion::IncludeCandidate &I) {
                             return !I.Insertion.has_value();
                           });

clang-tools-extra/clangd/Headers.h

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;
 };
 
+/// A header and directives as stored in a Symbol.
+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.
+  Symbol::IncludeDirective Directive;
+};
+
 /// Creates a `HeaderFile` from \p Header which can be either a URI or a literal
 /// include.
 llvm::Expected<HeaderFile> toHeaderFile(llvm::StringRef Header,
                                         llvm::StringRef HintPath);
 
 // Returns include headers for \p Sym sorted by popularity. If two headers are
 // equally popular, prefer the shorter one.
-llvm::SmallVector<llvm::StringRef, 1> getRankedIncludes(const Symbol &Sym);
+llvm::SmallVector<SymbolInclude, 1> getRankedIncludes(const Symbol &Sym);
 
 // An #include directive that we found in the main file.
 struct Inclusion {
   tok::PPKeywordKind Directive; // Directive used for inclusion, e.g. import
   std::string Written;          // Inclusion name as written e.g. <vector>.
   Path Resolved; // Resolved path of included file. Empty if not resolved.
   unsigned HashOffset = 0; // Byte offset from start of file to #.
   int HashLine = 0;        // Line number containing the directive, 0-indexed.

clang-tools-extra/clangd/Headers.cpp

if (!IncludePath->empty())
     return HeaderFile{std::move(*IncludePath), /*Verbatim=*/true};
 
   auto Resolved = URI::resolve(*U, HintPath);
   if (!Resolved)
     return Resolved.takeError();
   return HeaderFile{std::move(*Resolved), /*Verbatim=*/false};
 }
 
-llvm::SmallVector<llvm::StringRef, 1> getRankedIncludes(const Symbol &Sym) {
+llvm::SmallVector<SymbolInclude, 1> getRankedIncludes(const Symbol &Sym) {
   auto Includes = Sym.IncludeHeaders;
   // Sort in descending order by reference count and header length.
   llvm::sort(Includes, [](const Symbol::IncludeHeaderWithReferences &LHS,
                           const Symbol::IncludeHeaderWithReferences &RHS) {
     if (LHS.References == RHS.References)
       return LHS.IncludeHeader.size() < RHS.IncludeHeader.size();
     return LHS.References > RHS.References;
   });
-  llvm::SmallVector<llvm::StringRef, 1> Headers;
+  llvm::SmallVector<SymbolInclude, 1> Headers;
   for (const auto &Include : Includes)
-    Headers.push_back(Include.IncludeHeader);
+    Headers.push_back({Include.IncludeHeader, Include.supportedDirectives()});
   return Headers;
 }
 
 void IncludeStructure::collect(const CompilerInstance &CI) {
   auto &SM = CI.getSourceManager();
   MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
   auto Collector = std::make_unique<RecordHeaders>(CI, this);
   CI.getPreprocessor().addCommentHandler(Collector.get());

clang-tools-extra/clangd/IncludeFixer.cpp

std::vector<Fix> IncludeFixer::fixesForSymbols(const SymbolSlab &Syms) const {
 
   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)) {
-      if (auto ToInclude = Inserted(Sym, Inc)) {
+      // FIXME: We should support #import directives here.
+      if ((Inc.Directive & clang::clangd::Symbol::Include) == 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()))
             Fixes.push_back(std::move(*Fix));
         }
       } else {
-        vlog("Failed to calculate include insertion for {0} into {1}: {2}", Inc,
-             File, ToInclude.takeError());
+        vlog("Failed to calculate include insertion for {0} into {1}: {2}",
+             Inc.Header, File, ToInclude.takeError());
       }
     }
   }
   return Fixes;
 }
 
 // Returns the identifiers qualified by an unresolved name. \p Loc is the
 // start location of the unresolved name. For the example below, this returns

clang-tools-extra/clangd/index/Merge.cpp

Symbol mergeSymbol(const Symbol &L, const Symbol &R) {
   if (S.Type == "")
     S.Type = O.Type;
   for (const auto &OI : O.IncludeHeaders) {
     bool Found = false;
     for (auto &SI : S.IncludeHeaders) {
       if (SI.IncludeHeader == OI.IncludeHeader) {
         Found = true;
         SI.References += OI.References;
+        SI.SupportedDirectives |= OI.SupportedDirectives;
         break;
       }
     }
     if (!Found && MergeIncludes)
-      S.IncludeHeaders.emplace_back(OI.IncludeHeader, OI.References);
+      S.IncludeHeaders.emplace_back(OI.IncludeHeader, OI.References,
+                                    OI.supportedDirectives());
   }
 
   S.Origin |= O.Origin | SymbolOrigin::Merge;
   S.Flags |= O.Flags;
   return S;
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/Serialization.cpp

void writeSymbol(const Symbol &Sym, const StringTableOut &Strings,
   writeVar(Strings.index(Sym.Signature), OS);
   writeVar(Strings.index(Sym.CompletionSnippetSuffix), OS);
   writeVar(Strings.index(Sym.Documentation), OS);
   writeVar(Strings.index(Sym.ReturnType), OS);
   writeVar(Strings.index(Sym.Type), OS);
 
   auto WriteInclude = [&](const Symbol::IncludeHeaderWithReferences &Include) {
     writeVar(Strings.index(Include.IncludeHeader), OS);
-    writeVar(Include.References, OS);
+    writeVar((Include.References << 2) | Include.SupportedDirectives, OS);
   };
   writeVar(Sym.IncludeHeaders.size(), OS);
   for (const auto &Include : Sym.IncludeHeaders)
     WriteInclude(Include);
 }
 
 Symbol readSymbol(Reader &Data, llvm::ArrayRef<llvm::StringRef> Strings,
                   SymbolOrigin Origin) {
Symbol readSymbol(Reader &Data, llvm::ArrayRef<llvm::StringRef> Strings,
   Sym.CompletionSnippetSuffix = Data.consumeString(Strings);
   Sym.Documentation = Data.consumeString(Strings);
   Sym.ReturnType = Data.consumeString(Strings);
   Sym.Type = Data.consumeString(Strings);
   if (!Data.consumeSize(Sym.IncludeHeaders))
     return Sym;
   for (auto &I : Sym.IncludeHeaders) {
     I.IncludeHeader = Data.consumeString(Strings);
-    I.References = Data.consumeVar();
+    uint32_t RefsWithDirectives = Data.consumeVar();
+    I.References = RefsWithDirectives >> 2;
+    I.SupportedDirectives = RefsWithDirectives & 0x3;
   }
   return Sym;
 }
 
 // REFS ENCODING
 // A refs section has data grouped by Symbol. Each symbol has:
 //  - SymbolID: 8 bytes
 //  - NumRefs: varint

 //   - srcs: information related to include graph
 //   - stri: string table
 //   - symb: symbols
 //   - refs: references to symbols
 
 // The current versioning scheme is simple - non-current versions are rejected.
 // If you make a breaking change, bump this version number to invalidate stored
 // data. Later we may want to support some backward compatibility.
-constexpr static uint32_t Version = 17;
+constexpr static uint32_t Version = 18;
 
 llvm::Expected<IndexFileIn> readRIFF(llvm::StringRef Data,
                                      SymbolOrigin Origin) {
   auto RIFF = riff::readFile(Data);
   if (!RIFF)
     return RIFF.takeError();
   if (RIFF->Type != riff::fourCC("CdIx"))
     return error("wrong RIFF filetype: {0}", riff::fourCCStr(RIFF->Type));

clang-tools-extra/clangd/index/Symbol.h

 //===--- Symbol.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_INDEX_SYMBOL_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOL_H
 
 #include "index/SymbolID.h"
 #include "index/SymbolLocation.h"
 #include "index/SymbolOrigin.h"
 #include "clang/Index/IndexSymbol.h"
+#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/StringSaver.h"
 
 namespace clang {
 namespace clangd {
 
+LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
+
 /// The class presents a C++ symbol, e.g. class, function.
 ///
 /// WARNING: Symbols do not own much of their underlying data - typically
 /// strings are owned by a SymbolSlab. They should be treated as non-owning
 /// references. Copies are shallow.
 ///
 /// When adding new unowned data fields to Symbol, remember to update:
 ///   - SymbolSlab::Builder in Index.cpp, to copy them to the slab's storage.
struct Symbol {
   /// Only set when the symbol is indexed for completion.
   llvm::StringRef ReturnType;
 
   /// Raw representation of the OpaqueType of the symbol, used for scoring
   /// purposes.
   /// Only set when the symbol is indexed for completion.
   llvm::StringRef Type;
 
+  enum IncludeDirective : uint8_t {
+    Invalid = 0,
+    /// `#include "header.h"`
+    Include = 1,
+    /// `#import "header.h"`
+    Import = 2,
+
+    LLVM_MARK_AS_BITMASK_ENUM(Import)
+  };
+
   struct IncludeHeaderWithReferences {
     IncludeHeaderWithReferences() = default;
 
     IncludeHeaderWithReferences(llvm::StringRef IncludeHeader,
-                                unsigned References)
-        : IncludeHeader(IncludeHeader), References(References) {}
+                                uint32_t References,
+                                IncludeDirective SupportedDirectives)
+        : IncludeHeader(IncludeHeader), References(References),
+          SupportedDirectives(SupportedDirectives) {}
 
     /// This can be either a URI of the header to be #include'd
     /// for this symbol, or a literal header quoted with <> or "" that is
     /// suitable to be included directly. When it is a URI, the exact #include
     /// path needs to be calculated according to the URI scheme.
     ///
     /// Note that the include header is a canonical include for the symbol and
     /// can be different from FileURI in the CanonicalDeclaration.
     llvm::StringRef IncludeHeader = "";
     /// The number of translation units that reference this symbol and include
     /// this header. This number is only meaningful if aggregated in an index.
-    unsigned References = 0;
+    uint32_t References : 30;
+    /// Bitfield of supported directives (IncludeDirective) that can be used
+    /// when including this header.
+    uint32_t SupportedDirectives : 2;
+
+    IncludeDirective supportedDirectives() const {
+      return static_cast<IncludeDirective>(SupportedDirectives);
+    }
   };
   /// One Symbol can potentially be included via different headers.
   ///   - If we haven't seen a definition, this covers all declarations.
   ///   - If we have seen a definition, this covers declarations visible from
   ///   any definition.
   /// Only set when the symbol is indexed for completion.
   llvm::SmallVector<IncludeHeaderWithReferences, 1> IncludeHeaders;
 

clang-tools-extra/clangd/index/SymbolCollector.h

private:
   SymbolID getSymbolIDCached(const llvm::StringRef MacroName,
                              const MacroInfo *MI, const SourceManager &SM);
 
   // All Symbols collected from the AST.
   SymbolSlab::Builder Symbols;
   // File IDs for Symbol.IncludeHeaders.
   // The final spelling is calculated in finish().
   llvm::DenseMap<SymbolID, FileID> IncludeFiles;
+  // Files which contain ObjC symbols.
+  // This is finalized and used in finish().
+  llvm::DenseSet<FileID> FilesWithObjCConstructs;
   void setIncludeLocation(const Symbol &S, SourceLocation);
   // Indexed macros, to be erased if they turned out to be include guards.
   llvm::DenseSet<const IdentifierInfo *> IndexedMacros;
   // All refs collected from the AST. It includes:
   //   1) symbols declared in the preamble and referenced from the main file (
   //     which is not a header), or
   //   2) symbols declared and referenced from the main file (which is a header)
   RefSlab::Builder Refs;

clang-tools-extra/clangd/index/SymbolCollector.cpp

bool isPrivateProtoDecl(const NamedDecl &ND) {
   // FIXME: the heuristic relies on naming style (i.e. no underscore in
   // user-defined names) and can be improved.
   return (ND.getKind() != Decl::EnumConstant) || llvm::any_of(Name, islower);
 }
 
 // We only collect #include paths for symbols that are suitable for global code
 // completion, except for namespaces since #include path for a namespace is hard
 // to define.
-bool shouldCollectIncludePath(index::SymbolKind Kind) {
+Symbol::IncludeDirective shouldCollectIncludePath(index::SymbolKind Kind) {
   using SK = index::SymbolKind;
   switch (Kind) {
   case SK::Macro:
   case SK::Enum:
   case SK::Struct:
   case SK::Class:
   case SK::Union:
   case SK::TypeAlias:
   case SK::Using:
   case SK::Function:
   case SK::Variable:
   case SK::EnumConstant:
   case SK::Concept:
-    return true;
+    return Symbol::Include | Symbol::Import;
+  case SK::Protocol:
+    return Symbol::Import;
   default:
-    return false;
+    return Symbol::Invalid;
   }
 }
 
 // Return the symbol range of the token at \p TokLoc.
 std::pair<SymbolLocation::Position, SymbolLocation::Position>
 getTokenRange(SourceLocation TokLoc, const SourceManager &SM,
               const LangOptions &LangOpts) {
   auto CreatePosition = [&SM](SourceLocation Loc) {
     auto LSPLoc = sourceLocToPosition(SM, Loc);
     SymbolLocation::Position Pos;
     Pos.setLine(LSPLoc.line);
     Pos.setColumn(LSPLoc.character);
     return Pos;
   };
 
   auto TokenLength = clang::Lexer::MeasureTokenLength(TokLoc, SM, LangOpts);
   return {CreatePosition(TokLoc),
for (const auto &R : Relations) {
     // We currently do (B) because it's simpler.
     if (*RKind == RelationKind::BaseOf)
       this->Relations.insert({ID, *RKind, ObjectID});
     else if (*RKind == RelationKind::OverriddenBy)
       this->Relations.insert({ObjectID, *RKind, ID});
   }
 }
 
 void SymbolCollector::setIncludeLocation(const Symbol &S, SourceLocation Loc) {
-  if (Opts.CollectIncludePath)
-    if (shouldCollectIncludePath(S.SymInfo.Kind))
-      // Use the expansion location to get the #include header since this is
-      // where the symbol is exposed.
-      IncludeFiles[S.ID] =
-          PP->getSourceManager().getDecomposedExpansionLoc(Loc).first;
+  if (Opts.CollectIncludePath &&
+      shouldCollectIncludePath(S.SymInfo.Kind) != Symbol::Invalid)
 }
 
 void SymbolCollector::finish() {
   // At the end of the TU, add 1 to the refcount of all referenced symbols.
   for (const auto &ID : ReferencedSymbols) {
     if (const auto *S = Symbols.find(ID)) {
       // SymbolSlab::Builder returns const symbols because strings are interned
       // and modifying returned symbols without inserting again wouldn't go
       // well. const_cast is safe here as we're modifying a data owned by the
       // Symbol. This reduces time spent in SymbolCollector by ~1%.
       ++const_cast<Symbol *>(S)->References;
     }
   }
   if (Opts.CollectMacro) {
     assert(PP);
     // First, drop header guards. We can't identify these until EOF.
     for (const IdentifierInfo *II : IndexedMacros) {
       if (const auto *MI = PP->getMacroDefinition(II).getMacroInfo())
         if (auto ID =
                 getSymbolIDCached(II->getName(), MI, PP->getSourceManager()))
           if (MI->isUsedForHeaderGuard())
             Symbols.erase(ID);
     }
   }
+  llvm::DenseMap<FileID, bool> FileToContainsImportsOrObjC;
   // Fill in IncludeHeaders.
   // We delay this until end of TU so header guards are all resolved.
   llvm::SmallString<128> QName;
-  for (const auto &Entry : IncludeFiles) {
-    if (const Symbol *S = Symbols.find(Entry.first)) {
+  for (const auto &[SID, FID] : IncludeFiles) {
+    if (const Symbol *S = Symbols.find(SID)) {
       llvm::StringRef IncludeHeader;
       // Look for an overridden include header for this symbol specifically.
       if (Opts.Includes) {
         QName = S->Scope;
         QName.append(S->Name);
         IncludeHeader = Opts.Includes->mapSymbol(QName);
         if (!IncludeHeader.empty()) {
           if (IncludeHeader.front() != '"' && IncludeHeader.front() != '<')
             IncludeHeader = HeaderFileURIs->toURI(IncludeHeader);
           else if (IncludeHeader == "<utility>" && QName == "std::move" &&
                    S->Signature.contains(','))
             IncludeHeader = "<algorithm>";
         }
       }
       // Otherwise find the approprate include header for the defining file.
       if (IncludeHeader.empty())
-        IncludeHeader = HeaderFileURIs->getIncludeHeader(Entry.second);
+        IncludeHeader = HeaderFileURIs->getIncludeHeader(FID);
 
       // Symbols in slabs aren't mutable, insert() has to walk all the strings
       if (!IncludeHeader.empty()) {
+        Symbol::IncludeDirective Directives = Symbol::Invalid;
+        auto CollectDirectives = shouldCollectIncludePath(S->SymInfo.Kind);
+        if ((CollectDirectives & Symbol::Include) != 0)
+          Directives |= Symbol::Include;
+        // Only allow #import for symbols from ObjC-like files.
+        if ((CollectDirectives & Symbol::Import) != 0) {
+          auto [It, Inserted] = FileToContainsImportsOrObjC.try_emplace(FID);
+          if (Inserted)
+            It->second = FilesWithObjCConstructs.contains(FID) ||
+                         tooling::codeContainsImports(
+                             ASTCtx->getSourceManager().getBufferData(FID));
+          if (It->second)
+            Directives |= Symbol::Import;
+        }
+        if (Directives != Symbol::Invalid) {
           Symbol NewSym = *S;
-        NewSym.IncludeHeaders.push_back({IncludeHeader, 1});
+          NewSym.IncludeHeaders.push_back({IncludeHeader, 1, Directives});
           Symbols.insert(NewSym);
         }
       }
     }
+  }
 
   ReferencedSymbols.clear();
   IncludeFiles.clear();
+  FilesWithObjCConstructs.clear();
 }
 
 const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
                                               bool IsMainFileOnly) {
   auto &Ctx = ND.getASTContext();
   auto &SM = Ctx.getSourceManager();
 
   Symbol S;
const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
   if (isImplementationDetail(&ND))
     S.Flags |= Symbol::ImplementationDetail;
   if (!IsMainFileOnly)
     S.Flags |= Symbol::VisibleOutsideFile;
   S.SymInfo = index::getSymbolInfo(&ND);
   auto Loc = nameLocation(ND, SM);
   assert(Loc.isValid() && "Invalid source location for NamedDecl");
   // FIXME: use the result to filter out symbols.
-  shouldIndexFile(SM.getFileID(Loc));
+  auto FID = SM.getFileID(Loc);
+  shouldIndexFile(FID);
   if (auto DeclLoc = getTokenLocation(Loc))
     S.CanonicalDeclaration = *DeclLoc;
 
   S.Origin = Opts.Origin;
   if (ND.getAvailability() == AR_Deprecated)
     S.Flags |= Symbol::Deprecated;
 
   // Add completion info.
const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
   if (S.Flags & Symbol::IndexedForCodeCompletion) {
     TypeStorage = OpaqueType::fromCompletionResult(*ASTCtx, SymbolCompletion);
     if (TypeStorage)
       S.Type = TypeStorage->raw();
   }
 
   Symbols.insert(S);
   setIncludeLocation(S, ND.getLocation());
+  if (S.SymInfo.Lang == index::SymbolLanguage::ObjC)
+    FilesWithObjCConstructs.insert(FID);
   return Symbols.find(S.ID);
 }
 
 void SymbolCollector::addDefinition(const NamedDecl &ND,
                                     const Symbol &DeclSym) {
   if (DeclSym.Definition)
     return;
   const auto &SM = ND.getASTContext().getSourceManager();

clang-tools-extra/clangd/index/YAMLSerialization.cpp

 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 
+namespace {
+struct YIncludeHeaderWithReferences;
+}
+
 LLVM_YAML_IS_SEQUENCE_VECTOR(clang::clangd::Symbol::IncludeHeaderWithReferences)
 LLVM_YAML_IS_SEQUENCE_VECTOR(clang::clangd::Ref)
+LLVM_YAML_IS_SEQUENCE_VECTOR(YIncludeHeaderWithReferences)
 
 namespace {
 using RefBundle =
     std::pair<clang::clangd::SymbolID, std::vector<clang::clangd::Ref>>;
 // This is a pale imitation of std::variant<Symbol, RefBundle, Relation>
 struct VariantEntry {
   llvm::Optional<clang::clangd::Symbol> Symbol;
   llvm::Optional<RefBundle> Refs;
   llvm::Optional<clang::clangd::Relation> Relation;
   llvm::Optional<clang::clangd::IncludeGraphNode> Source;
   llvm::Optional<clang::tooling::CompileCommand> Cmd;
 };
 // A class helps YAML to serialize the 32-bit encoded position (Line&Column),
 // as YAMLIO can't directly map bitfields.
 struct YPosition {
   uint32_t Line;
   uint32_t Column;
 };
+// A class helps YAML to serialize the IncludeHeaderWithReferences as YAMLIO
+// can't directly map bitfields.
+struct YIncludeHeaderWithReferences {
+  llvm::StringRef IncludeHeader;
+  uint32_t References;
+  clang::clangd::Symbol::IncludeDirective SupportedDirectives;
+
+  YIncludeHeaderWithReferences() = default;
+
+  YIncludeHeaderWithReferences(
+      llvm::StringRef IncludeHeader, uint32_t References,
+      clang::clangd::Symbol::IncludeDirective SupportedDirectives)
+      : IncludeHeader(IncludeHeader), References(References),
+        SupportedDirectives(SupportedDirectives) {}
+};
 
 // avoid ODR violation of specialization for non-owned CompileCommand
 struct CompileCommandYAML : clang::tooling::CompileCommand {};
 
 } // namespace
 namespace llvm {
 namespace yaml {
 

 template <> struct MappingTraits<SymbolInfo> {
   static void mapping(IO &IO, SymbolInfo &SymInfo) {
     // FIXME: expose other fields?
     IO.mapRequired("Kind", SymInfo.Kind);
     IO.mapRequired("Lang", SymInfo.Lang);
   }
 };
 
-template <>
-struct MappingTraits<clang::clangd::Symbol::IncludeHeaderWithReferences> {
-  static void mapping(IO &IO,
-                      clang::clangd::Symbol::IncludeHeaderWithReferences &Inc) {
+template <> struct ScalarBitSetTraits<clang::clangd::Symbol::IncludeDirective> {
+  static void bitset(IO &IO, clang::clangd::Symbol::IncludeDirective &Value) {
+    IO.bitSetCase(Value, "Include", clang::clangd::Symbol::Include);
+    IO.bitSetCase(Value, "Import", clang::clangd::Symbol::Import);
+  }
+};
+
+template <> struct MappingTraits<YIncludeHeaderWithReferences> {
+  static void mapping(IO &IO, YIncludeHeaderWithReferences &Inc) {
     IO.mapRequired("Header", Inc.IncludeHeader);
     IO.mapRequired("References", Inc.References);
+    IO.mapOptional("Directives", Inc.SupportedDirectives,
+                   clang::clangd::Symbol::Include);
+  }
+};
+
+struct NormalizedIncludeHeaders {
+  using IncludeHeader = clang::clangd::Symbol::IncludeHeaderWithReferences;
+  NormalizedIncludeHeaders(IO &) {}
+  NormalizedIncludeHeaders(
+      IO &, const llvm::SmallVector<IncludeHeader, 1> &IncludeHeaders) {
+    for (auto &I : IncludeHeaders) {
+      Headers.emplace_back(I.IncludeHeader, I.References,
+                           I.supportedDirectives());
+    }
+  }
+
+  llvm::SmallVector<IncludeHeader, 1> denormalize(IO &) {
+    llvm::SmallVector<IncludeHeader, 1> Result;
+    for (auto &H : Headers)
+      Result.emplace_back(H.IncludeHeader, H.References, H.SupportedDirectives);
+    return Result;
   }
+  llvm::SmallVector<YIncludeHeaderWithReferences, 1> Headers;
 };
 
 template <> struct MappingTraits<Symbol> {
   static void mapping(IO &IO, Symbol &Sym) {
     MappingNormalization<NormalizedSymbolID, SymbolID> NSymbolID(IO, Sym.ID);
     MappingNormalization<NormalizedSymbolFlag, Symbol::SymbolFlag> NSymbolFlag(
         IO, Sym.Flags);
+    MappingNormalization<
+        NormalizedIncludeHeaders,
+        llvm::SmallVector<Symbol::IncludeHeaderWithReferences, 1>>
+        NIncludeHeaders(IO, Sym.IncludeHeaders);
     IO.mapRequired("ID", NSymbolID->HexString);
     IO.mapRequired("Name", Sym.Name);
     IO.mapRequired("Scope", Sym.Scope);
     IO.mapRequired("SymInfo", Sym.SymInfo);
     IO.mapOptional("CanonicalDeclaration", Sym.CanonicalDeclaration,
                    SymbolLocation());
     IO.mapOptional("Definition", Sym.Definition, SymbolLocation());
     IO.mapOptional("References", Sym.References, 0u);
     IO.mapOptional("Flags", NSymbolFlag->Flag);
     IO.mapOptional("Signature", Sym.Signature);
     IO.mapOptional("TemplateSpecializationArgs",
                    Sym.TemplateSpecializationArgs);
     IO.mapOptional("CompletionSnippetSuffix", Sym.CompletionSnippetSuffix);
     IO.mapOptional("Documentation", Sym.Documentation);
     IO.mapOptional("ReturnType", Sym.ReturnType);
     IO.mapOptional("Type", Sym.Type);
-    IO.mapOptional("IncludeHeaders", Sym.IncludeHeaders);
+    IO.mapOptional("IncludeHeaders", NIncludeHeaders->Headers);
   }
 };
 
 template <> struct ScalarEnumerationTraits<SymbolLanguage> {
   static void enumeration(IO &IO, SymbolLanguage &Value) {
     IO.enumCase(Value, "C", SymbolLanguage::C);
     IO.enumCase(Value, "Cpp", SymbolLanguage::CXX);
     IO.enumCase(Value, "ObjC", SymbolLanguage::ObjC);

clang-tools-extra/clangd/index/remote/Index.proto

 message Position {
   optional uint32 line = 1;
   optional uint32 column = 2;
 }
 
 message HeaderWithReferences {
   optional string header = 1;
   optional uint32 references = 2;
+  optional uint32 supported_directives = 3;
 }
 
 message RelationsRequest {
   repeated string subjects = 1;
   optional uint32 predicate = 2;
   optional uint32 limit = 3;
 }
 

clang-tools-extra/clangd/index/remote/marshalling/Marshalling.cpp

Marshaller::toProtobuf(const clangd::SymbolLocation &Location) {
   *Result.mutable_end() = toProtobuf(Location.End);
   return Result;
 }
 
 llvm::Expected<HeaderWithReferences> Marshaller::toProtobuf(
     const clangd::Symbol::IncludeHeaderWithReferences &IncludeHeader) {
   HeaderWithReferences Result;
   Result.set_references(IncludeHeader.References);
+  Result.set_supported_directives(IncludeHeader.SupportedDirectives);
   const std::string Header = IncludeHeader.IncludeHeader.str();
   if (isLiteralInclude(Header)) {
     Result.set_header(Header);
     return Result;
   }
   auto RelativePath = uriToRelativePath(Header);
   if (!RelativePath)
     return RelativePath.takeError();
   Result.set_header(*RelativePath);
   return Result;
 }
 
 llvm::Expected<clangd::Symbol::IncludeHeaderWithReferences>
 Marshaller::fromProtobuf(const HeaderWithReferences &Message) {
   std::string Header = Message.header();
   if (!isLiteralInclude(Header)) {
     auto URIString = relativePathToURI(Header);
     if (!URIString)
       return URIString.takeError();
     Header = *URIString;
   }
-  return clangd::Symbol::IncludeHeaderWithReferences{Strings.save(Header),
-                                                     Message.references()};
+  auto Directives = clangd::Symbol::IncludeDirective::Include;
+  if (Message.has_supported_directives())
+    Directives = static_cast<clangd::Symbol::IncludeDirective>(
+        Message.supported_directives());
+  return clangd::Symbol::IncludeHeaderWithReferences{
+      Strings.save(Header), Message.references(), Directives};
 }
 
 } // namespace remote
 } // namespace clangd
 } // namespace clang

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

 MATCHER_P(returnType, D, "") { return arg.ReturnType == D; }
 MATCHER_P(hasInclude, IncludeHeader, "") {
   return !arg.Includes.empty() && arg.Includes[0].Header == IncludeHeader;
 }
 MATCHER_P(insertInclude, IncludeHeader, "") {
   return !arg.Includes.empty() && arg.Includes[0].Header == IncludeHeader &&
          bool(arg.Includes[0].Insertion);
 }
+MATCHER_P(insertIncludeText, InsertedText, "") {
+  return !arg.Includes.empty() && arg.Includes[0].Insertion &&
+         arg.Includes[0].Insertion->newText == InsertedText;
+}
 MATCHER(insertInclude, "") {
   return !arg.Includes.empty() && bool(arg.Includes[0].Insertion);
 }
 MATCHER_P(snippetSuffix, Text, "") { return arg.SnippetSuffix == Text; }
 MATCHER_P(origin, OriginSet, "") { return arg.Origin == OriginSet; }
 MATCHER_P(signature, S, "") { return arg.Signature == S; }
 MATCHER_P(replacesRange, Range, "") {
   return arg.CompletionTokenRange == Range;

 TEST(CompletionTest, IncludeInsertionPreprocessorIntegrationTests) {
   TestTU TU;
   TU.ExtraArgs.push_back("-I" + testPath("sub"));
   TU.AdditionalFiles["sub/bar.h"] = "";
   auto BarURI = URI::create(testPath("sub/bar.h")).toString();
 
   Symbol Sym = cls("ns::X");
   Sym.CanonicalDeclaration.FileURI = BarURI.c_str();
-  Sym.IncludeHeaders.emplace_back(BarURI, 1);
+  Sym.IncludeHeaders.emplace_back(BarURI, 1, Symbol::Include);
   // Shorten include path based on search directory and insert.
   Annotations Test("int main() { ns::^ }");
   TU.Code = Test.code().str();
   auto Results = completions(TU, Test.point(), {Sym});
   EXPECT_THAT(Results.Completions,
               ElementsAre(AllOf(named("X"), insertInclude("\"bar.h\""))));
   // Can be disabled via option.
   CodeCompleteOptions NoInsertion;

 
 TEST(CompletionTest, NoIncludeInsertionWhenDeclFoundInFile) {
   Symbol SymX = cls("ns::X");
   Symbol SymY = cls("ns::Y");
   std::string BarHeader = testPath("bar.h");
   auto BarURI = URI::create(BarHeader).toString();
   SymX.CanonicalDeclaration.FileURI = BarURI.c_str();
   SymY.CanonicalDeclaration.FileURI = BarURI.c_str();
-  SymX.IncludeHeaders.emplace_back("<bar>", 1);
-  SymY.IncludeHeaders.emplace_back("<bar>", 1);
+  SymX.IncludeHeaders.emplace_back("<bar>", 1, Symbol::Include);
+  SymY.IncludeHeaders.emplace_back("<bar>", 1, Symbol::Include);
   // Shorten include path based on search directory and insert.
   auto Results = completions(R"cpp(
           namespace ns {
             class X;
             class Y {};
           }
           int main() { ns::^ }
       )cpp",
TEST(CompletionTest, OverloadBundling) {
   Symbol NoArgsGFunc = func("GFuncC");
   EXPECT_THAT(
       completions(Context + "int y = GFunc^", {NoArgsGFunc}, Opts).Completions,
       UnorderedElementsAre(labeled("GFuncC(…)"), labeled("GFuncD(int)")));
 
   // Differences in header-to-insert suppress bundling.
   std::string DeclFile = URI::create(testPath("foo")).toString();
   NoArgsGFunc.CanonicalDeclaration.FileURI = DeclFile.c_str();
-  NoArgsGFunc.IncludeHeaders.emplace_back("<foo>", 1);
+  NoArgsGFunc.IncludeHeaders.emplace_back("<foo>", 1, Symbol::Include);
   EXPECT_THAT(
       completions(Context + "int y = GFunc^", {NoArgsGFunc}, Opts).Completions,
       UnorderedElementsAre(AllOf(named("GFuncC"), insertInclude("<foo>")),
                            labeled("GFuncC(int)"), labeled("GFuncD(int)")));
 
   // Examine a bundled completion in detail.
   auto A =
       completions(Context + "int y = X().a^", {}, Opts).Completions.front();
TEST(CompletionTest, OverloadBundlingSameFileDifferentURI) {
 
   Symbol SymX = sym("ns::X", index::SymbolKind::Function, "@F@\\0#");
   Symbol SymY = sym("ns::X", index::SymbolKind::Function, "@F@\\0#I#");
   std::string BarHeader = testPath("bar.h");
   auto BarURI = URI::create(BarHeader).toString();
   SymX.CanonicalDeclaration.FileURI = BarURI.c_str();
   SymY.CanonicalDeclaration.FileURI = BarURI.c_str();
   // The include header is different, but really it's the same file.
-  SymX.IncludeHeaders.emplace_back("\"bar.h\"", 1);
-  SymY.IncludeHeaders.emplace_back(BarURI.c_str(), 1);
+  SymX.IncludeHeaders.emplace_back("\"bar.h\"", 1, Symbol::Include);
+  SymY.IncludeHeaders.emplace_back(BarURI.c_str(), 1, Symbol::Include);
 
   auto Results = completions("void f() { ::ns::^ }", {SymX, SymY}, Opts);
   // Expect both results are bundled, despite the different-but-same
   // IncludeHeader.
   ASSERT_EQ(1u, Results.Completions.size());
   const auto &R = Results.Completions.front();
   EXPECT_EQ("X", R.Name);
   EXPECT_EQ(2u, R.BundleSize);
TEST(CompletionTest, EnableSpeculativeIndexRequest) {
   auto Reqs3 = Requests.consumeRequests(2);
   ASSERT_EQ(Reqs3.size(), 2u);
 }
 
 TEST(CompletionTest, InsertTheMostPopularHeader) {
   std::string DeclFile = URI::create(testPath("foo")).toString();
   Symbol Sym = func("Func");
   Sym.CanonicalDeclaration.FileURI = DeclFile.c_str();
-  Sym.IncludeHeaders.emplace_back("\"foo.h\"", 2);
-  Sym.IncludeHeaders.emplace_back("\"bar.h\"", 1000);
+  Sym.IncludeHeaders.emplace_back("\"foo.h\"", 2, Symbol::Include);
+  Sym.IncludeHeaders.emplace_back("\"bar.h\"", 1000, Symbol::Include);
 
   auto Results = completions("Fun^", {Sym}).Completions;
   assert(!Results.empty());
   EXPECT_THAT(Results[0], AllOf(named("Func"), insertInclude("\"bar.h\"")));
   EXPECT_EQ(Results[0].Includes.size(), 2u);
 }
 
+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;
+  assert(!Results.empty());
+  EXPECT_THAT(Results[0],
+              AllOf(named("Func"), insertIncludeText("#include \"bar.h\"\n")));
+
+  Results = completions("Fun^", {Sym}, {}, "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")));
+
+  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) {
   Annotations Test(R"cpp(
     #include "foo.h"
     Fun^
   )cpp");
   auto TU = TestTU::withCode(Test.code());
   TU.AdditionalFiles["foo.h"] = "";
 
   std::string DeclFile = URI::create(testPath("foo")).toString();
   Symbol Sym = func("Func");
   Sym.CanonicalDeclaration.FileURI = DeclFile.c_str();
-  Sym.IncludeHeaders.emplace_back("\"foo.h\"", 2);
-  Sym.IncludeHeaders.emplace_back("\"bar.h\"", 1000);
+  Sym.IncludeHeaders.emplace_back("\"foo.h\"", 2, Symbol::Include);
+  Sym.IncludeHeaders.emplace_back("\"bar.h\"", 1000, Symbol::Include);
 
   EXPECT_THAT(completions(TU, Test.point(), {Sym}).Completions,
               UnorderedElementsAre(AllOf(named("Func"), hasInclude("\"foo.h\""),
                                          Not(insertInclude()))));
 }
 
 TEST(CompletionTest, MergeMacrosFromIndexAndSema) {
   Symbol Sym;

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

 std::unique_ptr<SymbolIndex>
 buildIndexWithSymbol(llvm::ArrayRef<SymbolWithHeader> Syms) {
   SymbolSlab::Builder Slab;
   for (const auto &S : Syms) {
     Symbol Sym = cls(S.QName);
     Sym.Flags |= Symbol::IndexedForCodeCompletion;
     Sym.CanonicalDeclaration.FileURI = S.DeclaringFile.c_str();
     Sym.Definition.FileURI = S.DeclaringFile.c_str();
-    Sym.IncludeHeaders.emplace_back(S.IncludeHeader, 1);
+    Sym.IncludeHeaders.emplace_back(S.IncludeHeader, 1, Symbol::Include);
     Slab.insert(Sym);
   }
   return MemIndex::build(std::move(Slab).build(), RefSlab(), RelationSlab());
 }
 
 TEST(IncludeFixerTest, IncompleteType) {
   auto TU = TestTU::withHeaderCode("namespace ns { class X; } ns::X *x;");
   TU.ExtraArgs.push_back("-std=c++20");
EXPECT_THAT(
         << Case.second;
   }
 }
 
 TEST(IncludeFixerTest, IncompleteEnum) {
   Symbol Sym = enm("X");
   Sym.Flags |= Symbol::IndexedForCodeCompletion;
   Sym.CanonicalDeclaration.FileURI = Sym.Definition.FileURI = "unittest:///x.h";
-  Sym.IncludeHeaders.emplace_back("\"x.h\"", 1);
+  Sym.IncludeHeaders.emplace_back("\"x.h\"", 1, Symbol::Include);
   SymbolSlab::Builder Slab;
   Slab.insert(Sym);
   auto Index =
       MemIndex::build(std::move(Slab).build(), RefSlab(), RelationSlab());
 
   TestTU TU;
   TU.ExternalIndex = Index.get();
   TU.ExtraArgs.push_back("-std=c++20");
int main() {
   x$access[[->]]f();
 }
   )cpp");
   auto TU = TestTU::withCode(Test.code());
   Symbol Sym = cls("ns::X");
   Sym.Flags |= Symbol::IndexedForCodeCompletion;
   Sym.CanonicalDeclaration.FileURI = "unittest:///x.h";
   Sym.Definition.FileURI = "unittest:///x.cc";
-  Sym.IncludeHeaders.emplace_back("\"x.h\"", 1);
+  Sym.IncludeHeaders.emplace_back("\"x.h\"", 1, Symbol::Include);
 
   SymbolSlab::Builder Slab;
   Slab.insert(Sym);
   auto Index =
       MemIndex::build(std::move(Slab).build(), RefSlab(), RelationSlab());
   TU.ExternalIndex = Index.get();
 
   EXPECT_THAT(*TU.build().getDiagnostics(),
TEST(IncludeFixerTest, CImplicitFunctionDecl) {
   Annotations Test("void x() { [[foo]](); }");
   auto TU = TestTU::withCode(Test.code());
   TU.Filename = "test.c";
   TU.ExtraArgs = {"-std=c99", "-Wno-error=implicit-function-declaration"};
 
   Symbol Sym = func("foo");
   Sym.Flags |= Symbol::IndexedForCodeCompletion;
   Sym.CanonicalDeclaration.FileURI = "unittest:///foo.h";
-  Sym.IncludeHeaders.emplace_back("\"foo.h\"", 1);
+  Sym.IncludeHeaders.emplace_back("\"foo.h\"", 1, Symbol::Include);
 
   SymbolSlab::Builder Slab;
   Slab.insert(Sym);
   auto Index =
       MemIndex::build(std::move(Slab).build(), RefSlab(), RelationSlab());
   TU.ExternalIndex = Index.get();
 
   EXPECT_THAT(

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

MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References, "") {
   return (arg.IncludeHeader == IncludeHeader) && (arg.References == References);
 }
 
 TEST(MergeTest, MergeIncludesOnDifferentDefinitions) {
   Symbol L, R;
   L.Name = "left";
   R.Name = "right";
   L.ID = R.ID = SymbolID("hello");
-  L.IncludeHeaders.emplace_back("common", 1);
-  R.IncludeHeaders.emplace_back("common", 1);
-  R.IncludeHeaders.emplace_back("new", 1);
+  L.IncludeHeaders.emplace_back("common", 1, Symbol::Include);
+  R.IncludeHeaders.emplace_back("common", 1, Symbol::Include);
+  R.IncludeHeaders.emplace_back("new", 1, Symbol::Include);
 
   // Both have no definition.
   Symbol M = mergeSymbol(L, R);
   EXPECT_THAT(M.IncludeHeaders,
               UnorderedElementsAre(IncludeHeaderWithRef("common", 2u),
                                    IncludeHeaderWithRef("new", 1u)));
 
   // Only merge references of the same includes but do not merge new #includes.
TEST(MergeIndexTest, IncludeHeadersMerged) {
   auto DynSize = DynSymbols.bytes() + DynRefs.bytes();
   auto DynData = std::make_pair(std::move(DynSymbols), std::move(DynRefs));
   llvm::StringSet<> DynFiles = {S.Definition.FileURI};
   MemIndex DynIndex(std::move(DynData.first), std::move(DynData.second),
                     RelationSlab(), std::move(DynFiles), IndexContents::Symbols,
                     std::move(DynData), DynSize);
 
   SymbolSlab::Builder StaticB;
-  S.IncludeHeaders.push_back({"<header>", 0});
+  S.IncludeHeaders.push_back({"<header>", 0, Symbol::Include});
   StaticB.insert(S);
   auto StaticIndex =
       MemIndex::build(std::move(StaticB).build(), RefSlab(), RelationSlab());
   MergedIndex Merge(&DynIndex, StaticIndex.get());
 
   EXPECT_THAT(runFuzzyFind(Merge, S.Name),
               ElementsAre(testing::Field(
                   &Symbol::IncludeHeaders,

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

End:
     Line: 1
     Column: 1
 Flags:    129
 Documentation:    'Foo doc'
 ReturnType:    'int'
 IncludeHeaders:
   - Header:    'include1'
     References:    7
+    Directives:      [ Include ]
   - Header:    'include2'
     References:    3
+    Directives:      [ Import ]
+  - Header:    'include3'
+    References:    2
+    Directives:      [ Include, Import ]
+  - Header:    'include4'
+    References:    1
+    Directives:      [ ]
 ...
 ---
 !Symbol
 ID: 057557CEBF6E6B2E
 Name:   'Foo2'
 Scope:   'clang::'
 SymInfo:
   Kind:            Function

 DirectIncludes:
   - 'file:///path/inc1.h'
   - 'file:///path/inc2.h'
 ...
 )";
 
 MATCHER_P(id, I, "") { return arg.ID == cantFail(SymbolID::fromStr(I)); }
 MATCHER_P(qName, Name, "") { return (arg.Scope + arg.Name).str() == Name; }
-MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References, "") {
-  return (arg.IncludeHeader == IncludeHeader) && (arg.References == References);
+MATCHER_P3(IncludeHeaderWithRefAndDirectives, IncludeHeader, References,
+           SupportedDirectives, "") {
+  return (arg.IncludeHeader == IncludeHeader) &&
+         (arg.References == References) &&
+         (arg.SupportedDirectives == SupportedDirectives);
 }
 
 auto readIndexFile(llvm::StringRef Text) {
   return readIndexFile(Text, SymbolOrigin::Static);
 }
 
 TEST(SerializationTest, NoCrashOnEmptyYAML) {
   EXPECT_TRUE(bool(readIndexFile("")));
TEST(SerializationTest, YAMLConversions) {
   EXPECT_EQ(Sym1.Signature, "");
   EXPECT_EQ(Sym1.Documentation, "Foo doc");
   EXPECT_EQ(Sym1.ReturnType, "int");
   EXPECT_EQ(StringRef(Sym1.CanonicalDeclaration.FileURI), "file:///path/foo.h");
   EXPECT_EQ(Sym1.Origin, SymbolOrigin::Static);
   EXPECT_EQ(static_cast<uint8_t>(Sym1.Flags), 129);
   EXPECT_TRUE(Sym1.Flags & Symbol::IndexedForCodeCompletion);
   EXPECT_FALSE(Sym1.Flags & Symbol::Deprecated);
-  EXPECT_THAT(Sym1.IncludeHeaders,
-              UnorderedElementsAre(IncludeHeaderWithRef("include1", 7u),
-                                   IncludeHeaderWithRef("include2", 3u)));
+  EXPECT_THAT(
+      Sym1.IncludeHeaders,
+      UnorderedElementsAre(
+          IncludeHeaderWithRefAndDirectives("include1", 7u, Symbol::Include),
+          IncludeHeaderWithRefAndDirectives("include2", 3u, Symbol::Import),
+          IncludeHeaderWithRefAndDirectives("include3", 2u,
+                                            Symbol::Include | Symbol::Import),
+          IncludeHeaderWithRefAndDirectives("include4", 1u, Symbol::Invalid)));
 
   EXPECT_THAT(Sym2, qName("clang::Foo2"));
   EXPECT_EQ(Sym2.Signature, "-sig");
   EXPECT_EQ(Sym2.ReturnType, "");
   EXPECT_EQ(llvm::StringRef(Sym2.CanonicalDeclaration.FileURI),
             "file:///path/bar.h");
   EXPECT_FALSE(Sym2.Flags & Symbol::IndexedForCodeCompletion);
   EXPECT_TRUE(Sym2.Flags & Symbol::Deprecated);

clang/include/clang/Tooling/Inclusions/HeaderAnalysis.h

 class SourceManager;
 class HeaderSearch;
 
 namespace tooling {
 
 /// Returns true if the given physical file is a self-contained header.
 ///
 /// A header is considered self-contained if
-//   - it has a proper header guard or has been #imported
+//   - it has a proper header guard or has been #imported or contains #import(s)
 //   - *and* it doesn't have a dont-include-me pattern.
 ///
 /// This function can be expensive as it may scan the source code to find out
 /// dont-include-me pattern heuristically.
 bool isSelfContainedHeader(const FileEntry *FE, const SourceManager &SM,
                            HeaderSearch &HeaderInfo);
 
+/// This scans the given source code to see if it contains #import(s).
+bool codeContainsImports(llvm::StringRef Code);
+
 /// If Text begins an Include-What-You-Use directive, returns it.
 /// Given "// IWYU pragma: keep", returns "keep".
 /// Input is a null-terminated char* as provided by SM.getCharacterData().
 /// (This should not be StringRef as we do *not* want to scan for its length).
 /// For multi-line comments, we return only the first line.
 llvm::Optional<llvm::StringRef> parseIWYUPragma(const char *Text);
 
 } // namespace tooling
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLING_INCLUSIONS_HEADER_ANALYSIS_H

clang/lib/Tooling/Inclusions/HeaderAnalysis.cpp

bool isErrorAboutInclude(llvm::StringRef Line) {
   Line = Line.ltrim();
   if (!Line.startswith("error"))
     return false;
   return Line.contains_insensitive(
       "includ"); // Matches "include" or "including".
 }
 
 // Heuristically headers that only want to be included via an umbrella.
-bool isDontIncludeMeHeader(llvm::MemoryBufferRef Buffer) {
-  StringRef Content = Buffer.getBuffer();
+bool isDontIncludeMeHeader(StringRef Content) {
   llvm::StringRef Line;
   // Only sniff up to 100 lines or 10KB.
   Content = Content.take_front(100 * 100);
   for (unsigned I = 0; I < 100 && !Content.empty(); ++I) {
     std::tie(Line, Content) = Content.split('\n');
     if (isIf(Line) && isErrorAboutInclude(Content.split('\n').first))
       return true;
   }
   return false;
 }
 
+bool isImportLine(llvm::StringRef Line) {
+  Line = Line.ltrim();
+  if (!Line.consume_front("#"))
+    return false;
+  Line = Line.ltrim();
+  return Line.startswith("import");
+}
+
+llvm::StringRef getFileContents(const FileEntry *FE, const SourceManager &SM) {
+  return const_cast<SourceManager &>(SM)
+      .getMemoryBufferForFileOrNone(FE)
+      .value_or(llvm::MemoryBufferRef())
+      .getBuffer();
+}
+
 } // namespace
 
 bool isSelfContainedHeader(const FileEntry *FE, const SourceManager &SM,
                            HeaderSearch &HeaderInfo) {
   assert(FE);
   if (!HeaderInfo.isFileMultipleIncludeGuarded(FE) &&
-      !HeaderInfo.hasFileBeenImported(FE))
+      !HeaderInfo.hasFileBeenImported(FE) &&
+      // Any header that contains #imports is supposed to be #import'd so no
+      // need to check for anything but the main-file.
+      (SM.getFileEntryForID(SM.getMainFileID()) != FE ||
+       !codeContainsImports(getFileContents(FE, SM))))
     return false;
   // This pattern indicates that a header can't be used without
   // particular preprocessor state, usually set up by another header.
-  return !isDontIncludeMeHeader(
-      const_cast<SourceManager &>(SM).getMemoryBufferForFileOrNone(FE).value_or(
-          llvm::MemoryBufferRef()));
+  return !isDontIncludeMeHeader(getFileContents(FE, SM));
+}
+
+bool codeContainsImports(llvm::StringRef Code) {
+  // Only sniff up to 100 lines or 10KB.
+  Code = Code.take_front(100 * 100);
+  llvm::StringRef Line;
+  for (unsigned I = 0; I < 100 && !Code.empty(); ++I) {
+    std::tie(Line, Code) = Code.split('\n');
+    if (isImportLine(Line))
+      return true;
+  }
+  return false;
 }
 
 llvm::Optional<StringRef> parseIWYUPragma(const char *Text) {
   // Skip the comment start, // or /*.
   if (Text[0] != '/' || (Text[1] != '/' && Text[1] != '*'))
     return std::nullopt;
   bool BlockComment = Text[1] == '*';
   Text += 2;

clang/unittests/Tooling/HeaderAnalysisTest.cpp

TEST(HeaderAnalysisTest, IsSelfContained) {
 
   TestAST AST(Inputs);
   const auto &SM = AST.sourceManager();
   auto &FM = SM.getFileManager();
   auto &HI = AST.preprocessor().getHeaderSearchInfo();
   EXPECT_TRUE(isSelfContainedHeader(FM.getFile("headerguard.h").get(), SM, HI));
   EXPECT_TRUE(isSelfContainedHeader(FM.getFile("pragmaonce.h").get(), SM, HI));
   EXPECT_TRUE(isSelfContainedHeader(FM.getFile("imported.h").get(), SM, HI));
+  EXPECT_TRUE(
+      isSelfContainedHeader(SM.getFileEntryForID(SM.getMainFileID()), SM, HI));
 
   EXPECT_FALSE(isSelfContainedHeader(FM.getFile("unguarded.h").get(), SM, HI));
   EXPECT_FALSE(isSelfContainedHeader(FM.getFile("bad.h").get(), SM, HI));
 }
 
+TEST(HeaderAnalysisTest, CodeContainsImports) {
+  EXPECT_TRUE(codeContainsImports(R"cpp(
+  #include "foo.h"
+  #import "NSFoo.h"
+
+  int main() {
+    foo();
+  }
+  )cpp"));
+
+  EXPECT_TRUE(codeContainsImports(R"cpp(
+  #include "foo.h"
+
+  int main() {
+    foo();
+  }
+
+  #import "NSFoo.h"
+  )cpp"));
+
+  EXPECT_FALSE(codeContainsImports(R"cpp(
+  #include "foo.h"
+
+  int main() {
+    foo();
+  }
+  )cpp"));
+}
+
 TEST(HeaderAnalysisTest, ParseIWYUPragma) {
   EXPECT_THAT(parseIWYUPragma("// IWYU pragma: keep"), ValueIs(Eq("keep")));
   EXPECT_THAT(parseIWYUPragma("// IWYU pragma:   keep  me\netc"),
               ValueIs(Eq("keep  me")));
   EXPECT_THAT(parseIWYUPragma("/* IWYU pragma: keep */"), ValueIs(Eq("keep")));
   EXPECT_EQ(parseIWYUPragma("//  IWYU pragma: keep"), std::nullopt)
       << "Prefix is sensitive to whitespace";
   EXPECT_EQ(parseIWYUPragma("// IWYU pragma:keep"), std::nullopt)
       << "Prefix is sensitive to whitespace";
   EXPECT_EQ(parseIWYUPragma("/\n* IWYU pragma: keep */"), std::nullopt)
       << "Must start with /* or //";
 }
 
 } // namespace
 } // namespace tooling
 } // namespace clang