[clangd] Support include canonicalization in symbol leve.

[ioeric]

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>

[hokein]

clangd/index/CanonicalIncludes.cpp

83

Looks like the list only contains stream-related symbols, there are some symbols in <iosfwd> not covered, is it intended?

Maybe we keep the order of this map the same as the one listed in http://en.cppreference.com/w/cpp/header/iosfwd? That would make it easier to see the difference?

unittests/clangd/SymbolCollectorTests.cpp

561

The STL implementation of ios is a typedef typedef basic_ios<char> ios; , I think we should make the test align with it?

[ioeric]

clangd/index/CanonicalIncludes.cpp

83

Although allocator could be exposed by <iosfwd>, it's not defined in <iosfwd>, so we don't need to handle it specially.

I am sorting the symbols by their canonical headers, which I think would make it easier to check the header mapping.

unittests/clangd/SymbolCollectorTests.cpp

561

The symbol mapping doesn't make a difference between symbol types, and symbol kind test is not in the scope of this patch, so I think this wouldn't be necessary here.

[sammccall]

clangd/index/CanonicalIncludes.h

47

Nit: "Symbol mappings take precedence over header mappings"?

51

Nit: rephrase as "returns the canonical include for \p Symbol, which is declared in \p Header"?

53–56

Why can't we always provide this?

unittests/clangd/SymbolCollectorTests.cpp

561

I'd be +1 on matching the real header a bit more so it's more obvious how this relates to the standard library. As long as it's not too hard!

[ioeric]

• Merge with origin/master
• Address review comments.

clangd/index/CanonicalIncludes.h

53–56

I think it would be better to keep supporting the header-only mapping use case.

unittests/clangd/SymbolCollectorTests.cpp

561

I'd prefer to keep the test case minimum so that it's clear that we are only matching on names.

clangd/index/CanonicalIncludes.h

public:
   CanonicalIncludes() = default;
 
   /// Adds a string-to-string mapping from \p Path to \p CanonicalPath.
   void addMapping(llvm::StringRef Path, llvm::StringRef CanonicalPath);
 
   /// 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
   // that we can match again a Regex in a const function member.
   // FIXME(ioeric): All the regexes we have so far are suffix matches. The
   // performance could be improved by allowing only suffix matches instead of
   // 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;
 };
 
 /// Returns a CommentHandler that parses pragma comment on include files to
 /// determine when we should include a different header from the header that
 /// directly defines a symbol. Mappinps are registered with \p Includes.
 ///
 /// Currently it only supports IWYU private pragma:
 /// https://github.com/include-what-you-use/include-what-you-use/blob/master/docs/IWYUPragmas.md#iwyu-pragma-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>
 /// The mapping is hardcoded and hand-maintained, so it might not cover all
 /// headers.
 void addSystemHeadersMapping(CanonicalIncludes *Includes);
 
 } // namespace clangd
 } // namespace clang
 
 #endif //LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_HEADERMAPCOLLECTOR_H

clangd/index/CanonicalIncludes.cpp

addRegexMapping((llvm::Twine("^") + llvm::Regex::escape(Path) + "$").str(),
                   CanonicalPath);
 }
 
 void CanonicalIncludes::addRegexMapping(llvm::StringRef RE,
                                         llvm::StringRef CanonicalPath) {
   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
     std::string Dummy;
     assert(Entry.first.isValid(Dummy) && "Regex should never be invalid!");
 #endif
     if (Entry.first.match(Header))
       return Entry.second;
collectIWYUHeaderMaps(CanonicalIncludes *Includes) {
 
   private:
     CanonicalIncludes *const Includes;
   };
   return llvm::make_unique<PragmaCommentHandler>(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>"},
           {"include/__wmmintrin_aes.h$", "<wmmintrin.h>"},
           {"include/__wmmintrin_pclmul.h$", "<wmmintrin.h>"},
           {"include/adxintrin.h$", "<immintrin.h>"},
           {"include/ammintrin.h$", "<ammintrin.h>"},
           {"include/avx2intrin.h$", "<immintrin.h>"},

clangd/index/SymbolCollector.cpp

 }
 
 /// Gets a canonical include (URI of the header or <header>  or "header") for
 /// header of \p Loc.
 /// Returns None if fails to get include header for \p Loc.
 /// 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()
                  : ("\"" + Mapped + "\"").str();
   }
 
   return toURI(SM, SM.getFilename(Loc), Opts);
 }
const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND,
   std::string FilterText = getFilterText(*CCS);
   std::string Documentation = getDocumentation(*CCS);
   std::string CompletionDetail = getDetail(*CCS);
 
   std::string Include;
   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;
   S.CompletionLabel = Label;
   S.CompletionPlainInsertText = PlainInsertText;
   S.CompletionSnippetInsertText = SnippetInsertText;
   Symbol::Details Detail;
   Detail.Documentation = Documentation;

unittests/clangd/SymbolCollectorTests.cpp

TEST_F(SymbolCollectorTest, CanonicalSTLHeader) {
   TestHeaderURI = URI::createFile(TestHeaderName).toString();
   runSymbolCollector("class string {};", /*Main=*/"");
   EXPECT_THAT(Symbols, UnorderedElementsAre(AllOf(QName("string"),
                                                   DeclURI(TestHeaderURI),
                                                   IncludeHeader("<string>"))));
 }
 #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;
   PragmaHandler = collectIWYUHeaderMaps(&Includes);
   CollectorOpts.Includes = &Includes;
   const std::string Header = R"(
     // IWYU pragma: private, include the/good/header.h
     class Foo {};