[Serialization] Improve module loading performance

When looking for a Swift module on disk, we were scanning all module search paths if they contain the module we are searching for. In a setup where each module is contained in its own framework search path, this scaled quadratically with the number of modules being imported. E.g. a setup with 100 modules being imported form 100 module search paths could cause on the order of 10,000 checks of `FileSystem::exists`. While these checks are fairly fast (~10µs), they add up to ~100ms.

To improve this, perform a first scan of all module search paths and list the files they contain. From this, create a lookup map that maps filenames to the search paths they can be found in. E.g. for
```
searchPath1/
  Module1.framework

searchPath2/
  Module1.framework
  Module2.swiftmodule
```
we create the following lookup table
```
Module1.framework -> [searchPath1, searchPath2]
Module2.swiftmodule -> [searchPath2]
```

Author: ahoppen

include/swift/AST/SearchPathOptions.h

@@ -16,27 +16,155 @@
 #include "swift/Basic/ArrayRefView.h"
 #include "swift/Basic/PathRemapper.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/VirtualFileSystem.h"
 
 #include <string>
 #include <vector>
 
 namespace swift {
 
-/// Options for controlling search path behavior.
-class SearchPathOptions {
+/// The kind of a module search path. The order of this enum is important
+/// because import search paths should be considered before framework search
+/// paths etc.
+enum class ModuleSearchPathKind {
+  Import,
+  Framework,
+  DarwinImplictFramework,
+  RuntimeLibrary,
+};
+
+/// A single module search path that can come from different sources, e.g.
+/// framework search paths, import search path etc.
+struct ModuleSearchPath {
+  /// The actual path of the module search path. References a search path string
+  /// stored inside \c SearchPathOptions, which must outlive this reference.
+  StringRef Path;
+
+  /// The kind of the search path.
+  ModuleSearchPathKind Kind;
+
+  bool IsSystem;
+
+  /// An index that describes the order this search path should be considered
+  /// in within its \c ModuleSearchPathKind. This allows us to reconstruct the
+  /// user-defined search path order when merging search paths containing
+  /// different file names in \c searchPathsContainingFile.
+  unsigned Index;
+
+  bool operator<(const ModuleSearchPath &Other) const {
+    if (this->Kind == Other.Kind) {
+      return this->Index < Other.Index;
+    } else {
+      return this->Kind < Other.Kind;
+    }
+  }
+};
+
+class SearchPathOptions;
+
+/// Maintains a mapping of filenames to search paths that contain a file with
+/// this name (non-recursively). E.g. if we have a directory structure as
+/// follows.
+///
+/// \code
+/// searchPath1/
+///   Module1.framework
+///
+/// searchPath2/
+///   Module1.framework
+///   Module2.swiftmodule
+/// \endcode
+///
+/// We have the following lookup table
+///
+/// \code
+/// Module1.framework -> [searchPath1, searchPath2]
+/// Module2.swiftmodule -> [searchPath2]
+/// \endcode
+///
+/// When searching for a module this allows an efficient search of only those
+/// search paths that are relevant. In a naive implementation, we would need
+/// to scan all search paths for every module we import.
+class ModuleSearchPathLookup {
+  /// Parameters for which the \c LookupTable has been built. If one if these
+  /// changes, the lookup table needs to be rebuilt. It is not expected that any
+  /// of these change frequently.
+  struct {
+    llvm::vfs::FileSystem *FileSystem;
+    bool IsOSDarwin;
+    bool IsPopulated;
+    const SearchPathOptions *Opts;
+  } State;
+
+  llvm::StringMap<SmallVector<ModuleSearchPath, 4>> LookupTable;
+
+  /// Scan the directory at \p SearchPath for files and add those files to the
+  /// lookup table. \p Kind specifies the search path kind and \p Index the
+  /// index of \p SearchPath within that search path kind. Search paths with
+  /// lower indicies are considered first.
+  /// The \p SearchPath is stored by as a \c StringRef, so the string backing it
+  /// must be alive as long as this lookup table is alive and not cleared.
+  void addFilesInPathToLookupTable(llvm::vfs::FileSystem *FS,
+                                   StringRef SearchPath,
+                                   ModuleSearchPathKind Kind, bool IsSystem,
+                                   unsigned Index);
+
+  /// Discard the current lookup table and rebuild a new one.
+  void rebuildLookupTable(const SearchPathOptions *Opts,
+                          llvm::vfs::FileSystem *FS, bool IsOsDarwin);
+
+  /// Discard the current lookup table.
+  void clearLookupTable() {
+    LookupTable.clear();
+    State.IsPopulated = false;
+    State.FileSystem = nullptr;
+    State.IsOSDarwin = false;
+    State.Opts = nullptr;
+  }
+
 public:
-  /// Path to the SDK which is being built against.
-  std::string SDKPath;
+  /// Called by \p SearchPathOptions when search paths indexed by this \c
+  /// SearchPathLookup have changed in an unknown way. Causes the lookup table
+  /// to be rebuilt at the next request.
+  void searchPathsDidChange() { clearLookupTable(); }
 
-  /// Path(s) which should be searched for modules.
-  ///
-  /// Do not add values to this directly. Instead, use
-  /// \c ASTContext::addSearchPath.
-  std::vector<std::string> ImportSearchPaths;
+  /// Called by \p SearchPathOptions when an import or framework search path has
+  /// been added.
+  /// \p Index is the index of the search path within its kind and is used to
+  /// make sure this search path is considered last (within its kind).
+  void searchPathAdded(llvm::vfs::FileSystem *FS, StringRef SearchPath,
+                       ModuleSearchPathKind Kind, bool IsSystem,
+                       unsigned Index) {
+    if (!State.IsPopulated) {
+      // If the lookup table hasn't been built yet, we will scan the search path
+      // once the lookup table is requested. Nothing to do yet.
+      return;
+    }
+    if (State.FileSystem != FS) {
+      // We would be using a different file system to augment the lookup table
+      // than we initially used to build it. Discard everything to be safe.
+      clearLookupTable();
+      return;
+    }
+    addFilesInPathToLookupTable(FS, SearchPath, Kind, IsSystem, Index);
+  }
 
-  /// Path(s) to virtual filesystem overlay YAML files.
-  std::vector<std::string> VFSOverlayFiles;
+  /// Returns all search paths that non-recursively contain a file whose name
+  /// is in \p Filenames.
+  SmallVector<const ModuleSearchPath *, 4>
+  searchPathsContainingFile(const SearchPathOptions *Opts,
+                            llvm::ArrayRef<std::string> Filenames,
+                            llvm::vfs::FileSystem *FS, bool IsOSDarwin);
+};
+
+/// Options for controlling search path behavior.
+class SearchPathOptions {
+  /// To call \c addImportSearchPath and \c addFrameworkSearchPath from
+  /// \c ASTContext::addSearchPath.
+  friend class ASTContext;
 
+public:
   struct FrameworkSearchPath {
     std::string Path;
     bool IsSystem = false;
@@ -52,12 +180,116 @@ class SearchPathOptions {
       return !(LHS == RHS);
     }
   };
+
+private:
+  ModuleSearchPathLookup Lookup;
+
+  /// Path to the SDK which is being built against.
+  ///
+  /// Must me modified through setter to keep \c SearchPathLookup in sync.
+  std::string SDKPath;
+
+  /// Path(s) which should be searched for modules.
+  ///
+  /// Must me modified through setter to keep \c SearchPathLookup in sync.
+  std::vector<std::string> ImportSearchPaths;
+
   /// Path(s) which should be searched for frameworks.
   ///
-  /// Do not add values to this directly. Instead, use
-  /// \c ASTContext::addSearchPath.
+  /// Must me modified through setter to keep \c SearchPathLookup in sync.
   std::vector<FrameworkSearchPath> FrameworkSearchPaths;
 
+  /// Paths to search for stdlib modules. One of these will be
+  /// compiler-relative.
+  ///
+  /// Must me modified through setter to keep \c SearchPathLookup in sync.
+  std::vector<std::string> RuntimeLibraryImportPaths;
+
+  /// When on Darwin the framework paths that are implicitly imported.
+  /// $SDKROOT/System/Library/Frameworks/ and $SDKROOT/Library/Frameworks/.
+  ///
+  /// On non-Darwin platforms these are populated, but ignored.
+  ///
+  /// Computed when the SDK path is set and cached so we can reference the
+  /// Darwin implicit framework search paths as \c StringRef from
+  /// \c ModuleSearchPath.
+  std::vector<std::string> DarwinImplicitFrameworkSearchPaths;
+
+  /// Add a single import search path. Must only be called from
+  /// \c ASTContext::addSearchPath.
+  void addImportSearchPath(StringRef Path, llvm::vfs::FileSystem *FS) {
+    ImportSearchPaths.push_back(Path.str());
+    Lookup.searchPathAdded(FS, ImportSearchPaths.back(),
+                           ModuleSearchPathKind::Import, /*isSystem=*/false,
+                           ImportSearchPaths.size() - 1);
+  }
+
+  /// Add a single framework search path. Must only be called from
+  /// \c ASTContext::addSearchPath.
+  void addFrameworkSearchPath(FrameworkSearchPath NewPath,
+                              llvm::vfs::FileSystem *FS) {
+    FrameworkSearchPaths.push_back(NewPath);
+    Lookup.searchPathAdded(FS, FrameworkSearchPaths.back().Path,
+                           ModuleSearchPathKind::Framework, NewPath.IsSystem,
+                           FrameworkSearchPaths.size() - 1);
+  }
+
+public:
+  StringRef getSDKPath() const { return SDKPath; }
+
+  void setSDKPath(std::string NewSDKPath) {
+    SDKPath = NewSDKPath;
+
+    // Compute Darwin implicit framework search paths.
+    SmallString<128> systemFrameworksScratch(NewSDKPath);
+    llvm::sys::path::append(systemFrameworksScratch, "System", "Library",
+                            "Frameworks");
+    SmallString<128> frameworksScratch(NewSDKPath);
+    llvm::sys::path::append(frameworksScratch, "Library", "Frameworks");
+    DarwinImplicitFrameworkSearchPaths = {systemFrameworksScratch.str().str(),
+                                          frameworksScratch.str().str()};
+
+    Lookup.searchPathsDidChange();
+  }
+
+  ArrayRef<std::string> getImportSearchPaths() const {
+    return ImportSearchPaths;
+  }
+
+  void setImportSearchPaths(std::vector<std::string> NewImportSearchPaths) {
+    ImportSearchPaths = NewImportSearchPaths;
+    Lookup.searchPathsDidChange();
+  }
+
+  ArrayRef<FrameworkSearchPath> getFrameworkSearchPaths() const {
+    return FrameworkSearchPaths;
+  }
+
+  void setFrameworkSearchPaths(
+      std::vector<FrameworkSearchPath> NewFrameworkSearchPaths) {
+    FrameworkSearchPaths = NewFrameworkSearchPaths;
+    Lookup.searchPathsDidChange();
+  }
+
+  /// The extra implicit framework search paths on Apple platforms:
+  /// $SDKROOT/System/Library/Frameworks/ and $SDKROOT/Library/Frameworks/.
+  ArrayRef<std::string> getDarwinImplicitFrameworkSearchPaths() const {
+    return DarwinImplicitFrameworkSearchPaths;
+  }
+
+  ArrayRef<std::string> getRuntimeLibraryImportPaths() const {
+    return RuntimeLibraryImportPaths;
+  }
+
+  void setRuntimeLibraryImportPaths(
+      std::vector<std::string> NewRuntimeLibraryImportPaths) {
+    RuntimeLibraryImportPaths = NewRuntimeLibraryImportPaths;
+    Lookup.searchPathsDidChange();
+  }
+
+  /// Path(s) to virtual filesystem overlay YAML files.
+  std::vector<std::string> VFSOverlayFiles;
+
   /// Path(s) which should be searched for libraries.
   ///
   /// This is used in immediate modes. It is safe to add paths to this directly.
@@ -70,9 +302,6 @@ class SearchPathOptions {
   /// preference.
   std::vector<std::string> RuntimeLibraryPaths;
 
-  /// Paths to search for stdlib modules. One of these will be compiler-relative.
-  std::vector<std::string> RuntimeLibraryImportPaths;
-
   /// Don't look in for compiler-provided modules.
   bool SkipRuntimeLibraryImportPaths = false;
 
@@ -107,6 +336,14 @@ class SearchPathOptions {
   /// original form.
   PathObfuscator DeserializedPathRecoverer;
 
+  /// Return all module search paths that (non-recursively) contain a file whose
+  /// name is in \p Filenames.
+  SmallVector<const ModuleSearchPath *, 4>
+  moduleSearchPathsContainingFile(llvm::ArrayRef<std::string> Filenames,
+                                  llvm::vfs::FileSystem *FS, bool IsOSDarwin) {
+    return Lookup.searchPathsContainingFile(this, Filenames, FS, IsOSDarwin);
+  }
+
 private:
   static StringRef
   pathStringFromFrameworkSearchPath(const FrameworkSearchPath &next) {
@@ -141,7 +378,6 @@ class SearchPathOptions {
                         DisableModulesValidateSystemDependencies);
   }
 };
-
 }
 
 #endif

include/swift/Basic/PathRemapper.h

@@ -24,6 +24,7 @@
 #ifndef SWIFT_BASIC_PATHREMAPPER_H
 #define SWIFT_BASIC_PATHREMAPPER_H
 
+#include "swift/Basic/LLVM.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Twine.h"
 

include/swift/Frontend/Frontend.h

@@ -170,20 +170,20 @@ class CompilerInvocation {
   }
 
   void setImportSearchPaths(const std::vector<std::string> &Paths) {
-    SearchPathOpts.ImportSearchPaths = Paths;
+    SearchPathOpts.setImportSearchPaths(Paths);
   }
 
   ArrayRef<std::string> getImportSearchPaths() const {
-    return SearchPathOpts.ImportSearchPaths;
+    return SearchPathOpts.getImportSearchPaths();
   }
 
   void setFrameworkSearchPaths(
              const std::vector<SearchPathOptions::FrameworkSearchPath> &Paths) {
-    SearchPathOpts.FrameworkSearchPaths = Paths;
+    SearchPathOpts.setFrameworkSearchPaths(Paths);
   }
 
   ArrayRef<SearchPathOptions::FrameworkSearchPath> getFrameworkSearchPaths() const {
-    return SearchPathOpts.FrameworkSearchPaths;
+    return SearchPathOpts.getFrameworkSearchPaths();
   }
 
   void setExtraClangArgs(const std::vector<std::string> &Args) {
@@ -229,9 +229,7 @@ class CompilerInvocation {
 
   void setSDKPath(const std::string &Path);
 
-  StringRef getSDKPath() const {
-    return SearchPathOpts.SDKPath;
-  }
+  StringRef getSDKPath() const { return SearchPathOpts.getSDKPath(); }
 
   LangOptions &getLangOptions() {
     return LangOpts;

include/swift/Serialization/ModuleDependencyScanner.h

@@ -110,14 +110,14 @@ namespace swift {
         }
       }
 
-      std::error_code findModuleFilesInDirectory(
-          ImportPath::Element ModuleID,
-          const SerializedModuleBaseName &BaseName,
-          SmallVectorImpl<char> *ModuleInterfacePath,
-          std::unique_ptr<llvm::MemoryBuffer> *ModuleBuffer,
-          std::unique_ptr<llvm::MemoryBuffer> *ModuleDocBuffer,
-          std::unique_ptr<llvm::MemoryBuffer> *ModuleSourceInfoBuffer,
-          bool skipBuildingInterface, bool IsFramework) override;
+      virtual bool
+      findModule(ImportPath::Element moduleID,
+                 SmallVectorImpl<char> *moduleInterfacePath,
+                 std::unique_ptr<llvm::MemoryBuffer> *moduleBuffer,
+                 std::unique_ptr<llvm::MemoryBuffer> *moduleDocBuffer,
+                 std::unique_ptr<llvm::MemoryBuffer> *moduleSourceInfoBuffer,
+                 bool skipBuildingInterface, bool &isFramework,
+                 bool &isSystemModule) override;
 
       static bool classof(const ModuleDependencyScanner *MDS) {
         return MDS->getKind() == MDS_placeholder;