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

Preparation for D108194.

Reviewed By: sammccall

Differential Revision: https://reviews.llvm.org/D110386

Author: kirillbobyrev

clang-tools-extra/clangd/CodeComplete.cpp

@@ -1379,14 +1379,16 @@ class CodeCompleteFlow {
       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);

clang-tools-extra/clangd/Headers.cpp

@@ -67,8 +67,9 @@ class RecordHeaders : public PPCallbacks {
         // 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);
     }
   }
 
@@ -154,49 +155,49 @@ collectIncludeStructureCallback(const SourceManager &SM,
   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;
         }
       }
     }

clang-tools-extra/clangd/Headers.h

@@ -12,6 +12,7 @@
 #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"
@@ -62,7 +63,7 @@ struct Inclusion {
 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 {
@@ -112,34 +113,42 @@ operator|=(IncludeGraphNode::SourceFlag &A, IncludeGraphNode::SourceFlag B) {
 // 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;
+  // 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;
 
-  // This updates IncludeDepth(), but not MainFileIncludes.
-  void recordInclude(llvm::StringRef IncludingName,
-                     llvm::StringRef IncludedName,
-                     llvm::StringRef IncludedRealName);
+  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.
@@ -205,4 +214,31 @@ class IncludeInserter {
 } // 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/unittests/HeadersTests.cpp

@@ -17,6 +17,7 @@
 #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"
@@ -29,8 +30,10 @@ 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:
@@ -64,8 +67,15 @@ class HeadersTest : public ::testing::Test {
   }
 
 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]));
@@ -81,7 +91,7 @@ class HeadersTest : public ::testing::Test {
   // 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]));
@@ -107,7 +117,7 @@ class HeadersTest : public ::testing::Test {
   }
 
   llvm::Optional<TextEdit> insert(llvm::StringRef VerbatimHeader) {
-    auto Clang = setupClang();
+    Clang = setupClang();
     PreprocessOnlyAction Action;
     EXPECT_TRUE(
         Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
@@ -126,6 +136,7 @@ class HeadersTest : public ::testing::Test {
   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; }
@@ -134,11 +145,11 @@ 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) {
@@ -148,12 +159,14 @@ TEST_F(HeadersTest, CollectRewrittenAndResolved) {
   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) {
@@ -166,17 +179,21 @@ TEST_F(HeadersTest, OnlyCollectInclusionsInMain) {
   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) {
@@ -202,8 +219,32 @@ TEST_F(HeadersTest, UnResolvedInclusion) {
 
   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) {