[clang][Lexer] Speedup HeaderSearch when there are many HeaderMaps

[troyj]

Summary

HeaderSearch already uses a caching system to avoid duplicate searches, but the initial cold miss can take a long time if a build system has supplied thousands of HeaderMaps. For this case, the SearchDirs vector begins with those HeaderMaps, so a cache miss requires testing if the sought filename is present in each of those maps. Instead, we can consolidate the keys of those HeaderMaps into one StringMap and then each cache miss can skip directly to the correct HeaderMap or continue its search beyond the initial sequence of HeaderMaps. In testing on TUs with ~15000 SearchDirs where the initial 99% are HeaderMaps, time spent in Clang was reduced by 15%. This patch is expected to be neutral for SearchDir vectors that do not begin with HeaderMaps.

[bruno]

Thanks for working on this Troy, nice speed up.

A bit more context for reviewers: we're probably unique on how we use hmaps, we have tons of them in a single compiler invocation. The other users of hmaps I've seen don't use more than a handful. This means that general users shouldn't expect much disruption here.

clang/include/clang/Lex/HeaderMap.h

52

Please remove the curly braces for this if and similar ones in the rest of the patch (if present).

clang/include/clang/Lex/HeaderSearch.h

255

Did you end up getting data on the memory footprint of this map? It seems to me that it should be fine if header maps aren't much populated.

clang/lib/Lex/HeaderSearch.cpp

391

Remove this empty line.

392

What happens if different header maps contain the same header name?

1034

Looks like we could have only one call to CacheLookup.reset instead?

[troyj]

clang/include/clang/Lex/HeaderSearch.h

255

Right, with ~15000 header maps, this map has ~92000 entries, so each map contributes about 6 unique keys. The patch does not increase the peak memory usage of Clang.

clang/lib/Lex/HeaderSearch.cpp

392

The first one sticks because try_emplace doesn't insert if the key already exists. I'll add a comment.

1034

The reset call has a side effect of setting CacheLookup.MappedName to nullptr, so we can't call it unconditionally because the Hit case should not call it. The unpatched code combined the !SkipCache check and the Hit check, so it still called reset when SkipCache == true. I preserved that behavior because I think it's necessary. SkipCache seems to imply "don't consult the cache for this call" instead of "don't cache anything" so it still caches the result for a future LookupFile call that passes SkipCache == false.

[troyj]

Addressed Bruno's comments.

[bruno]

Thanks for the detailed explanations, LGTM!

clang/include/clang/Lex/HeaderMap.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_LEX_HEADERMAP_H
 #define LLVM_CLANG_LEX_HEADERMAP_H
 
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/LLVM.h"
+#include "clang/Lex/HeaderMapTypes.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <memory>
 
 namespace clang {
 
 struct HMapBucket;
 struct HMapHeader;
 
 /// Implementation for \a HeaderMap that doesn't depend on \a FileManager.
 class HeaderMapImpl {
   std::unique_ptr<const llvm::MemoryBuffer> FileBuffer;
   bool NeedsBSwap;
   mutable llvm::StringMap<StringRef> ReverseMap;
 
 public:
   HeaderMapImpl(std::unique_ptr<const llvm::MemoryBuffer> File, bool NeedsBSwap)
       : FileBuffer(std::move(File)), NeedsBSwap(NeedsBSwap) {}
 
   // Check for a valid header and extract the byte swap.
   static bool checkHeader(const llvm::MemoryBuffer &File, bool &NeedsByteSwap);
 
+  // Make a call for every Key in the map.
+  template <typename Func> void forEachKey(Func Callback) const {
+    const HMapHeader &Hdr = getHeader();
+    unsigned NumBuckets = getEndianAdjustedWord(Hdr.NumBuckets);
+
+    for (unsigned Bucket = 0; Bucket < NumBuckets; ++Bucket) {
+      HMapBucket B = getBucket(Bucket);
+      if (B.Key != HMAP_EmptyBucketKey) {
+        Optional<StringRef> Key = getString(B.Key);
+        if (Key)
+          Callback(Key.value());
+      }
+    }
+  }
+
   /// If the specified relative filename is located in this HeaderMap return
   /// the filename it is mapped to, otherwise return an empty StringRef.
   StringRef lookupFilename(StringRef Filename,
                            SmallVectorImpl<char> &DestPath) const;
 
   /// Return the filename of the headermap.
   StringRef getFileName() const;
 

 
 public:
   /// This attempts to load the specified file as a header map.  If it doesn't
   /// look like a HeaderMap, it gives up and returns null.
   static std::unique_ptr<HeaderMap> Create(const FileEntry *FE,
                                            FileManager &FM);
 
   using HeaderMapImpl::dump;
+  using HeaderMapImpl::forEachKey;
   using HeaderMapImpl::getFileName;
   using HeaderMapImpl::lookupFilename;
   using HeaderMapImpl::reverseLookupFilename;
 };
 
 } // end namespace clang.
 
 #endif

clang/include/clang/Lex/HeaderSearch.h

class HeaderSearch {
   std::vector<DirectoryLookup> SearchDirs;
   /// Whether the DirectoryLookup at the corresponding index in SearchDirs has
   /// been successfully used to lookup a file.
   std::vector<bool> SearchDirsUsage;
   unsigned AngledDirIdx = 0;
   unsigned SystemDirIdx = 0;
   bool NoCurDirSearch = false;
 
+  /// Maps HeaderMap keys to SearchDir indices. When HeaderMaps are used
+  /// heavily, SearchDirs can start with thousands of HeaderMaps, so this Index
+  /// lets us avoid scanning them all to find a match.
+  llvm::StringMap<unsigned, llvm::BumpPtrAllocator> SearchDirHeaderMapIndex;
+
+  /// The index of the first SearchDir that isn't a header map.
+  unsigned FirstNonHeaderMapSearchDirIdx = 0;
+
   /// \#include prefixes for which the 'system header' property is
   /// overridden.
   ///
   /// For a \#include "x" or \#include \<x> directive, the last string in this
   /// list which is a prefix of 'x' determines whether the file is treated as
   /// a system header.
   std::vector<std::pair<std::string, bool>> SystemHeaderPrefixes;
 
class HeaderSearch {
   /// Entity used to resolve the identifier IDs of controlling
   /// macros into IdentifierInfo pointers, and keep the identifire up to date,
   /// as needed.
   ExternalPreprocessorSource *ExternalLookup = nullptr;
 
   /// Entity used to look up stored header file information.
   ExternalHeaderFileInfoSource *ExternalSource = nullptr;
 
+  /// Scan all of the header maps at the beginning of SearchDirs and
+  /// map their keys to the SearchDir index of their header map.
+  void indexInitialHeaderMaps();
+
 public:
   HeaderSearch(std::shared_ptr<HeaderSearchOptions> HSOpts,
                SourceManager &SourceMgr, DiagnosticsEngine &Diags,
                const LangOptions &LangOpts, const TargetInfo *Target);
   HeaderSearch(const HeaderSearch &) = delete;
   HeaderSearch &operator=(const HeaderSearch &) = delete;
 
   /// Retrieve the header-search options with which this header search
public:
 
   SearchDirIterator search_dir_begin() { return {*this, 0}; }
   SearchDirIterator search_dir_end() { return {*this, SearchDirs.size()}; }
   SearchDirRange search_dir_range() {
     return {search_dir_begin(), search_dir_end()};
   }
 
   ConstSearchDirIterator search_dir_begin() const { return quoted_dir_begin(); }
+  ConstSearchDirIterator search_dir_nth(size_t n) const {
+    assert(n < SearchDirs.size());
+    return {*this, n};
+  }
   ConstSearchDirIterator search_dir_end() const { return system_dir_end(); }
   ConstSearchDirRange search_dir_range() const {
     return {search_dir_begin(), search_dir_end()};
   }
 
   unsigned search_dir_size() const { return SearchDirs.size(); }
 
   ConstSearchDirIterator quoted_dir_begin() const { return {*this, 0}; }

clang/lib/Lex/HeaderSearch.cpp

assert(angledDirIdx <= systemDirIdx && systemDirIdx <= dirs.size() &&
          "Directory indices are unordered");
   SearchDirs = std::move(dirs);
   SearchDirsUsage.assign(SearchDirs.size(), false);
   AngledDirIdx = angledDirIdx;
   SystemDirIdx = systemDirIdx;
   NoCurDirSearch = noCurDirSearch;
   SearchDirToHSEntry = std::move(searchDirToHSEntry);
   //LookupFileCache.clear();
+  indexInitialHeaderMaps();
 }
 
 void HeaderSearch::AddSearchPath(const DirectoryLookup &dir, bool isAngled) {
   unsigned idx = isAngled ? SystemDirIdx : AngledDirIdx;
   SearchDirs.insert(SearchDirs.begin() + idx, dir);
   SearchDirsUsage.insert(SearchDirsUsage.begin() + idx, false);
   if (!isAngled)
     AngledDirIdx++;
for (DirectoryLookup Dir : search_dir_range()) {
     Module = ModMap.findModule(ModuleName);
     if (Module)
       break;
   }
 
   return Module;
 }
 
+void HeaderSearch::indexInitialHeaderMaps() {
+  llvm::StringMap<unsigned, llvm::BumpPtrAllocator> Index(SearchDirs.size());
+
+  // Iterate over all filename keys and associate them with the index i.
+  unsigned i = 0;
+  for (; i != SearchDirs.size(); ++i) {
+    auto &Dir = SearchDirs[i];
+
+    // We're concerned with only the initial contiguous run of header
+    // maps within SearchDirs, which can be 99% of SearchDirs when
+    // SearchDirs.size() is ~10000.
+    if (!Dir.isHeaderMap())
+      break;
+
+    // Give earlier keys precedence over identical later keys.
+    auto Callback = [&](StringRef Filename) { Index.try_emplace(Filename, i); };
+    Dir.getHeaderMap()->forEachKey(Callback);
+  }
+
+  SearchDirHeaderMapIndex = std::move(Index);
+  FirstNonHeaderMapSearchDirIdx = i;
+}
+
 //===----------------------------------------------------------------------===//
 // File lookup within a DirectoryLookup scope
 //===----------------------------------------------------------------------===//
 
 /// getName - Return the directory or filename corresponding to this lookup
 /// object.
 StringRef DirectoryLookup::getName() const {
   // FIXME: Use the name from \c DirectoryEntryRef.
Optional<FileEntryRef> HeaderSearch::LookupFile(
   // Cache all of the lookups performed by this method.  Many headers are
   // multiply included, and the "pragma once" optimization prevents them from
   // being relex/pp'd, but they would still have to search through a
   // (potentially huge) series of SearchDirs to find it.
   LookupFileCacheInfo &CacheLookup = LookupFileCache[Filename];
 
   ConstSearchDirIterator NextIt = std::next(It);
 
-  // If the entry has been previously looked up, the first value will be
-  // non-zero.  If the value is equal to i (the start point of our search), then
-  // this is a matching hit.
-  if (!SkipCache && CacheLookup.StartIt == NextIt) {
-    // Skip querying potentially lots of directories for this lookup.
+  if (!SkipCache) {
+    if (CacheLookup.StartIt == NextIt) {
+      // HIT: Skip querying potentially lots of directories for this lookup.
       if (CacheLookup.HitIt)
         It = CacheLookup.HitIt;
       if (CacheLookup.MappedName) {
         Filename = CacheLookup.MappedName;
         if (IsMapped)
           *IsMapped = true;
       }
     } else {
-    // Otherwise, this is the first query, or the previous query didn't match
-    // our search start.  We will fill in our found location below, so prime the
-    // start point value.
+      // MISS: This is the first query, or the previous query didn't match
+      // our search start.  We will fill in our found location below, so prime
+      // the start point value.
       CacheLookup.reset(/*NewStartIt=*/NextIt);
+
+      if (It == search_dir_begin() && FirstNonHeaderMapSearchDirIdx > 0) {
+        // Handle cold misses of user includes in the presence of many header
+        // maps.  We avoid searching perhaps thousands of header maps by
+        // jumping directly to the correct one or jumping beyond all of them.
+        auto Iter = SearchDirHeaderMapIndex.find(Filename);
+        if (Iter == SearchDirHeaderMapIndex.end())
+          // Not in index => Skip to first SearchDir after initial header maps
+          It = search_dir_nth(FirstNonHeaderMapSearchDirIdx);
+        else
+          // In index => Start with a specific header map
+          It = search_dir_nth(Iter->second);
       }
+    }
+  } else
+    CacheLookup.reset(/*NewStartIt=*/NextIt);
 
   SmallString<64> MappedName;
 
   // Check each directory in sequence to see if it contains this file.
   for (; It != search_dir_end(); ++It) {
     bool InUserSpecifiedSystemFramework = false;
     bool IsInHeaderMap = false;
     bool IsFrameworkFoundInDir = false;