[clangd] Support include canonicalization in symbol leve.

Summary:
Symbols with different canonical includes might be defined in the same header
(e.g. symbols defined in STL <iosfwd>). This patch adds support for mapping from
qualified symbol names to canonical headers and special mapping for symbols in <iosfwd>

Reviewers: sammccall, hokein

Reviewed By: sammccall

Subscribers: klimek, ilya-biryukov, jkorous-apple, cfe-commits

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

llvm-svn: 326456

clang-tools-extra/clangd/index/CanonicalIncludes.cpp

@@ -27,7 +27,19 @@ void CanonicalIncludes::addRegexMapping(llvm::StringRef RE,
   this->RegexHeaderMappingTable.emplace_back(llvm::Regex(RE), CanonicalPath);
 }
 
-llvm::StringRef CanonicalIncludes::mapHeader(llvm::StringRef Header) const {
+void CanonicalIncludes::addSymbolMapping(llvm::StringRef QualifiedName,
+                                         llvm::StringRef CanonicalPath) {
+  this->SymbolMapping[QualifiedName] = CanonicalPath;
+}
+
+llvm::StringRef
+CanonicalIncludes::mapHeader(llvm::StringRef Header,
+                             llvm::StringRef QualifiedName) const {
+  if (!QualifiedName.empty()) {
+    auto SE = SymbolMapping.find(QualifiedName);
+    if (SE != SymbolMapping.end())
+      return SE->second;
+  }
   std::lock_guard<std::mutex> Lock(RegexMutex);
   for (auto &Entry : RegexHeaderMappingTable) {
 #ifndef NDEBUG
@@ -67,6 +79,53 @@ collectIWYUHeaderMaps(CanonicalIncludes *Includes) {
 }
 
 void addSystemHeadersMapping(CanonicalIncludes *Includes) {
+  static const std::vector<std::pair<const char *, const char *>> SymbolMap = {
+      // Map symbols in <iosfwd> to their preferred includes.
+      {"std::basic_filebuf", "<fstream>"},
+      {"std::basic_fstream", "<fstream>"},
+      {"std::basic_ifstream", "<fstream>"},
+      {"std::basic_ofstream", "<fstream>"},
+      {"std::filebuf", "<fstream>"},
+      {"std::fstream", "<fstream>"},
+      {"std::ifstream", "<fstream>"},
+      {"std::ofstream", "<fstream>"},
+      {"std::wfilebuf", "<fstream>"},
+      {"std::wfstream", "<fstream>"},
+      {"std::wifstream", "<fstream>"},
+      {"std::wofstream", "<fstream>"},
+      {"std::basic_ios", "<ios>"},
+      {"std::ios", "<ios>"},
+      {"std::wios", "<ios>"},
+      {"std::basic_iostream", "<iostream>"},
+      {"std::iostream", "<iostream>"},
+      {"std::wiostream", "<iostream>"},
+      {"std::basic_istream", "<istream>"},
+      {"std::istream", "<istream>"},
+      {"std::wistream", "<istream>"},
+      {"std::istreambuf_iterator", "<iterator>"},
+      {"std::ostreambuf_iterator", "<iterator>"},
+      {"std::basic_ostream", "<ostream>"},
+      {"std::ostream", "<ostream>"},
+      {"std::wostream", "<ostream>"},
+      {"std::basic_istringstream", "<sstream>"},
+      {"std::basic_ostringstream", "<sstream>"},
+      {"std::basic_stringbuf", "<sstream>"},
+      {"std::basic_stringstream", "<sstream>"},
+      {"std::istringstream", "<sstream>"},
+      {"std::ostringstream", "<sstream>"},
+      {"std::stringbuf", "<sstream>"},
+      {"std::stringstream", "<sstream>"},
+      {"std::wistringstream", "<sstream>"},
+      {"std::wostringstream", "<sstream>"},
+      {"std::wstringbuf", "<sstream>"},
+      {"std::wstringstream", "<sstream>"},
+      {"std::basic_streambuf", "<streambuf>"},
+      {"std::streambuf", "<streambuf>"},
+      {"std::wstreambuf", "<streambuf>"},
+  };
+  for (const auto &Pair : SymbolMap)
+    Includes->addSymbolMapping(Pair.first, Pair.second);
+
   static const std::vector<std::pair<const char *, const char *>>
       SystemHeaderMap = {
           {"include/__stddef_max_align_t.h$", "<cstddef>"},

clang-tools-extra/clangd/index/CanonicalIncludes.h

@@ -43,9 +43,17 @@ public:
   /// Maps all files matching \p RE to \p CanonicalPath
   void addRegexMapping(llvm::StringRef RE, llvm::StringRef CanonicalPath);
 
+  /// Sets the canonical include for any symbol with \p QualifiedName.
+  /// Symbol mappings take precedence over header mappings.
+  void addSymbolMapping(llvm::StringRef QualifiedName,
+                        llvm::StringRef CanonicalPath);
+
   /// \return \p Header itself if there is no mapping for it; otherwise, return
   /// a canonical header name.
-  llvm::StringRef mapHeader(llvm::StringRef Header) const;
+  /// \p QualifiedName of a symbol declared in \p Header can be provided to
+  /// check against the symbol mapping.
+  llvm::StringRef mapHeader(llvm::StringRef Header,
+                            llvm::StringRef QualifiedName = "") const;
 
 private:
   // A map from header patterns to header names. This needs to be mutable so
@@ -55,6 +63,8 @@ private:
   // arbitrary regexes.
   mutable std::vector<std::pair<llvm::Regex, std::string>>
       RegexHeaderMappingTable;
+  // A map from fully qualified symbol names to header names.
+  llvm::StringMap<std::string> SymbolMapping;
   // Guards Regex matching as it's not thread-safe.
   mutable std::mutex RegexMutex;
 };
@@ -68,8 +78,9 @@ private:
 std::unique_ptr<CommentHandler>
 collectIWYUHeaderMaps(CanonicalIncludes *Includes);
 
-/// Adds mapping for system headers. Approximately, the following system headers
-/// are handled:
+/// Adds mapping for system headers and some special symbols (e.g. STL symbols
+/// in <iosfwd> need to be mapped individually). Approximately, the following
+/// system headers are handled:
 ///   - C++ standard library e.g. bits/basic_string.h$ -> <string>
 ///   - Posix library e.g. bits/pthreadtypes.h$ -> <pthread.h>
 ///   - Compiler extensions, e.g. include/avx512bwintrin.h$ -> <immintrin.h>

clang-tools-extra/clangd/index/SymbolCollector.cpp

@@ -154,13 +154,13 @@ bool shouldCollectIncludePath(index::SymbolKind Kind) {
 /// FIXME: we should handle .inc files whose symbols are expected be exported by
 /// their containing headers.
 llvm::Optional<std::string>
-getIncludeHeader(const SourceManager &SM, SourceLocation Loc,
-                 const SymbolCollector::Options &Opts) {
+getIncludeHeader(llvm::StringRef QName, const SourceManager &SM,
+                 SourceLocation Loc, const SymbolCollector::Options &Opts) {
   llvm::StringRef FilePath = SM.getFilename(Loc);
   if (FilePath.empty())
     return llvm::None;
   if (Opts.Includes) {
-    llvm::StringRef Mapped = Opts.Includes->mapHeader(FilePath);
+    llvm::StringRef Mapped = Opts.Includes->mapHeader(FilePath, QName);
     if (Mapped != FilePath)
       return (Mapped.startswith("<") || Mapped.startswith("\""))
                  ? Mapped.str()
@@ -316,8 +316,8 @@ const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND,
   if (Opts.CollectIncludePath && shouldCollectIncludePath(S.SymInfo.Kind)) {
     // Use the expansion location to get the #include header since this is
     // where the symbol is exposed.
-    if (auto Header =
-            getIncludeHeader(SM, SM.getExpansionLoc(ND.getLocation()), Opts))
+    if (auto Header = getIncludeHeader(
+            QName, SM, SM.getExpansionLoc(ND.getLocation()), Opts))
       Include = std::move(*Header);
   }
   S.CompletionFilterText = FilterText;

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

@@ -547,6 +547,32 @@ TEST_F(SymbolCollectorTest, CanonicalSTLHeader) {
 }
 #endif
 
+TEST_F(SymbolCollectorTest, STLiosfwd) {
+  CollectorOpts.CollectIncludePath = true;
+  CanonicalIncludes Includes;
+  addSystemHeadersMapping(&Includes);
+  CollectorOpts.Includes = &Includes;
+  // Symbols from <iosfwd> should be mapped individually.
+  TestHeaderName = testPath("iosfwd");
+  TestFileName = testPath("iosfwd.cpp");
+  std::string Header = R"(
+    namespace std {
+      class no_map {};
+      class ios {};
+      class ostream {};
+      class filebuf {};
+    } // namespace std
+  )";
+  runSymbolCollector(Header, /*Main=*/"");
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(
+                  QName("std"),
+                  AllOf(QName("std::no_map"), IncludeHeader("<iosfwd>")),
+                  AllOf(QName("std::ios"), IncludeHeader("<ios>")),
+                  AllOf(QName("std::ostream"), IncludeHeader("<ostream>")),
+                  AllOf(QName("std::filebuf"), IncludeHeader("<fstream>"))));
+}
+
 TEST_F(SymbolCollectorTest, IWYUPragma) {
   CollectorOpts.CollectIncludePath = true;
   CanonicalIncludes Includes;