[clangd] Follow-up on rGdea48079b90d

Reviewed By: sammccall

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

Author: kirillbobyrev

clang-tools-extra/clangd/CodeComplete.cpp

@@ -1203,7 +1203,7 @@ bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
   loadMainFilePreambleMacros(Clang->getPreprocessor(), Input.Preamble);
   if (Includes)
     Clang->getPreprocessor().addPPCallbacks(
-        collectIncludeStructureCallback(Clang->getSourceManager(), Includes));
+        Includes->collect(Clang->getSourceManager()));
   if (llvm::Error Err = Action.Execute()) {
     log("Execute() failed when running codeComplete for {0}: {1}",
         Input.FileName, toString(std::move(Err)));
@@ -1380,7 +1380,7 @@ class CodeCompleteFlow {
       const auto &SM = Recorder->CCSema->getSourceManager();
       llvm::StringMap<SourceParams> ProxSources;
       auto MainFileID =
-          Includes.getID(SM.getFileEntryForID(SM.getMainFileID()), SM);
+          Includes.getID(SM.getFileEntryForID(SM.getMainFileID()));
       assert(MainFileID);
       for (auto &HeaderIDAndDepth : Includes.includeDepth(*MainFileID)) {
         auto &Source =

clang-tools-extra/clangd/Headers.cpp

@@ -67,8 +67,8 @@ class RecordHeaders : public PPCallbacks {
         // Treat as if included from the main file.
         IncludingFileEntry = SM.getFileEntryForID(MainFID);
       }
-      auto IncludingID = Out->getOrCreateID(IncludingFileEntry, SM),
-           IncludedID = Out->getOrCreateID(File, SM);
+      auto IncludingID = Out->getOrCreateID(IncludingFileEntry),
+           IncludedID = Out->getOrCreateID(File);
       Out->IncludeChildren[IncludingID].push_back(IncludedID);
     }
   }
@@ -150,16 +150,15 @@ llvm::SmallVector<llvm::StringRef, 1> getRankedIncludes(const Symbol &Sym) {
 }
 
 std::unique_ptr<PPCallbacks>
-collectIncludeStructureCallback(const SourceManager &SM,
-                                IncludeStructure *Out) {
-  return std::make_unique<RecordHeaders>(SM, Out);
+IncludeStructure::collect(const SourceManager &SM) {
+  MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
+  return std::make_unique<RecordHeaders>(SM, this);
 }
 
 llvm::Optional<IncludeStructure::HeaderID>
-IncludeStructure::getID(const FileEntry *Entry,
-                        const SourceManager &SM) const {
+IncludeStructure::getID(const FileEntry *Entry) const {
   // HeaderID of the main file is always 0;
-  if (SM.getMainFileID() == SM.translateFile(Entry)) {
+  if (Entry == MainFileEntry) {
     return static_cast<IncludeStructure::HeaderID>(0u);
   }
   auto It = UIDToIndex.find(Entry->getUniqueID());
@@ -169,12 +168,12 @@ IncludeStructure::getID(const FileEntry *Entry,
 }
 
 IncludeStructure::HeaderID
-IncludeStructure::getOrCreateID(const FileEntry *Entry,
-                                const SourceManager &SM) {
-  // Main file's FileID was not known at IncludeStructure creation time.
-  if (SM.getMainFileID() == SM.translateFile(Entry)) {
-    UIDToIndex[Entry->getUniqueID()] =
-        static_cast<IncludeStructure::HeaderID>(0u);
+IncludeStructure::getOrCreateID(const FileEntry *Entry) {
+  // Main file's FileEntry was not known at IncludeStructure creation time.
+  if (Entry == MainFileEntry) {
+    if (RealPathNames.front().empty())
+      RealPathNames.front() = MainFileEntry->tryGetRealPathName().str();
+    return MainFileID;
   }
   auto R = UIDToIndex.try_emplace(
       Entry->getUniqueID(),

clang-tools-extra/clangd/Headers.h

@@ -14,6 +14,7 @@
 #include "index/Symbol.h"
 #include "support/Logger.h"
 #include "support/Path.h"
+#include "clang/Basic/FileEntry.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Format/Format.h"
 #include "clang/Lex/HeaderSearch.h"
@@ -119,15 +120,22 @@ class IncludeStructure {
     RealPathNames.emplace_back();
   }
 
+  // Returns a PPCallback that visits all inclusions in the main file and
+  // populates the structure.
+  std::unique_ptr<PPCallbacks> collect(const SourceManager &SM);
+
+  void setMainFileEntry(const FileEntry *Entry) {
+    assert(Entry && Entry->isValid());
+    this->MainFileEntry = Entry;
+  }
+
   // 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 SourceManager &SM) const;
-  HeaderID getOrCreateID(const FileEntry *Entry,
-                         const SourceManager &SM);
+  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());
@@ -141,32 +149,33 @@ class IncludeStructure {
   // All transitive includes (absolute paths), with their minimum include depth.
   // Root --> 0, #included file --> 1, etc.
   // Root is the ID of the header being visited first.
-  // Usually it is getID(SM.getFileEntryForID(SM.getMainFileID()), SM).
-  llvm::DenseMap<HeaderID, unsigned> includeDepth(HeaderID Root) const;
+  llvm::DenseMap<HeaderID, unsigned>
+  includeDepth(HeaderID Root = MainFileID) const;
 
   // Maps HeaderID to the ids of the files included from it.
   llvm::DenseMap<HeaderID, SmallVector<HeaderID>> IncludeChildren;
 
   std::vector<Inclusion> MainFileIncludes;
 
+  // We reserve HeaderID(0) for the main file and will manually check for that
+  // in getID and getOrCreateID because the UniqueID is not stable when the
+  // content of the main file changes.
+  static const HeaderID MainFileID = HeaderID(0u);
+
 private:
+  // MainFileEntry will be used to check if the queried file is the main file
+  // or not.
+  const FileEntry *MainFileEntry = nullptr;
+
   std::vector<std::string> RealPathNames; // In HeaderID order.
-  // HeaderID maps the FileEntry::UniqueID to the internal representation.
+  // FileEntry::UniqueID is mapped to the internal representation (HeaderID).
   // 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
+  // builds is surprisingly hard. FileID is unavailable in InclusionDirective(),
+  // and RealPathName and UniqueID are not preserved in
   // the preamble.
-  //
-  // We reserve 0 to the main file and will manually check for that in getID
-  // and getOrCreateID because llvm::sys::fs::UniqueID is not stable when their
-  // content of the main file changes.
   llvm::DenseMap<llvm::sys::fs::UniqueID, HeaderID> UIDToIndex;
 };
 
-/// 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 {
 public:
@@ -228,7 +237,7 @@ class IncludeInserter {
 
 namespace llvm {
 
-// Support Tokens as DenseMap keys.
+// Support HeaderIDs as DenseMap keys.
 template <> struct DenseMapInfo<clang::clangd::IncludeStructure::HeaderID> {
   static inline clang::clangd::IncludeStructure::HeaderID getEmptyKey() {
     return static_cast<clang::clangd::IncludeStructure::HeaderID>(
@@ -251,30 +260,6 @@ template <> struct DenseMapInfo<clang::clangd::IncludeStructure::HeaderID> {
   }
 };
 
-// Support Tokens as DenseMap keys.
-template <> struct DenseMapInfo<llvm::sys::fs::UniqueID> {
-  static inline llvm::sys::fs::UniqueID getEmptyKey() {
-    auto EmptyKey = DenseMapInfo<std::pair<unsigned, unsigned>>::getEmptyKey();
-    return {EmptyKey.first, EmptyKey.second};
-  }
-
-  static inline llvm::sys::fs::UniqueID getTombstoneKey() {
-    auto TombstoneKey =
-        DenseMapInfo<std::pair<unsigned, unsigned>>::getTombstoneKey();
-    return {TombstoneKey.first, TombstoneKey.second};
-  }
-
-  static unsigned getHashValue(const llvm::sys::fs::UniqueID &Tag) {
-    return hash_value(
-        std::pair<unsigned, unsigned>(Tag.getDevice(), Tag.getFile()));
-  }
-
-  static bool isEqual(const llvm::sys::fs::UniqueID &LHS,
-                      const llvm::sys::fs::UniqueID &RHS) {
-    return LHS == RHS;
-  }
-};
-
 } // namespace llvm
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H

clang-tools-extra/clangd/ParsedAST.cpp

@@ -439,7 +439,7 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   // Otherwise we would collect the replayed includes again...
   // (We can't *just* use the replayed includes, they don't have Resolved path).
   Clang->getPreprocessor().addPPCallbacks(
-      collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
+      Includes.collect(Clang->getSourceManager()));
   // Copy over the macros in the preamble region of the main file, and combine
   // with non-preamble macros below.
   MainFileMacros Macros;

clang-tools-extra/clangd/Preamble.cpp

@@ -104,7 +104,7 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
            "SourceMgr and LangOpts must be set at this point");
 
     return std::make_unique<PPChainedCallbacks>(
-        collectIncludeStructureCallback(*SourceMgr, &Includes),
+        Includes.collect(*SourceMgr),
         std::make_unique<PPChainedCallbacks>(
             std::make_unique<CollectMainFileMacros>(*SourceMgr, Macros),
             collectPragmaMarksCallback(*SourceMgr, Marks)));
@@ -285,7 +285,7 @@ scanPreamble(llvm::StringRef Contents, const tooling::CompileCommand &Cmd) {
   const auto &SM = Clang->getSourceManager();
   Preprocessor &PP = Clang->getPreprocessor();
   IncludeStructure Includes;
-  PP.addPPCallbacks(collectIncludeStructureCallback(SM, &Includes));
+  PP.addPPCallbacks(Includes.collect(SM));
   ScannedPreamble SP;
   SP.Bounds = Bounds;
   PP.addPPCallbacks(

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

@@ -72,7 +72,7 @@ class HeadersTest : public ::testing::Test {
     auto &SM = Clang->getSourceManager();
     auto Entry = SM.getFileManager().getFile(Filename);
     EXPECT_TRUE(Entry);
-    return Includes.getOrCreateID(*Entry, SM);
+    return Includes.getOrCreateID(*Entry);
   }
 
   IncludeStructure collectIncludes() {
@@ -82,7 +82,7 @@ class HeadersTest : public ::testing::Test {
         Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
     IncludeStructure Includes;
     Clang->getPreprocessor().addPPCallbacks(
-        collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
+        Includes.collect(Clang->getSourceManager()));
     EXPECT_FALSE(Action.Execute());
     Action.EndSourceFile();
     return Includes;
@@ -180,9 +180,9 @@ TEST_F(HeadersTest, OnlyCollectInclusionsInMain) {
 #include "bar.h"
 )cpp";
   auto Includes = collectIncludes();
-  EXPECT_THAT(Includes.MainFileIncludes,
-              UnorderedElementsAre(
-                  AllOf(Written("\"bar.h\""), Resolved(BarHeader))));
+  EXPECT_THAT(
+      Includes.MainFileIncludes,
+      UnorderedElementsAre(AllOf(Written("\"bar.h\""), Resolved(BarHeader))));
   EXPECT_THAT(Includes.includeDepth(getID(MainFile, Includes)),
               UnorderedElementsAre(Distance(getID(MainFile, Includes), 0u),
                                    Distance(getID(BarHeader, Includes), 1u),

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

@@ -516,10 +516,10 @@ TEST(ParsedASTTest, PatchesAdditionalIncludes) {
   IncludeStructure Includes = PatchedAST->getIncludeStructure();
   auto MainFE = FM.getFile(testPath("foo.cpp"));
   ASSERT_TRUE(MainFE);
-  auto MainID = Includes.getID(*MainFE, SM);
+  auto MainID = Includes.getID(*MainFE);
   auto AuxFE = FM.getFile(testPath("sub/aux.h"));
   ASSERT_TRUE(AuxFE);
-  auto AuxID = Includes.getID(*AuxFE, SM);
+  auto AuxID = Includes.getID(*AuxFE);
   EXPECT_THAT(Includes.IncludeChildren[*MainID], Contains(*AuxID));
 }
 
@@ -560,12 +560,12 @@ TEST(ParsedASTTest, PatchesDeletedIncludes) {
   IncludeStructure Includes = ExpectedAST.getIncludeStructure();
   auto MainFE = FM.getFile(testPath("foo.cpp"));
   ASSERT_TRUE(MainFE);
-  auto MainID = Includes.getOrCreateID(*MainFE, SM);
+  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, SM);
+  auto PatchedMainID = PatchedIncludes.getOrCreateID(*PatchedMainFE);
   EXPECT_EQ(Includes.includeDepth(MainID)[MainID],
             PatchedIncludes.includeDepth(PatchedMainID)[PatchedMainID]);
 }

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

@@ -83,7 +83,7 @@ collectPatchedIncludes(llvm::StringRef ModifiedContents,
   }
   IncludeStructure Includes;
   Clang->getPreprocessor().addPPCallbacks(
-      collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
+      Includes.collect(Clang->getSourceManager()));
   if (llvm::Error Err = Action.Execute()) {
     ADD_FAILURE() << "failed to execute action: " << std::move(Err);
     return {};

llvm/include/llvm/Support/FileSystem/UniqueID.h

@@ -14,7 +14,9 @@
 #ifndef LLVM_SUPPORT_FILESYSTEM_UNIQUEID_H
 #define LLVM_SUPPORT_FILESYSTEM_UNIQUEID_H
 
+#include "llvm/ADT/DenseMap.h"
 #include <cstdint>
+#include <utility>
 
 namespace llvm {
 namespace sys {
@@ -47,6 +49,31 @@ class UniqueID {
 
 } // end namespace fs
 } // end namespace sys
+
+// Support UniqueIDs as DenseMap keys.
+template <> struct DenseMapInfo<llvm::sys::fs::UniqueID> {
+  static inline llvm::sys::fs::UniqueID getEmptyKey() {
+    auto EmptyKey = DenseMapInfo<std::pair<unsigned, unsigned>>::getEmptyKey();
+    return {EmptyKey.first, EmptyKey.second};
+  }
+
+  static inline llvm::sys::fs::UniqueID getTombstoneKey() {
+    auto TombstoneKey =
+        DenseMapInfo<std::pair<unsigned, unsigned>>::getTombstoneKey();
+    return {TombstoneKey.first, TombstoneKey.second};
+  }
+
+  static unsigned getHashValue(const llvm::sys::fs::UniqueID &Tag) {
+    return hash_value(
+        std::pair<unsigned, unsigned>(Tag.getDevice(), Tag.getFile()));
+  }
+
+  static bool isEqual(const llvm::sys::fs::UniqueID &LHS,
+                      const llvm::sys::fs::UniqueID &RHS) {
+    return LHS == RHS;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_SUPPORT_FILESYSTEM_UNIQUEID_H