[clangd] Perform self-containedness check at EOF (#75965)

Header gurads are not detected until we hit EOF. Make sure we postpone
any such detection until then.

3 files changed, 59 insertions, 18 deletions

diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index cf6102d..7ef4b15 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -826,22 +826,8 @@ void SymbolCollector::setIncludeLocation(const Symbol &S, SourceLocation DefLoc,
   // We update providers for a symbol with each occurence, as SymbolCollector
   // might run while parsing, rather than at the end of a translation unit.
   // Hence we see more and more redecls over time.
-  auto [It, Inserted] = SymbolProviders.try_emplace(S.ID);
-  auto Headers =
+  SymbolProviders[S.ID] =
       include_cleaner::headersForSymbol(Sym, SM, Opts.PragmaIncludes);
-  if (Headers.empty())
-    return;
-
-  auto *HeadersIter = Headers.begin();
-  include_cleaner::Header H = *HeadersIter;
-  while (HeadersIter != Headers.end() &&
-         H.kind() == include_cleaner::Header::Physical &&
-         !tooling::isSelfContainedHeader(H.physical(), SM,
-                                         PP->getHeaderSearchInfo())) {
-    H = *HeadersIter;
-    HeadersIter++;
-  }
-  It->second = H;
 }
 
 llvm::StringRef getStdHeader(const Symbol *S, const LangOptions &LangOpts) {
@@ -889,7 +875,7 @@ void SymbolCollector::finish() {
   llvm::DenseMap<include_cleaner::Header, std::string> HeaderSpelling;
   // Fill in IncludeHeaders.
   // We delay this until end of TU so header guards are all resolved.
-  for (const auto &[SID, OptionalProvider] : SymbolProviders) {
+  for (const auto &[SID, Providers] : SymbolProviders) {
     const Symbol *S = Symbols.find(SID);
     if (!S)
       continue;
@@ -931,9 +917,27 @@ void SymbolCollector::finish() {
       continue;
     }
 
-    assert(Directives == Symbol::Include);
     // For #include's, use the providers computed by the include-cleaner
     // library.
+    assert(Directives == Symbol::Include);
+    // Ignore providers that are not self-contained, this is especially
+    // important for symbols defined in the main-file. We want to prefer the
+    // header, if possible.
+    // TODO: Limit this to specifically ignore main file, when we're indexing a
+    // non-header file?
+    auto SelfContainedProvider =
+        [this](llvm::ArrayRef<include_cleaner::Header> Providers)
+        -> std::optional<include_cleaner::Header> {
+      for (const auto &H : Providers) {
+        if (H.kind() != include_cleaner::Header::Physical)
+          return H;
+        if (tooling::isSelfContainedHeader(H.physical(), PP->getSourceManager(),
+                                           PP->getHeaderSearchInfo()))
+          return H;
+      }
+      return std::nullopt;
+    };
+    const auto OptionalProvider = SelfContainedProvider(Providers);
     if (!OptionalProvider)
       continue;
     const auto &H = *OptionalProvider;
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.h b/clang-tools-extra/clangd/index/SymbolCollector.h
index 1076502..20116fc 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.h
+++ b/clang-tools-extra/clangd/index/SymbolCollector.h
@@ -15,18 +15,25 @@
 #include "index/Relation.h"
 #include "index/Symbol.h"
 #include "index/SymbolID.h"
+#include "index/SymbolLocation.h"
 #include "index/SymbolOrigin.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
+#include "clang/Basic/LLVM.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Sema/CodeCompleteConsumer.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include <functional>
 #include <memory>
 #include <optional>
+#include <string>
+#include <utility>
 
 namespace clang {
 namespace clangd {
@@ -177,7 +184,7 @@ private:
 
   // Providers for Symbol.IncludeHeaders.
   // The final spelling is calculated in finish().
-  llvm::DenseMap<SymbolID, std::optional<include_cleaner::Header>>
+  llvm::DenseMap<SymbolID, llvm::SmallVector<include_cleaner::Header>>
       SymbolProviders;
   // Files which contain ObjC symbols.
   // This is finalized and used in finish().
diff --git a/clang-tools-extra/clangd/unittests/IndexActionTests.cpp b/clang-tools-extra/clangd/unittests/IndexActionTests.cpp
index fa3d9c3..2a9b8c9 100644
--- a/clang-tools-extra/clangd/unittests/IndexActionTests.cpp
+++ b/clang-tools-extra/clangd/unittests/IndexActionTests.cpp
@@ -341,6 +341,36 @@ TEST_F(IndexActionTest, SymbolFromCC) {
                   hasName("foo"),
                   includeHeader(URI::create(testPath("main.h")).toString()))));
 }
+
+TEST_F(IndexActionTest, IncludeHeaderForwardDecls) {
+  std::string MainFilePath = testPath("main.cpp");
+  addFile(MainFilePath, R"cpp(
+#include "fwd.h"
+#include "full.h"
+ )cpp");
+  addFile(testPath("fwd.h"), R"cpp(
+#ifndef _FWD_H_
+#define _FWD_H_
+struct Foo;
+#endif
+ )cpp");
+  addFile(testPath("full.h"), R"cpp(
+#ifndef _FULL_H_
+#define _FULL_H_
+struct Foo {};
+
+// This decl is important, as otherwise we detect control macro for the file,
+// before handling definition of Foo.
+void other();
+#endif
+ )cpp");
+  IndexFileIn IndexFile = runIndexingAction(MainFilePath);
+  EXPECT_THAT(*IndexFile.Symbols,
+              testing::Contains(AllOf(
+                  hasName("Foo"),
+                  includeHeader(URI::create(testPath("full.h")).toString()))))
+      << *IndexFile.Symbols->begin();
+}
 } // namespace
 } // namespace clangd
 } // namespace clang