From b8a073cb9f78fd16c4a561adbef00a4098c1c6aa Mon Sep 17 00:00:00 2001 From: Rintaro Ishizaki Date: Fri, 16 Jun 2023 14:40:12 -0700 Subject: [PATCH] [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. --- include/swift/AST/PluginLoader.h | 21 +++-- lib/AST/PluginLoader.cpp | 128 +++++++++++++++++++++---------- lib/Sema/TypeCheckMacros.cpp | 18 ++--- 3 files changed, 106 insertions(+), 61 deletions(-) diff --git a/include/swift/AST/PluginLoader.h b/include/swift/AST/PluginLoader.h index 5ef31029af38f..229161515f867 100644 --- a/include/swift/AST/PluginLoader.h +++ b/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 ExecutablePluginPaths; + /// Map a module name to an plugin entry that provides the module. + llvm::Optional> PluginMap; - void createModuleToExecutablePluginMap(); + /// Get or lazily create and populate 'PluginMap'. + llvm::DenseMap &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 - 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 diff --git a/lib/AST/PluginLoader.cpp b/lib/AST/PluginLoader.cpp index 49368dd9a0fa6..2f89336577107 100644 --- a/lib/AST/PluginLoader.cpp +++ b/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()) { - // 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 -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 & +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 '. case PluginSearchOption::Kind::LoadPluginLibrary: { auto &val = entry.get(); - 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 #, ...'. case PluginSearchOption::Kind::LoadPluginExecutable: { auto &val = entry.get(); - 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 '. case PluginSearchOption::Kind::PluginPath: { auto &val = entry.get(); - 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 #'. case PluginSearchOption::Kind::ExternalPluginPath: { auto &val = entry.get(); - 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) { diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp index 7621a64563e49..36a13fbd2fd87 100644 --- a/lib/Sema/TypeCheckMacros.cpp +++ b/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; } }