[clangd] Refactor IncludeStructure: use File (unsigned) for most computations

[kbobyrev]

Summary

Preparation for D108194.

[sammccall]

clang-tools-extra/clangd/Headers.h

116

this concept needs documentation.

116

there's a naming issue here:

• "File" is an overloaded enough concept that you still end up referring to these as "file indexes" below, which isn't clear.
• obviously this integer isn't either the file itself nor our record of it, so the metaphor is a bit confusing

What do you think about just ID, which would be spelled from the outside as IncludeStructure::ID?
Maybe HeaderID is better. I don't think it's a big problem that we have one for the main file too...

116–127

this public member is now sandwiched between methods

118

"Name" certainly needs to be clarified here.
Do you have any use cases in mind where you wouldn't have a FileEntry on hand? That would make this typesafe (and in particular avoid passing absolute filenames, or strings like "utility" or "<utility>)

125

it's unclear from this patch whether you want to expose the whole graph or just allow querying per-file - I doubt we need both.

125

tests?

135–140

The public interface now uses the File abstraction, so this method sticks out. (It has an awkward interface in the first place due to the need to use strings)

Suggest replacing it with DenseMap<File, unsigned> includeDepth(File Root). We'd need to add StringRef RealPath(File) for the caller, but that seems to make sense anyway.

144

this isn't abstract :-)

also the public nature of it isn't used, in this patch at least.

145

File should be capitalized.

FWIW, for a private method I prefer the old name here

[kbobyrev]

Resolve review comments.

[kbobyrev]

s/ToVector/Flatten/g

[kbobyrev]

Update non-working test

[kbobyrev]

I resolved the comments we discussed offline, there's a weird SIGSEGV (out of boundary error) somewhere in the tests but it looks fine otherwise.

[kbobyrev]

Fix the last failing test. All tests are green now.

[kbobyrev]

Clean tests.

[tschuett]

Do you want to wrap the unsigned in a struct to make it slightly safer, but more complicated to use?

[sammccall]

Just nits left really.

In D110386#3020890, @tschuett wrote:

Do you want to wrap the unsigned in a struct to make it slightly safer, but more complicated to use?

I like this idea, though it's probably not a huge deal either way.
(I marginally prefer enum class HeaderID : unsigned {}; over a struct, but it's basically the same thing)

clang-tools-extra/clangd/CodeComplete.cpp

1392

this isn't the ID of the include structure, it's the ID of the main file

1393

IncludeStructure is mutable here, I think it'd be safe to getOrCreate here.

Alternatively, make this infallible having IncludeStructure record the HeaderID for the main file in its constructor, and just use Includes.Root here.

Either way, I'd rather not deal with Expected<> error handling at runtime if we can define this error away.

clang-tools-extra/clangd/Headers.h

116–127

As mentioned offline, this is an implementation comment, but it doesn't say *what the HeaderID is*.

Suggest something like:

HeaderID identifies file in the include graph.
It corresponds to a FileEntry rather than a FileID, but stays stable across preamble & main file builds.

I'd move this impl comment back down to the private StringMap, but it could also follow the description.

119

comment still refers to "file index" as if it were introducing a term we use elsewhere

125

I think this can be Optional rather than Expected - much more lightweight.

129

this should be an assertion (or just leave that up to the vector) - passing an out-of-range HeaderID is invalid

134–135

comment is out of date

136

getName() is not right here.
(Hopefully the example will fit on the "Usually" line when fixed?)

143

we don't need this method if IncludeChildren is exposed, right?

[kbobyrev]

Resolve all review comments.

[kbobyrev]

Rebase on top of main.

[sammccall]

clang-tools-extra/clangd/Headers.cpp

174

no auto here. I really hope this isn't std::string :-)

Actually I think you can avoid extracting this variable at all, because the !RealPathName.empty() check is redundant. The case where both sides are empty is rare enough that it's fine to assign as long as the *stored* name is empty.

176

extract std::string &RealPathName = RealPathNames[...]?

186

This is a "can't happen" condition - remove or replace with an assert?

199–200

This is no longer true, we don't need the check.

clang-tools-extra/clangd/Headers.h

116–127

This doesn't seem to be done.
You've added a line to the *bottom* of the implementation comment defining HeaderID in terms of the implementation.

A reader rather needs complete self-contained description of the concept at the *top*, optionally followed by an explanation about the implementation (separated by a blank line).

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

71

reparsing the code several times to obtain header IDs doesn't seem reasonable. It's surprising behavior from a test helper, and it requires that header IDs are stable across parses!

Can we just pass in the IncludeStructure instead?

241

Why are we asserting on every element of the map one at a time, instead of the whole map at once? Seems like it would be more regular and easier to read.

I'd probably just write:

DenseMap<HeaderID, SmallVector<HeaderID>> Expected = { ... };
EXPECT_EQ(Expected, Includes.IncludeChildren);

243

EXPECT_THAT(foo.empty()) -> EXPECT_THAT(foo, IsEmpty()).
(Much better error messages)

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

511–512

The ratio of setup to punchline of the rest of this test seems like it's gotten out of hand (it wasn't great before, but it's doubled).

What about just

auto &FM = PatchedAST->getSourceManager().getFileManager();
auto &Includes = PatchedAST->getIncludeStructure();
auto MainID = Includes.getID(FM.getFile(testPath("foo.cpp")));
auto AuxID = Includes.getID(FM.getFile(testPath("sub/aux.h")));
EXPECT_THAT(Includes.IncludeChildren.at(MainID), Contains(AuxID));

It tests a bit less I guess, but seems like enough

554–555

and here

[kbobyrev]

Resolve all but two comments, leave the clarification questions for these two.

clang-tools-extra/clangd/Headers.cpp

199–200

Wait, why not? We still have unresolved includes, e.g. preamble patches are like that, their RealPathNames stay empty.

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

241

This would expect the elements in the map to be in a particular order, isn't this something we don't want?

[sammccall]

clang-tools-extra/clangd/Headers.cpp

199–200

Right, but why do we need to filter them out here?
Can't we just drop them when we use them to seed the proximity sources?

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

241

== on DenseMap doesn't consider order.

If you mean the order within map values (i.e. lists of child edges): these are in the order the #includes appear in the file, which seems fine to depend on to me

[kbobyrev]

Switch to full DenseMap comparison in the tests.

clang-tools-extra/clangd/Headers.cpp

199–200

I feel like producing them in the first place is kind of redundant, what signal does it give to the user through such API?

Getting some "hidden"/"unresolved" includes certainly feels rather confusing to me, post-filtering outside probably requires understanding of the internals which does not look like a good idea. WDYT?

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

241

Yes, I initially meant the SmallVector order, but yeah, this seems reasonable, thank you!

[kbobyrev]

Resolve the last comment

[thakis]

Looks like this breaks tests on Windows: http://45.33.8.238/win/45971/step_9.txt

PTAL!

[thakis]

Reverted in 36dc5c048ac745ef62769dd824a9a13659afc2ef for now.

clang-tools-extra/clangd/CodeComplete.cpp

auto RecorderOwner = std::make_unique<CompletionRecorder>(Opts, [&]() {
 
       // Most of the cost of file proximity is in initializing the FileDistance
       // structures based on the observed includes, once per query. Conceptually
       // that happens here (though the per-URI-scheme initialization is lazy).
       // The per-result proximity scoring is (amortized) very cheap.
       FileDistanceOptions ProxOpts{}; // Use defaults.
       const auto &SM = Recorder->CCSema->getSourceManager();
       llvm::StringMap<SourceParams> ProxSources;
-      for (auto &Entry : Includes.includeDepth(
-               SM.getFileEntryForID(SM.getMainFileID())->getName())) {
-        auto &Source = ProxSources[Entry.getKey()];
-        Source.Cost = Entry.getValue() * ProxOpts.IncludeCost;
+      auto MainFileID =
+          Includes.getOrCreateID(SM.getFileEntryForID(SM.getMainFileID()));
+      for (auto &HeaderIDAndDepth : Includes.includeDepth(MainFileID)) {
+        auto &Source =
+            ProxSources[Includes.getRealPath(HeaderIDAndDepth.getFirst())];
+        Source.Cost = HeaderIDAndDepth.getSecond() * ProxOpts.IncludeCost;
         // Symbols near our transitive includes are good, but only consider
         // things in the same directory or below it. Otherwise there can be
         // many false positives.
-        if (Entry.getValue() > 0)
+        if (HeaderIDAndDepth.getSecond() > 0)
           Source.MaxUpTraversals = 1;
       }
       FileProximity.emplace(ProxSources, ProxOpts);
 
       Output = runWithSema();
       Inserter.reset(); // Make sure this doesn't out-live Clang.
       SPAN_ATTACH(Tracer, "sema_completion_kind",
                   getCompletionKindString(CCContextKind));
       log("Code complete: sema context {0}, query scopes [{1}] (AnyScope={2}), "
           "expected type {3}{4}",

clang-tools-extra/clangd/Headers.h

 //===--- Headers.h - Include headers -----------------------------*- 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_HEADERS_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H
 
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "index/Symbol.h"
+#include "support/Logger.h"
 #include "support/Path.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Format/Format.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Tooling/Inclusions/HeaderIncludes.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
struct Inclusion {
   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.
   SrcMgr::CharacteristicKind FileKind = SrcMgr::C_User;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Inclusion &);
 bool operator==(const Inclusion &LHS, const Inclusion &RHS);
 
-// Contains information about one file in the build grpah and its direct
+// Contains information about one file in the build graph and its direct
 // dependencies. Doesn't own the strings it references (IncludeGraph is
 // self-contained).
 struct IncludeGraphNode {
   enum class SourceFlag : uint8_t {
     None = 0,
     // Whether current file is a main file rather than a header.
     IsTU = 1 << 0,
     // Whether current file had any uncompilable errors during indexing.

 //
 // It should be built incrementally with collectIncludeStructureCallback().
 // When we build the preamble, we capture and store its include structure along
 // with the preamble data. When we use the preamble, we can copy its
 // IncludeStructure and use another collectIncludeStructureCallback() to fill
 // in any non-preamble inclusions.
 class IncludeStructure {
 public:
-  std::vector<Inclusion> MainFileIncludes;
+  // HeaderID identifies file in the include graph. It corresponds to a
+  // FileEntry rather than a FileID, but stays stable across preamble & main
+  // file builds.
+  enum class HeaderID : unsigned {};
+
+  llvm::Optional<HeaderID> getID(const FileEntry *Entry) const;
+  HeaderID getOrCreateID(const FileEntry *Entry);
+
+  StringRef getRealPath(HeaderID ID) const {
+    assert(static_cast<unsigned>(ID) <= RealPathNames.size());
+    return RealPathNames[static_cast<unsigned>(ID)];
+  }
 
   // Return all transitively reachable files.
   llvm::ArrayRef<std::string> allHeaders() const { return RealPathNames; }
 
   // Return all transitively reachable files, and their minimum include depth.
   // All transitive includes (absolute paths), with their minimum include depth.
   // Root --> 0, #included file --> 1, etc.
-  // Root is clang's name for a file, which may not be absolute.
-  // Usually it should be SM.getFileEntryForID(SM.getMainFileID())->getName().
-  llvm::StringMap<unsigned> includeDepth(llvm::StringRef Root) const;
-
-  // This updates IncludeDepth(), but not MainFileIncludes.
-  void recordInclude(llvm::StringRef IncludingName,
-                     llvm::StringRef IncludedName,
-                     llvm::StringRef IncludedRealName);
+  // Root is the ID of the header being visited first.
+  // Usually it is getID(SM.getFileEntryForID(SM.getMainFileID())->getName()).
+  llvm::DenseMap<HeaderID, unsigned> includeDepth(HeaderID Root) const;
+
+  // Maps HeaderID to the ids of the files included from it.
+  llvm::DenseMap<HeaderID, SmallVector<HeaderID>> IncludeChildren;
+
+  std::vector<Inclusion> MainFileIncludes;
 
 private:
+  std::vector<std::string> RealPathNames; // In HeaderID order.
+  // HeaderID maps the FileEntry::Name to the internal representation.
   // Identifying files in a way that persists from preamble build to subsequent
-  // builds is surprisingly hard. FileID is unavailable in InclusionDirective(),
-  // and RealPathName and UniqueID are not preserved in the preamble.
-  // We use the FileEntry::Name, which is stable, interned into a "file index".
-  // The paths we want to expose are the RealPathName, so store those too.
-  std::vector<std::string> RealPathNames; // In file index order.
-  unsigned fileIndex(llvm::StringRef Name);
-  llvm::StringMap<unsigned> NameToIndex; // Values are file indexes.
-  // Maps a file's index to that of the files it includes.
-  llvm::DenseMap<unsigned, llvm::SmallVector<unsigned>> IncludeChildren;
+  // builds is surprisingly hard. FileID is unavailable in
+  // InclusionDirective(), and RealPathName and UniqueID are not preserved in
+  // the preamble.
+  llvm::StringMap<HeaderID> NameToIndex;
 };
 
 /// Returns a PPCallback that visits all inclusions in the main file.
 std::unique_ptr<PPCallbacks>
 collectIncludeStructureCallback(const SourceManager &SM, IncludeStructure *Out);
 
 // Calculates insertion edit for including a new header in a file.
 class IncludeInserter {
private:
   HeaderSearch *HeaderSearchInfo = nullptr;
   llvm::StringSet<> IncludedHeaders; // Both written and resolved.
   tooling::HeaderIncludes Inserter;  // Computers insertion replacement.
 };
 
 } // namespace clangd
 } // namespace clang
 
+namespace llvm {
+
+// Support Tokens as DenseMap keys.
+template <> struct DenseMapInfo<clang::clangd::IncludeStructure::HeaderID> {
+  static inline clang::clangd::IncludeStructure::HeaderID getEmptyKey() {
+    return static_cast<clang::clangd::IncludeStructure::HeaderID>(
+        DenseMapInfo<unsigned>::getEmptyKey());
+  }
+
+  static inline clang::clangd::IncludeStructure::HeaderID getTombstoneKey() {
+    return static_cast<clang::clangd::IncludeStructure::HeaderID>(
+        DenseMapInfo<unsigned>::getTombstoneKey());
+  }
+
+  static unsigned
+  getHashValue(const clang::clangd::IncludeStructure::HeaderID &Tag) {
+    return hash_value(static_cast<unsigned>(Tag));
+  }
+
+  static bool isEqual(const clang::clangd::IncludeStructure::HeaderID &LHS,
+                      const clang::clangd::IncludeStructure::HeaderID &RHS) {
+    return LHS == RHS;
+  }
+};
+
+} // namespace llvm
+
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H

clang-tools-extra/clangd/Headers.cpp

void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
     if (File) {
       auto *IncludingFileEntry = SM.getFileEntryForID(SM.getFileID(HashLoc));
       if (!IncludingFileEntry) {
         assert(SM.getBufferName(HashLoc).startswith("<") &&
                "Expected #include location to be a file or <built-in>");
         // Treat as if included from the main file.
         IncludingFileEntry = SM.getFileEntryForID(MainFID);
       }
-      Out->recordInclude(IncludingFileEntry->getName(), File->getName(),
-                         File->tryGetRealPathName());
+      auto IncludingID = Out->getOrCreateID(IncludingFileEntry),
+           IncludedID = Out->getOrCreateID(File);
+      Out->IncludeChildren[IncludingID].push_back(IncludedID);
     }
   }
 
   void FileChanged(SourceLocation Loc, FileChangeReason Reason,
                    SrcMgr::CharacteristicKind FileType,
                    FileID PrevFID) override {
     switch (Reason) {
     case PPCallbacks::EnterFile:

 }
 
 std::unique_ptr<PPCallbacks>
 collectIncludeStructureCallback(const SourceManager &SM,
                                 IncludeStructure *Out) {
   return std::make_unique<RecordHeaders>(SM, Out);
 }
 
-void IncludeStructure::recordInclude(llvm::StringRef IncludingName,
-                                     llvm::StringRef IncludedName,
-                                     llvm::StringRef IncludedRealName) {
-  auto Child = fileIndex(IncludedName);
-  if (!IncludedRealName.empty() && RealPathNames[Child].empty())
-    RealPathNames[Child] = std::string(IncludedRealName);
-  auto Parent = fileIndex(IncludingName);
-  IncludeChildren[Parent].push_back(Child);
+llvm::Optional<IncludeStructure::HeaderID>
+IncludeStructure::getID(const FileEntry *Entry) const {
+  auto It = NameToIndex.find(Entry->getName());
+  if (It == NameToIndex.end())
+    return llvm::None;
+  return It->second;
 }
 
-unsigned IncludeStructure::fileIndex(llvm::StringRef Name) {
-  auto R = NameToIndex.try_emplace(Name, RealPathNames.size());
+IncludeStructure::HeaderID
+IncludeStructure::getOrCreateID(const FileEntry *Entry) {
+  auto R = NameToIndex.try_emplace(
+      Entry->getName(),
+      static_cast<IncludeStructure::HeaderID>(RealPathNames.size()));
   if (R.second)
     RealPathNames.emplace_back();
-  return R.first->getValue();
+  IncludeStructure::HeaderID Result = R.first->getValue();
+  std::string &RealPathName = RealPathNames[static_cast<unsigned>(Result)];
+  if (RealPathName.empty())
+    RealPathName = Entry->tryGetRealPathName().str();
+  return Result;
 }
 
-llvm::StringMap<unsigned>
-IncludeStructure::includeDepth(llvm::StringRef Root) const {
+llvm::DenseMap<IncludeStructure::HeaderID, unsigned>
+IncludeStructure::includeDepth(HeaderID Root) const {
   // Include depth 0 is the main file only.
-  llvm::StringMap<unsigned> Result;
+  llvm::DenseMap<HeaderID, unsigned> Result;
+  assert(static_cast<unsigned>(Root) < RealPathNames.size());
   Result[Root] = 0;
-  std::vector<unsigned> CurrentLevel;
-  llvm::DenseSet<unsigned> Seen;
-  auto It = NameToIndex.find(Root);
-  if (It != NameToIndex.end()) {
-    CurrentLevel.push_back(It->second);
-    Seen.insert(It->second);
-  }
+  std::vector<IncludeStructure::HeaderID> CurrentLevel;
+  CurrentLevel.push_back(Root);
+  llvm::DenseSet<IncludeStructure::HeaderID> Seen;
+  Seen.insert(Root);
 
   // Each round of BFS traversal finds the next depth level.
-  std::vector<unsigned> PreviousLevel;
+  std::vector<IncludeStructure::HeaderID> PreviousLevel;
   for (unsigned Level = 1; !CurrentLevel.empty(); ++Level) {
     PreviousLevel.clear();
     PreviousLevel.swap(CurrentLevel);
     for (const auto &Parent : PreviousLevel) {
       for (const auto &Child : IncludeChildren.lookup(Parent)) {
         if (Seen.insert(Child).second) {
           CurrentLevel.push_back(Child);
-          const auto &Name = RealPathNames[Child];
-          // Can't include files if we don't have their real path.
-          if (!Name.empty())
-            Result[Name] = Level;
+          Result[Child] = Level;
         }
       }
     }
   }
   return Result;
 }
 
 void IncludeInserter::addExisting(const Inclusion &Inc) {

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

 #include "TestFS.h"
 #include "TestTU.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::AllOf;
 using ::testing::Contains;
 using ::testing::ElementsAre;
+using ::testing::IsEmpty;
 using ::testing::Not;
 using ::testing::UnorderedElementsAre;
+using ::testing::UnorderedElementsAreArray;
 
 class HeadersTest : public ::testing::Test {
 public:
   HeadersTest() {
     CDB.ExtraClangFlags = {SearchDirArg.c_str()};
     FS.Files[MainFile] = "";
     // Make sure directory sub/ exists.
     FS.Files[testPath("sub/EMPTY")] = "";
auto Clang = prepareCompilerInstance(
         llvm::MemoryBuffer::getMemBuffer(FS.Files[MainFile], MainFile),
         std::move(VFS), IgnoreDiags);
 
     EXPECT_FALSE(Clang->getFrontendOpts().Inputs.empty());
     return Clang;
   }
 
 protected:
+  IncludeStructure::HeaderID getID(StringRef Filename,
+                                   IncludeStructure &Includes) {
+    auto Entry = Clang->getSourceManager().getFileManager().getFile(Filename);
+    EXPECT_TRUE(Entry);
+    return Includes.getOrCreateID(*Entry);
+  }
+
   IncludeStructure collectIncludes() {
-    auto Clang = setupClang();
+    Clang = setupClang();
     PreprocessOnlyAction Action;
     EXPECT_TRUE(
         Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
     IncludeStructure Includes;
     Clang->getPreprocessor().addPPCallbacks(
         collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
     EXPECT_FALSE(Action.Execute());
     Action.EndSourceFile();
     return Includes;
   }
 
   // Calculates the include path, or returns "" on error or header should not be
   // inserted.
   std::string calculate(PathRef Original, PathRef Preferred = "",
                         const std::vector<Inclusion> &Inclusions = {}) {
-    auto Clang = setupClang();
+    Clang = setupClang();
     PreprocessOnlyAction Action;
     EXPECT_TRUE(
         Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
 
     if (Preferred.empty())
       Preferred = Original;
     auto ToHeaderFile = [](llvm::StringRef Header) {
       return HeaderFile{std::string(Header),
std::string calculate(PathRef Original, PathRef Preferred = "",
     if (!Inserter.shouldInsertInclude(Original, Inserted))
       return "";
     auto Path = Inserter.calculateIncludePath(Inserted, MainFile);
     Action.EndSourceFile();
     return Path.getValueOr("");
   }
 
   llvm::Optional<TextEdit> insert(llvm::StringRef VerbatimHeader) {
-    auto Clang = setupClang();
+    Clang = setupClang();
     PreprocessOnlyAction Action;
     EXPECT_TRUE(
         Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
 
     IncludeInserter Inserter(MainFile, /*Code=*/"", format::getLLVMStyle(),
                              CDB.getCompileCommand(MainFile)->Directory,
                              &Clang->getPreprocessor().getHeaderSearchInfo());
     auto Edit = Inserter.insert(VerbatimHeader);
     Action.EndSourceFile();
     return Edit;
   }
 
   MockFS FS;
   MockCompilationDatabase CDB;
   std::string MainFile = testPath("main.cpp");
   std::string Subdir = testPath("sub");
   std::string SearchDirArg = (llvm::Twine("-I") + Subdir).str();
   IgnoringDiagConsumer IgnoreDiags;
+  std::unique_ptr<CompilerInstance> Clang;
 };
 
 MATCHER_P(Written, Name, "") { return arg.Written == Name; }
 MATCHER_P(Resolved, Name, "") { return arg.Resolved == Name; }
 MATCHER_P(IncludeLine, N, "") { return arg.HashLine == N; }
 MATCHER_P(Directive, D, "") { return arg.Directive == D; }
 
 MATCHER_P2(Distance, File, D, "") {
-  if (arg.getKey() != File)
-    *result_listener << "file =" << arg.getKey().str();
-  if (arg.getValue() != D)
-    *result_listener << "distance =" << arg.getValue();
-  return arg.getKey() == File && arg.getValue() == D;
+  if (arg.getFirst() != File)
+    *result_listener << "file =" << static_cast<unsigned>(arg.getFirst());
+  if (arg.getSecond() != D)
+    *result_listener << "distance =" << arg.getSecond();
+  return arg.getFirst() == File && arg.getSecond() == D;
 }
 
 TEST_F(HeadersTest, CollectRewrittenAndResolved) {
   FS.Files[MainFile] = R"cpp(
 #include "sub/bar.h" // not shortest
 )cpp";
   std::string BarHeader = testPath("sub/bar.h");
   FS.Files[BarHeader] = "";
 
-  EXPECT_THAT(collectIncludes().MainFileIncludes,
+  auto Includes = collectIncludes();
+  EXPECT_THAT(Includes.MainFileIncludes,
               UnorderedElementsAre(
                   AllOf(Written("\"sub/bar.h\""), Resolved(BarHeader))));
-  EXPECT_THAT(collectIncludes().includeDepth(MainFile),
-              UnorderedElementsAre(Distance(MainFile, 0u),
-                                   Distance(testPath("sub/bar.h"), 1u)));
+  EXPECT_THAT(collectIncludes().includeDepth(getID(MainFile, Includes)),
+              UnorderedElementsAre(
+                  Distance(getID(MainFile, Includes), 0u),
+                  Distance(getID(testPath("sub/bar.h"), Includes), 1u)));
 }
 
 TEST_F(HeadersTest, OnlyCollectInclusionsInMain) {
   std::string BazHeader = testPath("sub/baz.h");
   FS.Files[BazHeader] = "";
   std::string BarHeader = testPath("sub/bar.h");
   FS.Files[BarHeader] = R"cpp(
 #include "baz.h"
 )cpp";
   FS.Files[MainFile] = R"cpp(
 #include "bar.h"
 )cpp";
+  auto Includes = collectIncludes();
   EXPECT_THAT(
       collectIncludes().MainFileIncludes,
       UnorderedElementsAre(AllOf(Written("\"bar.h\""), Resolved(BarHeader))));
-  EXPECT_THAT(collectIncludes().includeDepth(MainFile),
-              UnorderedElementsAre(Distance(MainFile, 0u),
-                                   Distance(testPath("sub/bar.h"), 1u),
-                                   Distance(testPath("sub/baz.h"), 2u)));
+  EXPECT_THAT(Includes.includeDepth(getID(MainFile, Includes)),
+              UnorderedElementsAre(
+                  Distance(getID(MainFile, Includes), 0u),
+                  Distance(getID(testPath("sub/bar.h"), Includes), 1u),
+                  Distance(getID(testPath("sub/baz.h"), Includes), 2u)));
   // includeDepth() also works for non-main files.
-  EXPECT_THAT(collectIncludes().includeDepth(testPath("sub/bar.h")),
-              UnorderedElementsAre(Distance(testPath("sub/bar.h"), 0u),
-                                   Distance(testPath("sub/baz.h"), 1u)));
+  EXPECT_THAT(
+      collectIncludes().includeDepth(getID(testPath("sub/bar.h"), Includes)),
+      UnorderedElementsAre(
+          Distance(getID(testPath("sub/bar.h"), Includes), 0u),
+          Distance(getID(testPath("sub/baz.h"), Includes), 1u)));
 }
 
 TEST_F(HeadersTest, PreambleIncludesPresentOnce) {
   // We use TestTU here, to ensure we use the preamble replay logic.
   // We're testing that the logic doesn't crash, and doesn't result in duplicate
   // includes. (We'd test more directly, but it's pretty well encapsulated!)
   auto TU = TestTU::withCode(R"cpp(
     #include "a.h"

 
 TEST_F(HeadersTest, UnResolvedInclusion) {
   FS.Files[MainFile] = R"cpp(
 #include "foo.h"
 )cpp";
 
   EXPECT_THAT(collectIncludes().MainFileIncludes,
               UnorderedElementsAre(AllOf(Written("\"foo.h\""), Resolved(""))));
-  EXPECT_THAT(collectIncludes().includeDepth(MainFile),
-              UnorderedElementsAre(Distance(MainFile, 0u)));
+  EXPECT_THAT(collectIncludes().IncludeChildren, IsEmpty());
+}
+
+TEST_F(HeadersTest, IncludedFilesGraph) {
+  FS.Files[MainFile] = R"cpp(
+#include "bar.h"
+#include "foo.h"
+)cpp";
+  std::string BarHeader = testPath("bar.h");
+  FS.Files[BarHeader] = "";
+  std::string FooHeader = testPath("foo.h");
+  FS.Files[FooHeader] = R"cpp(
+#include "bar.h"
+#include "baz.h"
+)cpp";
+  std::string BazHeader = testPath("baz.h");
+  FS.Files[BazHeader] = "";
+
+  auto Includes = collectIncludes();
+  llvm::DenseMap<IncludeStructure::HeaderID,
+                 SmallVector<IncludeStructure::HeaderID>>
+      Expected = {{getID(MainFile, Includes),
+                   {getID(BarHeader, Includes), getID(FooHeader, Includes)}},
+                  {getID(FooHeader, Includes),
+                   {getID(BarHeader, Includes), getID(BazHeader, Includes)}}};
+  EXPECT_EQ(Includes.IncludeChildren, Expected);
 }
 
 TEST_F(HeadersTest, IncludeDirective) {
   FS.Files[MainFile] = R"cpp(
 #include "foo.h"
 #import "foo.h"
 #include_next "foo.h"
 )cpp";

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

 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::AllOf;
+using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::ElementsAreArray;
 using ::testing::IsEmpty;
+using ::testing::UnorderedElementsAreArray;
 
 MATCHER_P(DeclNamed, Name, "") {
   if (NamedDecl *ND = dyn_cast<NamedDecl>(arg))
     if (ND->getName() == Name)
       return true;
   if (auto *Stream = result_listener->stream()) {
     llvm::raw_os_ostream OS(*Stream);
     arg->dump(OS);
TEST(ParsedASTTest, PatchesAdditionalIncludes) {
   TU.Code = "";
   StoreDiags Diags;
   MockFS FS;
   auto Inputs = TU.inputs(FS);
   auto CI = buildCompilerInvocation(Inputs, Diags);
   auto EmptyPreamble =
       buildPreamble(testPath("foo.cpp"), *CI, Inputs, true, nullptr);
   ASSERT_TRUE(EmptyPreamble);
-  EXPECT_THAT(EmptyPreamble->Includes.MainFileIncludes, testing::IsEmpty());
+  EXPECT_THAT(EmptyPreamble->Includes.MainFileIncludes, IsEmpty());
 
   // Now build an AST using empty preamble and ensure patched includes worked.
   TU.Code = ModifiedContents.str();
   Inputs = TU.inputs(FS);
   auto PatchedAST = ParsedAST::build(testPath("foo.cpp"), Inputs, std::move(CI),
                                      {}, EmptyPreamble);
   ASSERT_TRUE(PatchedAST);
   ASSERT_FALSE(PatchedAST->getDiagnostics());
 
   // Ensure source location information is correct, including resolved paths.
   EXPECT_THAT(PatchedAST->getIncludeStructure().MainFileIncludes,
               testing::Pointwise(
                   EqInc(), ExpectedAST.getIncludeStructure().MainFileIncludes));
-  auto StringMapToVector = [](const llvm::StringMap<unsigned> SM) {
-    std::vector<std::pair<std::string, unsigned>> Res;
-    for (const auto &E : SM)
-      Res.push_back({E.first().str(), E.second});
-    llvm::sort(Res);
-    return Res;
-  };
   // Ensure file proximity signals are correct.
-  EXPECT_EQ(StringMapToVector(PatchedAST->getIncludeStructure().includeDepth(
-                testPath("foo.cpp"))),
-            StringMapToVector(ExpectedAST.getIncludeStructure().includeDepth(
-                testPath("foo.cpp"))));
+  auto &FM = PatchedAST->getSourceManager().getFileManager();
+  // Copy so that we can use operator[] to get the children.
+  IncludeStructure Includes = PatchedAST->getIncludeStructure();
+  auto MainFE = FM.getFile(testPath("foo.cpp"));
+  ASSERT_TRUE(MainFE);
+  auto MainID = Includes.getID(*MainFE);
+  auto AuxFE = FM.getFile(testPath("sub/aux.h"));
+  ASSERT_TRUE(AuxFE);
+  auto AuxID = Includes.getID(*AuxFE);
+  EXPECT_THAT(Includes.IncludeChildren[*MainID], Contains(AuxID));
 }
 
 TEST(ParsedASTTest, PatchesDeletedIncludes) {
   TestTU TU;
   TU.Filename = "foo.cpp";
   TU.Code = "";
   auto ExpectedAST = TU.build();
 
TEST(ParsedASTTest, PatchesDeletedIncludes) {
   Inputs = TU.inputs(FS);
   auto PatchedAST = ParsedAST::build(testPath("foo.cpp"), Inputs, std::move(CI),
                                      {}, BaselinePreamble);
   ASSERT_TRUE(PatchedAST);
 
   // Ensure source location information is correct.
   EXPECT_THAT(PatchedAST->getIncludeStructure().MainFileIncludes,
               testing::Pointwise(
                   EqInc(), ExpectedAST.getIncludeStructure().MainFileIncludes));
-  auto StringMapToVector = [](const llvm::StringMap<unsigned> SM) {
-    std::vector<std::pair<std::string, unsigned>> Res;
-    for (const auto &E : SM)
-      Res.push_back({E.first().str(), E.second});
-    llvm::sort(Res);
-    return Res;
-  };
   // Ensure file proximity signals are correct.
-  EXPECT_EQ(StringMapToVector(PatchedAST->getIncludeStructure().includeDepth(
-                testPath("foo.cpp"))),
-            StringMapToVector(ExpectedAST.getIncludeStructure().includeDepth(
-                testPath("foo.cpp"))));
+  auto &FM = ExpectedAST.getSourceManager().getFileManager();
+  // Copy so that we can getOrCreateID().
+  IncludeStructure Includes = ExpectedAST.getIncludeStructure();
+  auto MainFE = FM.getFile(testPath("foo.cpp"));
+  ASSERT_TRUE(MainFE);
+  auto MainID = Includes.getOrCreateID(*MainFE);
+  auto &PatchedFM = PatchedAST->getSourceManager().getFileManager();
+  IncludeStructure PatchedIncludes = PatchedAST->getIncludeStructure();
+  auto PatchedMainFE = PatchedFM.getFile(testPath("foo.cpp"));
+  ASSERT_TRUE(PatchedMainFE);
+  auto PatchedMainID = PatchedIncludes.getOrCreateID(*PatchedMainFE);
+  EXPECT_EQ(Includes.includeDepth(MainID)[MainID],
+            PatchedIncludes.includeDepth(PatchedMainID)[PatchedMainID]);
 }
 
 // Returns Code guarded by #ifndef guards
 std::string guard(llvm::StringRef Code) {
   static int GuardID = 0;
   std::string GuardName = ("GUARD_" + llvm::Twine(++GuardID)).str();
   return llvm::formatv("#ifndef {0}\n#define {0}\n{1}\n#endif\n", GuardName,
                        Code);