[Macros] Create plugin search lookup table.

Iterating all options and potential file system access is not great for
every plugin lookup request. Instead, lazily create a single lookup table
keyed by module name.

Author: rintaro

include/swift/AST/PluginLoader.h

@@ -28,6 +28,13 @@ class ASTContext;
 ///  * Load plugins resolving VFS paths
 ///  * Track plugin dependencies
 class PluginLoader {
+public:
+  struct PluginEntry {
+    StringRef libraryPath;
+    StringRef executablePath;
+  };
+
+private:
   /// Plugin registry. Lazily populated by get/setRegistry().
   /// NOTE: Do not reference this directly. Use getRegistry().
   PluginRegistry *Registry = nullptr;
@@ -38,16 +45,15 @@ class PluginLoader {
   ASTContext &Ctx;
   DependencyTracker *DepTracker;
 
-  /// Map a module name to an executable plugin path that provides the module.
-  llvm::DenseMap<swift::Identifier, llvm::StringRef> ExecutablePluginPaths;
+  /// Map a module name to an plugin entry that provides the module.
+  llvm::Optional<llvm::DenseMap<swift::Identifier, PluginEntry>> PluginMap;
 
-  void createModuleToExecutablePluginMap();
+  /// Get or lazily create and populate 'PluginMap'.
+  llvm::DenseMap<swift::Identifier, PluginEntry> &getPluginMap();
 
 public:
   PluginLoader(ASTContext &Ctx, DependencyTracker *DepTracker)
-      : Ctx(Ctx), DepTracker(DepTracker) {
-    createModuleToExecutablePluginMap();
-  }
+      : Ctx(Ctx), DepTracker(DepTracker) {}
 
   void setRegistry(PluginRegistry *newValue);
   PluginRegistry *getRegistry();
@@ -63,8 +69,7 @@ class PluginLoader {
   ///    'loadExecutablePlugin()'.
   ///  * (libPath: some,  execPath: some) - load the executable path by
   ///    'loadExecutablePlugin()' and let the plugin load the libPath via IPC.
-  std::pair<std::string, std::string>
-  lookupPluginByModuleName(Identifier moduleName);
+  const PluginEntry &lookupPluginByModuleName(Identifier moduleName);
 
   /// Load the specified dylib plugin path resolving the path with the
   /// current VFS. If it fails to load the plugin, a diagnostic is emitted, and

lib/AST/PluginLoader.cpp

@@ -15,23 +15,11 @@
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/Basic/SourceManager.h"
+#include "swift/Parse/Lexer.h"
 #include "llvm/Config/config.h"
 
 using namespace swift;
 
-void PluginLoader::createModuleToExecutablePluginMap() {
-  for (auto &elem : Ctx.SearchPathOpts.PluginSearchOpts) {
-    if (auto *arg = elem.dyn_cast<PluginSearchOption::LoadPluginExecutable>()) {
-      // Create a moduleName -> pluginPath mapping.
-      assert(!arg->ExecutablePath.empty() && "empty plugin path");
-      StringRef pathStr = Ctx.AllocateCopy(arg->ExecutablePath);
-      for (auto moduleName : arg->ModuleNames) {
-        ExecutablePluginPaths[Ctx.getIdentifier(moduleName)] = pathStr;
-      }
-    }
-  }
-}
-
 void PluginLoader::setRegistry(PluginRegistry *newValue) {
   assert(Registry == nullptr && "Too late to set a new plugin registry");
   Registry = newValue;
@@ -48,65 +36,121 @@ PluginRegistry *PluginLoader::getRegistry() {
   return Registry;
 }
 
-std::pair<std::string, std::string>
-PluginLoader::lookupPluginByModuleName(Identifier moduleName) {
-  auto fs = Ctx.SourceMgr.getFileSystem();
-
-  // Look for 'lib${module name}(.dylib|.so)'.
+/// Get plugin module name from \p path if the path looks like a shared library
+/// path. Otherwise, returns an empty string.
+static StringRef pluginModuleNameStringFromPath(StringRef path) {
+  // Plugin library must be named 'lib${module name}(.dylib|.so|.dll)'.
   // FIXME: Shared library prefix might be different between platforms.
-  SmallString<64> pluginLibBasename;
-  pluginLibBasename.append("lib");
-  pluginLibBasename.append(moduleName.str());
-  pluginLibBasename.append(LTDL_SHLIB_EXT);
+  constexpr StringRef libPrefix = "lib";
+  constexpr StringRef libSuffix = LTDL_SHLIB_EXT;
+
+  StringRef filename = llvm::sys::path::filename(path);
+  if (filename.starts_with(libPrefix) && filename.ends_with(libSuffix)) {
+    // We don't check if the result it a valid identifier. Even if we put
+    // invalid name in the lookup table, clients wound not be able to lookup
+    // that name, thus harmless.
+    return filename.drop_front(libPrefix.size()).drop_back(libSuffix.size());
+  }
+  return "";
+}
+
+llvm::DenseMap<Identifier, PluginLoader::PluginEntry> &
+PluginLoader::getPluginMap() {
+  if (PluginMap.has_value()) {
+    return PluginMap.value();
+  }
+
+  // Create and populate the map.
+  PluginMap.emplace();
+  auto &map = PluginMap.value();
+
+  // Helper function to try inserting an entry if there's no existing entry
+  // associated with the module name.
+  auto try_emplace = [&](StringRef moduleName, StringRef libPath,
+                         StringRef execPath) {
+    auto moduleNameIdentifier = Ctx.getIdentifier(moduleName);
+    if (map.find(moduleNameIdentifier) != map.end()) {
+      // Specified module name is already in the map.
+      return;
+    }
+
+    libPath = libPath.empty() ? "" : Ctx.AllocateCopy(libPath);
+    execPath = execPath.empty() ? "" : Ctx.AllocateCopy(execPath);
+    auto result = map.insert({moduleNameIdentifier, {libPath, execPath}});
+    assert(result.second);
+    (void)result;
+  };
+
+  auto fs = Ctx.SourceMgr.getFileSystem();
+  std::error_code ec;
 
-  // FIXME: Should we create a lookup table keyed by module name?
   for (auto &entry : Ctx.SearchPathOpts.PluginSearchOpts) {
     switch (entry.getKind()) {
-    // Try '-load-plugin-library'.
+
+    // '-load-plugin-library <library path>'.
     case PluginSearchOption::Kind::LoadPluginLibrary: {
       auto &val = entry.get<PluginSearchOption::LoadPluginLibrary>();
-      if (llvm::sys::path::filename(val.LibraryPath) == pluginLibBasename) {
-        return {val.LibraryPath, ""};
+      auto moduleName = pluginModuleNameStringFromPath(val.LibraryPath);
+      if (!moduleName.empty()) {
+        try_emplace(moduleName, val.LibraryPath, /*executablePath=*/"");
       }
       continue;
     }
 
-    // Try '-load-plugin-executable'.
+    // '-load-plugin-executable <executable path>#<module name>, ...'.
     case PluginSearchOption::Kind::LoadPluginExecutable: {
       auto &val = entry.get<PluginSearchOption::LoadPluginExecutable>();
-      auto found = ExecutablePluginPaths.find(moduleName);
-      if (found != ExecutablePluginPaths.end() &&
-          found->second == val.ExecutablePath) {
-        return {"", val.ExecutablePath};
+      assert(!val.ExecutablePath.empty() && "empty plugin path");
+      for (auto &moduleName : val.ModuleNames) {
+        try_emplace(moduleName, /*libraryPath=*/"", val.ExecutablePath);
       }
       continue;
     }
 
-    // Try '-plugin-path'.
+    // '-plugin-path <library search path>'.
     case PluginSearchOption::Kind::PluginPath: {
       auto &val = entry.get<PluginSearchOption::PluginPath>();
-      SmallString<128> fullPath(val.SearchPath);
-      llvm::sys::path::append(fullPath, pluginLibBasename);
-      if (fs->exists(fullPath)) {
-        return {std::string(fullPath), ""};
+      for (auto i = fs->dir_begin(val.SearchPath, ec);
+           i != llvm::vfs::directory_iterator(); i = i.increment(ec)) {
+        auto libPath = i->path();
+        auto moduleName = pluginModuleNameStringFromPath(libPath);
+        if (!moduleName.empty()) {
+          try_emplace(moduleName, libPath, /*executablePath=*/"");
+        }
       }
       continue;
     }
 
-    // Try '-external-plugin-path'.
+    // '-external-plugin-path <library search path>#<server path>'.
     case PluginSearchOption::Kind::ExternalPluginPath: {
       auto &val = entry.get<PluginSearchOption::ExternalPluginPath>();
-      SmallString<128> fullPath(val.SearchPath);
-      llvm::sys::path::append(fullPath, pluginLibBasename);
-      if (fs->exists(fullPath)) {
-        return {std::string(fullPath), val.ServerPath};
+      for (auto i = fs->dir_begin(val.SearchPath, ec);
+           i != llvm::vfs::directory_iterator(); i = i.increment(ec)) {
+        auto libPath = i->path();
+        auto moduleName = pluginModuleNameStringFromPath(libPath);
+        if (!moduleName.empty()) {
+          try_emplace(moduleName, libPath, val.ServerPath);
+        }
       }
       continue;
     }
     }
+    llvm_unreachable("unhandled PluginSearchOption::Kind");
   }
 
-  return {};
+  return map;
+}
+
+const PluginLoader::PluginEntry &
+PluginLoader::lookupPluginByModuleName(Identifier moduleName) {
+  auto &map = getPluginMap();
+  auto found = map.find(moduleName);
+  if (found != map.end()) {
+    return found->second;
+  } else {
+    static PluginEntry notFound{"", ""};
+    return notFound;
+  }
 }
 
 LoadedLibraryPlugin *PluginLoader::loadLibraryPlugin(StringRef path) {

lib/Sema/TypeCheckMacros.cpp

@@ -342,21 +342,17 @@ LoadedCompilerPlugin
 CompilerPluginLoadRequest::evaluate(Evaluator &evaluator, ASTContext *ctx,
                                     Identifier moduleName) const {
   PluginLoader &loader = ctx->getPluginLoader();
+  const auto &entry = loader.lookupPluginByModuleName(moduleName);
 
-  std::string libraryPath;
-  std::string executablePath;
-  std::tie(libraryPath, executablePath) =
-      loader.lookupPluginByModuleName(moduleName);
-
-  if (!executablePath.empty()) {
+  if (!entry.executablePath.empty()) {
     if (LoadedExecutablePlugin *executablePlugin =
-            loader.loadExecutablePlugin(executablePath)) {
-      return initializeExecutablePlugin(*ctx, executablePlugin, libraryPath,
-                                        moduleName);
+            loader.loadExecutablePlugin(entry.executablePath)) {
+      return initializeExecutablePlugin(*ctx, executablePlugin,
+                                        entry.libraryPath, moduleName);
     }
-  } else if (!libraryPath.empty()) {
+  } else if (!entry.libraryPath.empty()) {
     if (LoadedLibraryPlugin *libraryPlugin =
-            loader.loadLibraryPlugin(libraryPath)) {
+            loader.loadLibraryPlugin(entry.libraryPath)) {
       return libraryPlugin;
     }
   }