[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–130

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?

139–141

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.

151

this isn't abstract :-)

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

152

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

1382–1398

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

1383

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

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

138–139

comment is out of date

140

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

144

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

172

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.

174

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

189

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

202–203

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

clang-tools-extra/clangd/Headers.h

116–130

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

70

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?

236

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);

238

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

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

511–513

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

568–571

and here

[kbobyrev]

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

clang-tools-extra/clangd/Headers.cpp

202–203

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

236

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

202–203

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

236

== 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

202–203

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

236

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 IncludeStructureID =
+          Includes.getID(SM.getFileEntryForID(SM.getMainFileID()));
+      if (IncludeStructureID) {
+        for (auto &HeaderIDAndDepth :
+             Includes.includeDepth(*IncludeStructureID)) {
+          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;
         }
+      } else {
+        elog("Could not get includeDepth in CodeComplete: {0}",
+             llvm::toString(IncludeStructureID.takeError()));
+      }
       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;
+  // 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".
+  // HeaderID is mapping the FileEntry::Name to the internal representation:
+  // RealPathName stores header names in the order they were "discovered".
+  // Main file HeaderID is always 0, hence RealPathName.front() has its name.
+  using HeaderID = unsigned;
+
+  llvm::Expected<HeaderID> getID(const FileEntry *Entry) const;
+  HeaderID getOrCreateID(const FileEntry *Entry);
+
+  StringRef getRealPath(HeaderID ID) const {
+    return ID < RealPathNames.size() ? RealPathNames[ID] : StringRef();
+  }
 
   // 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;
+  // Usually it should be
+  // getFile(SM.getFileEntryForID(SM.getMainFileID())->getName()).
+  llvm::DenseMap<HeaderID, unsigned> includeDepth(HeaderID Root) const;
 
-  // This updates IncludeDepth(), but not MainFileIncludes.
-  void recordInclude(llvm::StringRef IncludingName,
-                     llvm::StringRef IncludedName,
-                     llvm::StringRef IncludedRealName);
+  // This updates IncludeChildren, but not MainFileIncludes.
+  void recordInclude(HeaderID Including, HeaderID Included);
+
+  // Maps HeaderID to the ids of the files included from it.
+  llvm::DenseMap<HeaderID, SmallVector<HeaderID>> IncludeChildren;
+
+  std::vector<Inclusion> MainFileIncludes;
 
 private:
-  // 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;
+  std::vector<std::string> RealPathNames; // In HeaderID order.
+  // Maps filenames to corresponding HeaderIDs.
+  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 {

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::Expected<IncludeStructure::HeaderID>
+IncludeStructure::getID(const FileEntry *Entry) const {
+  auto It = NameToIndex.find(Entry->getName());
+  if (It == NameToIndex.end())
+    return error("{0} is not in IncludeStructure", Entry->getName());
+  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(), RealPathNames.size());
   if (R.second)
     RealPathNames.emplace_back();
-  return R.first->getValue();
+  IncludeStructure::HeaderID Result = R.first->getValue();
+  auto RealPathName = Entry->tryGetRealPathName();
+  if (!RealPathName.empty() && RealPathNames[Result].empty())
+    RealPathNames[Result] = RealPathName.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;
+  if (Root >= RealPathNames.size()) {
+    elog("Requested includeDepth for {0} which doesn't exist IncludeStructure",
+         Root);
+    return Result;
+  }
   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);
-  }
+  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;
   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;
+          if (!RealPathNames[Child].empty())
+            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::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) {
+    auto Includes = collectIncludes();
+    auto Entry = Clang->getSourceManager().getFileManager().getFile(Filename);
+    EXPECT_TRUE(Entry);
+    llvm::Expected<IncludeStructure::HeaderID> EntryID = Includes.getID(*Entry);
+    EXPECT_TRUE(bool(EntryID)) << EntryID.takeError();
+    return *EntryID;
+  }
+
   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 =" << 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,
               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)),
+              UnorderedElementsAre(Distance(getID(MainFile), 0u),
+                                   Distance(getID(testPath("sub/bar.h")), 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";
   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(collectIncludes().includeDepth(getID(MainFile)),
+              UnorderedElementsAre(Distance(getID(MainFile), 0u),
+                                   Distance(getID(testPath("sub/bar.h")), 1u),
+                                   Distance(getID(testPath("sub/baz.h")), 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"))),
+              UnorderedElementsAre(Distance(getID(testPath("sub/bar.h")), 0u),
+                                   Distance(getID(testPath("sub/baz.h")), 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_TRUE(collectIncludes().IncludeChildren.empty());
+}
+
+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();
+  EXPECT_THAT(Includes.IncludeChildren[getID(MainFile)],
+              UnorderedElementsAreArray({getID(FooHeader), getID(BarHeader)}));
+  EXPECT_TRUE(Includes.IncludeChildren[getID(BarHeader)].empty());
+  EXPECT_THAT(Includes.IncludeChildren[getID(FooHeader)],
+              UnorderedElementsAreArray({getID(BarHeader), getID(BazHeader)}));
+  EXPECT_TRUE(Includes.IncludeChildren[getID(BazHeader)].empty());
 }
 
 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

 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::AllOf;
 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);
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) {
+  auto Flatten = [](const llvm::DenseMap<IncludeStructure::HeaderID, unsigned>
+                        &IncludeDepth,
+                    const IncludeStructure &Structure) {
     std::vector<std::pair<std::string, unsigned>> Res;
-    for (const auto &E : SM)
-      Res.push_back({E.first().str(), E.second});
+    for (const auto &E : IncludeDepth)
+      Res.emplace_back(Structure.getRealPath(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 PatchedEntry = PatchedAST->getSourceManager().getFileManager().getFile(
+      testPath("foo.cpp"));
+  ASSERT_TRUE(PatchedEntry);
+  auto PatchedEntryID = PatchedAST->getIncludeStructure().getID(*PatchedEntry);
+  ASSERT_TRUE(bool(PatchedEntryID)) << PatchedEntryID.takeError();
+  auto ExpectedEntry = PatchedAST->getSourceManager().getFileManager().getFile(
+      testPath("foo.cpp"));
+  ASSERT_TRUE(ExpectedEntry);
+  auto ExpectedEntryID = PatchedAST->getIncludeStructure().getID(*PatchedEntry);
+  ASSERT_TRUE(bool(ExpectedEntryID)) << ExpectedEntryID.takeError();
+  EXPECT_EQ(
+      Flatten(PatchedAST->getIncludeStructure().includeDepth(*PatchedEntryID),
+              PatchedAST->getIncludeStructure()),
+      Flatten(ExpectedAST.getIncludeStructure().includeDepth(*ExpectedEntryID),
+              ExpectedAST.getIncludeStructure()));
 }
 
 TEST(ParsedASTTest, PatchesDeletedIncludes) {
   TestTU TU;
   TU.Filename = "foo.cpp";
   TU.Code = "";
   auto ExpectedAST = TU.build();
 
TEST(ParsedASTTest, PatchesDeletedIncludes) {
   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) {
+  // Ensure file proximity signals are correct.
+  auto Flatten = [](const llvm::DenseMap<IncludeStructure::HeaderID, unsigned>
+                        &IncludeDepth,
+                    const IncludeStructure &Structure) {
     std::vector<std::pair<std::string, unsigned>> Res;
-    for (const auto &E : SM)
-      Res.push_back({E.first().str(), E.second});
+    for (const auto &E : IncludeDepth)
+      Res.emplace_back(Structure.getRealPath(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 PatchedEntry = PatchedAST->getSourceManager().getFileManager().getFile(
+      testPath("foo.cpp"));
+  ASSERT_TRUE(PatchedEntry);
+  auto PatchedEntryID = PatchedAST->getIncludeStructure().getID(*PatchedEntry);
+  ASSERT_TRUE(bool(PatchedEntryID)) << PatchedEntryID.takeError();
+  EXPECT_EQ(
+      Flatten(PatchedAST->getIncludeStructure().includeDepth(*PatchedEntryID),
+              PatchedAST->getIncludeStructure())
+          .size(),
+      1u);
+  EXPECT_TRUE(
+      Flatten(ExpectedAST.getIncludeStructure().includeDepth(*PatchedEntryID),
+              ExpectedAST.getIncludeStructure())
+          .empty());
 }
 
 // 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);