Move the BySpelling map to IncludeStructure.

[kadircet]

clang-tools-extra/clangd/Headers.h

164

instead of building this on-demand, what about building it as we're collecting directives around https://github.com/llvm/llvm-project/blob/main/clang-tools-extra/clangd/Headers.cpp#L52 ?

afterwards we can just have a lookup function exposed here, that returns an ArrayRef<HeaderID> ?

[VitaNuo]

Address review comments.

[VitaNuo]

Expose API for lookup instead of retrieving the whole map.

[VitaNuo]

Thanks for review!

clang-tools-extra/clangd/Headers.h

164

That's a great idea. I was thinking about this, but didn't know where to put it, so I thought if it was possible, you'd point out in review ;-)
Thanks!

[kadircet]

clang-tools-extra/clangd/Headers.cpp

77

right now we're only storing "resolved" includes from the main file and not all, this is creating a discrepancy between the view one gets through MainFileIncludes and through this map.
in addition to that only storing HeaderID gets the job done for IncludeCleaner, but won't really help anyone that wants to match main file includes apart from that (there's no easy way to go from a HeaderID to an Inclusion).

so instead of storing the HeaderID in the map values, we can actually store indexes into MainFileIncludes. later on during the lookup, we can build a SmallVector<Inclusion *> that contains pointers to the relevant includes. WDYT?

clang-tools-extra/clangd/Headers.h

164

what about renaming lookupHeaderIDsBySpelling to mainFileIncludesWithSpelling. with a comment explaining Returns includes inside the main file with the given Spelling. Spelling should include brackets or quotes, e.g. <foo>.

[kadircet]

oh forgot to mention, could you also please add some tests into llvm-project/clang-tools-extra/clangd/unittests/HeadersTests.cpp ?

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the review!

clang-tools-extra/clangd/Headers.cpp

77

Ok sure.

[kadircet]

clang-tools-extra/clangd/Headers.h

165

we're still returning just the HeaderID, the suggestion was to return llvm::SmallVector<Inclusion*> so that applications can work with other information available in the Inclusion. we also won't have any requirements around include being resolved that way. the application should figure out what to do if HeaderID is missing.

also can you move this function body to source file instead?

[VitaNuo]

Address comments

[VitaNuo]

Thanks for the comments!

[VitaNuo]

Move to source file.

[kadircet]

clang-tools-extra/clangd/Headers.cpp

303

nit:

llvm::SmallVector<Inclusion*> Includes;
for (auto Idx : MainFileIncludesBySpelling.lookup(Spelling))
  Includes.push_back(&MainFileIncludes[Idx]);
return Includes;

clang-tools-extra/clangd/Headers.h

165

almost, but not right. we're returning copies of Inclusions here, not pointers to Inclusions inside the main file. the suggested signature was llvm::SmallVector<Inclusion*> any reason for returning by value?

clang-tools-extra/clangd/IncludeCleaner.cpp

516

nit:

if(!Inc.HeaderID.has_value())
  continue;
Used.insert(static_cast<IncludeStructure::HeaderID>(*Inc.HeaderID));

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the comments!

clang-tools-extra/clangd/Headers.cpp

303

But lookup will return 0 if it doesn't find the spelling in the map (default value of int). And at the same time 0 is a valid index. Isn't that just wrong?

clang-tools-extra/clangd/Headers.h

165

Sure, thanks. No particular reason.

[kadircet]

clang-tools-extra/clangd/Headers.cpp

303

MainFileIncludesBySpelling is a map with values of type vector, not int. hence it'll be an empty vector, not 0.

[VitaNuo]

Address review comments.

[VitaNuo]

clang-tools-extra/clangd/Headers.cpp

303

Oh right. Sorry I guess I got confused.

[kadircet]

thanks (both for the patch and for bearing with me)!

clang-tools-extra/clangd/Headers.h

StringRef getRealPath(HeaderID ID) const {
     assert(static_cast<unsigned>(ID) <= RealPathNames.size());
     return RealPathNames[static_cast<unsigned>(ID)];
   }
 
   bool isSelfContained(HeaderID ID) const {
     return !NonSelfContained.contains(ID);
   }
 
-  bool hasIWYUExport(HeaderID ID) const {
-    return HasIWYUExport.contains(ID);
-  }
+  bool hasIWYUExport(HeaderID ID) const { return HasIWYUExport.contains(ID); }
 
   // Return all transitively reachable files.
   llvm::ArrayRef<std::string> allHeaders() const { return RealPathNames; }
 
+  // Returns includes inside the main file with the given spelling.
+  // Spelling should include brackets or quotes, e.g. <foo>.
+  llvm::SmallVector<const Inclusion *>
+  mainFileIncludesWithSpelling(llvm::StringRef Spelling) const;
+
   // 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 the ID of the header being visited first.
   llvm::DenseMap<HeaderID, unsigned>
   includeDepth(HeaderID Root = MainFileID) const;
 
   // Maps HeaderID to the ids of the files included from it.
private:
   // the preamble.
   llvm::DenseMap<llvm::sys::fs::UniqueID, HeaderID> UIDToIndex;
   // Contains HeaderIDs of all non self-contained entries in the
   // IncludeStructure.
   llvm::DenseSet<HeaderID> NonSelfContained;
   // Contains a set of headers that have either "IWYU pragma: export" or "IWYU
   // pragma: begin_exports".
   llvm::DenseSet<HeaderID> HasIWYUExport;
+
+  // Maps written includes to indices in MainFileInclude for easier lookup by
+  // spelling.
+  llvm::StringMap<llvm::SmallVector<unsigned>> MainFileIncludesBySpelling;
 };
 
 // Calculates insertion edit for including a new header in a file.
 class IncludeInserter {
 public:
   // If \p HeaderSearchInfo is nullptr (e.g. when compile command is
   // infeasible), this will only try to insert verbatim headers, and
   // include path of non-verbatim header will not be shortened.

clang-tools-extra/clangd/Headers.cpp

 #include "SourceCode.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Inclusions/HeaderAnalysis.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
 #include <cstring>
 #include <optional>
 
 namespace clang {
 namespace clangd {
 
if (isInsideMainFile(HashLoc, SM)) {
         if (IsAngled)
           if (auto StdlibHeader = tooling::stdlib::Header::named(Inc.Written)) {
             auto &IDs = Out->StdlibHeaders[*StdlibHeader];
             // Few physical files for one stdlib header name, linear scan is ok.
             if (!llvm::is_contained(IDs, HID))
               IDs.push_back(HID);
           }
       }
+      Out->MainFileIncludesBySpelling.try_emplace(Inc.Written)
+          .first->second.push_back(Out->MainFileIncludes.size() - 1);
     }
 
     // Record include graph (not just for main-file includes)
     if (File) {
       auto IncludingFileEntry = SM.getFileEntryRefForID(SM.getFileID(HashLoc));
       if (!IncludingFileEntry) {
         assert(SM.getBufferName(HashLoc).startswith("<") &&
                "Expected #include location to be a file or <built-in>");
for (const auto &Parent : PreviousLevel) {
           Result[Child] = Level;
         }
       }
     }
   }
   return Result;
 }
 
+llvm::SmallVector<const Inclusion *>
+IncludeStructure::mainFileIncludesWithSpelling(llvm::StringRef Spelling) const {
+  llvm::SmallVector<const Inclusion *> Includes;
+  for (auto Idx : MainFileIncludesBySpelling.lookup(Spelling))
+    Includes.push_back(&MainFileIncludes[Idx]);
+  return Includes;
+}
+
 void IncludeInserter::addExisting(const Inclusion &Inc) {
   IncludedHeaders.insert(Inc.Written);
   if (!Inc.Resolved.empty())
     IncludedHeaders.insert(Inc.Resolved);
 }
 
 /// FIXME(ioeric): we might not want to insert an absolute include path if the
 /// path is not shortened.

clang-tools-extra/clangd/IncludeCleaner.cpp

bool VisitCXXConstructExpr(CXXConstructExpr *CCE) {
     add(CCE->getConstructor());
     return true;
   }
 
   bool VisitTemplateSpecializationType(TemplateSpecializationType *TST) {
     // Using templateName case is handled by the override TraverseTemplateName.
     if (TST->getTemplateName().getKind() == TemplateName::UsingTemplate)
       return true;
-    add(TST->getAsCXXRecordDecl());                  // Specialization
+    add(TST->getAsCXXRecordDecl()); // Specialization
     return true;
   }
 
   // There is no VisitTemplateName in RAV, thus we override the Traverse version
   // to handle the Using TemplateName case.
   bool TraverseTemplateName(TemplateName TN) {
     VisitTemplateName(TN);
     return Base::TraverseTemplateName(TN);
std::vector<const Inclusion *> computeUnusedIncludes(ParsedAST &AST) {
   auto Refs = findReferencedLocations(AST);
   auto ReferencedFiles =
       findReferencedFiles(Refs, AST.getIncludeStructure(),
                           AST.getCanonicalIncludes(), AST.getSourceManager());
   auto ReferencedHeaders =
       translateToHeaderIDs(ReferencedFiles, AST.getIncludeStructure(), SM);
   return getUnused(AST, ReferencedHeaders, ReferencedFiles.SpelledUmbrellas);
 }
-std::vector<const Inclusion *> computeUnusedIncludesExperimental(ParsedAST &AST) {
-   const auto &SM = AST.getSourceManager();
-   const auto &Includes = AST.getIncludeStructure();
-   // FIXME: this map should probably be in IncludeStructure.
-   llvm::StringMap<llvm::SmallVector<IncludeStructure::HeaderID>> BySpelling;
-   for (const auto &Inc : Includes.MainFileIncludes) {
-    if (Inc.HeaderID)
-      BySpelling.try_emplace(Inc.Written)
-          .first->second.push_back(
-              static_cast<IncludeStructure::HeaderID>(*Inc.HeaderID));
-   }
-   // FIXME: !!this is a hacky way to collect macro references.
-   std::vector<include_cleaner::SymbolReference> Macros;
-    auto& PP = AST.getPreprocessor();
-   for (const syntax::Token &Tok :
-        AST.getTokens().spelledTokens(SM.getMainFileID())) {
+std::vector<const Inclusion *>
+computeUnusedIncludesExperimental(ParsedAST &AST) {
+
+  auto &PP = AST.getPreprocessor();
     auto Macro = locateMacroAt(Tok, PP);
     if (!Macro)
       continue;
     if (auto DefLoc = Macro->Info->getDefinitionLoc(); DefLoc.isValid())
       Macros.push_back(
           {Tok.location(),
            include_cleaner::Macro{/*Name=*/PP.getIdentifierInfo(Tok.text(SM)),
                                   DefLoc},
            include_cleaner::RefType::Explicit});
-   }
-   llvm::DenseSet<IncludeStructure::HeaderID> Used;
-   include_cleaner::walkUsed(
-       AST.getLocalTopLevelDecls(), /*MacroRefs=*/Macros,
-       AST.getPragmaIncludes(), SM,
-       [&](const include_cleaner::SymbolReference &Ref,
-           llvm::ArrayRef<include_cleaner::Header> Providers) {
-         for (const auto &H : Providers) {
-           switch (H.kind()) {
-           case include_cleaner::Header::Physical:
-             if (auto HeaderID = Includes.getID(H.physical()))
-               Used.insert(*HeaderID);
-             break;
-           case include_cleaner::Header::Standard:
-             for (auto HeaderID : Includes.StdlibHeaders.lookup(H.standard()))
-               Used.insert(HeaderID);
-             break;
-           case include_cleaner::Header::Verbatim:
-             for (auto HeaderID : BySpelling.lookup(H.verbatim()))
-               Used.insert(HeaderID);
-             break;
-           }
-         }
-       });
-   return getUnused(AST, Used, /*ReferencedPublicHeaders*/{});
+            for (auto *Inc :
+                 Includes.mainFileIncludesWithSpelling(H.verbatim())) {
+              if (!Inc->HeaderID.has_value())
+                continue;
+              IncludeStructure::HeaderID ID =
+                  static_cast<IncludeStructure::HeaderID>(*Inc->HeaderID);
+              Used.insert(ID);
+            }
+  return getUnused(AST, Used, /*ReferencedPublicHeaders*/ {});
 }
 
 std::vector<Diag> issueUnusedIncludesDiagnostics(ParsedAST &AST,
                                                  llvm::StringRef Code) {
   const Config &Cfg = Config::current();
   if (Cfg.Diagnostics.UnusedIncludes == Config::UnusedIncludesPolicy::None ||
       Cfg.Diagnostics.SuppressAll ||
       Cfg.Diagnostics.Suppress.contains("unused-includes"))

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

EXPECT_THAT(Includes.includeDepth(getID(MainFile, Includes)),
                                    Distance(getID(BarHeader, Includes), 1u),
                                    Distance(getID(BazHeader, Includes), 2u)));
   // includeDepth() also works for non-main files.
   EXPECT_THAT(Includes.includeDepth(getID(BarHeader, Includes)),
               UnorderedElementsAre(Distance(getID(BarHeader, Includes), 0u),
                                    Distance(getID(BazHeader, Includes), 1u)));
 }
 
+TEST_F(HeadersTest, CacheBySpellingIsBuiltForMainInclusions) {
+  std::string FooHeader = testPath("foo.h");
+  FS.Files[FooHeader] = R"cpp(
+  void foo();
+)cpp";
+  std::string BarHeader = testPath("bar.h");
+  FS.Files[BarHeader] = R"cpp(
+  void bar();
+)cpp";
+  std::string BazHeader = testPath("baz.h");
+  FS.Files[BazHeader] = R"cpp(
+  void baz();
+)cpp";
+  FS.Files[MainFile] = R"cpp(
+#include "foo.h"
+#include "bar.h"
+#include "baz.h"
+)cpp";
+  auto Includes = collectIncludes();
+  EXPECT_THAT(Includes.MainFileIncludes,
+              UnorderedElementsAre(written("\"foo.h\""), written("\"bar.h\""),
+                                   written("\"baz.h\"")));
+  EXPECT_THAT(Includes.mainFileIncludesWithSpelling("\"foo.h\""),
+              UnorderedElementsAre(&Includes.MainFileIncludes[0]));
+  EXPECT_THAT(Includes.mainFileIncludesWithSpelling("\"bar.h\""),
+              UnorderedElementsAre(&Includes.MainFileIncludes[1]));
+  EXPECT_THAT(Includes.mainFileIncludesWithSpelling("\"baz.h\""),
+              UnorderedElementsAre(&Includes.MainFileIncludes[2]));
+}
+
 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"
 
     #include "a.h"