Merge pull request #64610 from rintaro/macros-lazy-plugin-library

[Macros] Load '-load-plugin-library' plugins lazily

Author: rintaro

include/swift/AST/ASTContext.h

@@ -1461,12 +1461,21 @@ class ASTContext final {
   /// The declared interface type of Builtin.TheTupleType.
   BuiltinTupleType *getBuiltinTupleType();
 
-  /// Finds the address of the given symbol. If `libraryHandleHint` is non-null,
-  /// search within the library.
-  void *getAddressOfSymbol(const char *name, void *libraryHandleHint = nullptr);
-  
   Type getNamedSwiftType(ModuleDecl *module, StringRef name);
 
+  /// Lookup a library plugin that can handle \p moduleName and return the path
+  /// to it.
+  /// The path is valid within the VFS, use `FS.getRealPath()` for the
+  /// underlying path.
+  Optional<std::string> lookupLibraryPluginByModuleName(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
+  /// returns a nullptr.
+  /// NOTE: This method is idempotent. If the plugin is already loaded, the same
+  /// instance is simply returned.
+  void *loadLibraryPlugin(StringRef path);
+
   /// Lookup an executable plugin that is declared to handle \p moduleName
   /// module by '-load-plugin-executable'.
   /// The path is valid within the VFS, use `FS.getRealPath()` for the
@@ -1506,7 +1515,7 @@ class ASTContext final {
   Optional<StringRef> getBriefComment(const Decl *D);
   void setBriefComment(const Decl *D, StringRef Comment);
 
-  void loadCompilerPlugins();
+  void createModuleToExecutablePluginMap();
 
   friend TypeBase;
   friend ArchetypeType;

include/swift/AST/DiagnosticsFrontend.def

@@ -128,6 +128,9 @@ ERROR(error_invalid_source_location_str,none,
 ERROR(error_no_source_location_scope_map,none,
       "-dump-scope-maps argument must be 'expanded' or a list of "
       "source locations", ())
+ERROR(error_load_plugin_executable,none,
+      "invalid value '%0' in '-load-plugin-executable'; "
+      "make sure to use format '<plugin path>#<module names>'", (StringRef))
 
 NOTE(note_valid_swift_versions, none,
      "valid arguments to '-swift-version' are %0", (StringRef))

include/swift/AST/SearchPathOptions.h

@@ -174,6 +174,12 @@ class ModuleSearchPathLookup {
                             llvm::vfs::FileSystem *FS, bool IsOSDarwin);
 };
 
+/// Pair of a plugin path and the module name that the plugin provides.
+struct PluginExecutablePathAndModuleNames {
+  std::string ExecutablePath;
+  std::vector<std::string> ModuleNames;
+};
+
 /// Pair of a plugin search path and the corresponding plugin server executable
 /// path.
 struct ExternalPluginSearchPathAndServerPath {
@@ -242,8 +248,7 @@ class SearchPathOptions {
   std::vector<std::string> CompilerPluginLibraryPaths;
 
   /// Compiler plugin executable paths and providing module names.
-  /// Format: '<path>#<module names>'
-  std::vector<std::string> CompilerPluginExecutablePaths;
+  std::vector<PluginExecutablePathAndModuleNames> CompilerPluginExecutablePaths;
 
   /// Add a single import search path. Must only be called from
   /// \c ASTContext::addSearchPath.
@@ -361,12 +366,13 @@ class SearchPathOptions {
   }
 
   void setCompilerPluginExecutablePaths(
-      std::vector<std::string> NewCompilerPluginExecutablePaths) {
-    CompilerPluginExecutablePaths = NewCompilerPluginExecutablePaths;
+      std::vector<PluginExecutablePathAndModuleNames> &&newValue) {
+    CompilerPluginExecutablePaths = std::move(newValue);
     Lookup.searchPathsDidChange();
   }
 
-  ArrayRef<std::string> getCompilerPluginExecutablePaths() const {
+  ArrayRef<PluginExecutablePathAndModuleNames>
+  getCompilerPluginExecutablePaths() const {
     return CompilerPluginExecutablePaths;
   }
 

lib/AST/ASTContext.cpp

@@ -531,9 +531,6 @@ struct ASTContext::Implementation {
   /// `Plugins` storage if this ASTContext owns it.
   std::unique_ptr<PluginRegistry> OwnedPluginRegistry = nullptr;
 
-  /// Cache of loaded symbols.
-  llvm::StringMap<void *> LoadedSymbols;
-
   /// Map a module name to an executable plugin path that provides the module.
   llvm::DenseMap<Identifier, StringRef> ExecutablePluginPaths;
 
@@ -707,8 +704,7 @@ ASTContext::ASTContext(
   registerAccessRequestFunctions(evaluator);
   registerNameLookupRequestFunctions(evaluator);
 
-  // FIXME: Delay this so the client e.g. SourceKit can inject plugin registry.
-  loadCompilerPlugins();
+  createModuleToExecutablePluginMap();
 }
 
 ASTContext::~ASTContext() {
@@ -6264,66 +6260,17 @@ PluginRegistry *ASTContext::getPluginRegistry() const {
   return registry;
 }
 
-void ASTContext::loadCompilerPlugins() {
-  auto fs = this->SourceMgr.getFileSystem();
-  for (auto &path : SearchPathOpts.getCompilerPluginLibraryPaths()) {
-    SmallString<128> resolvedPath;
-    if (auto err = fs->getRealPath(path, resolvedPath)) {
-      Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
-                     err.message());
-      continue;
-    }
-    auto loaded = getPluginRegistry()->loadLibraryPlugin(resolvedPath);
-    if (!loaded) {
-      Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
-                     llvm::toString(loaded.takeError()));
-    }
-  }
-
+void ASTContext::createModuleToExecutablePluginMap() {
   for (auto &arg : SearchPathOpts.getCompilerPluginExecutablePaths()) {
-    // 'arg' is '<path to executable>#<module names>' where the module names are
-    // comma separated.
-
     // Create a moduleName -> pluginPath mapping.
-    StringRef path;
-    StringRef modulesStr;
-    std::tie(path, modulesStr) = StringRef(arg).rsplit('#');
-    SmallVector<StringRef, 1> modules;
-    modulesStr.split(modules, ',');
-
-    if (modules.empty() || path.empty()) {
-      // TODO: Error messsage.
-      Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, arg, "");
-    }
-    auto pathStr = AllocateCopy(path);
-    for (auto moduleName : modules) {
+    assert(!arg.ExecutablePath.empty() && "empty plugin path");
+    auto pathStr = AllocateCopy(arg.ExecutablePath);
+    for (auto moduleName : arg.ModuleNames) {
       getImpl().ExecutablePluginPaths[getIdentifier(moduleName)] = pathStr;
     }
   }
 }
 
-void *ASTContext::getAddressOfSymbol(const char *name,
-                                     void *libraryHandleHint) {
-  auto lookup = getImpl().LoadedSymbols.try_emplace(name, nullptr);
-  void *&address = lookup.first->getValue();
-#if !defined(_WIN32)
-  if (lookup.second) {
-    auto *handle = libraryHandleHint ? libraryHandleHint : RTLD_DEFAULT;
-    address = dlsym(handle, name);
-
-    // If we didn't know where to look, look specifically in each plugin.
-    if (!address && !libraryHandleHint) {
-      for (const auto &plugin : getPluginRegistry()->getLoadedLibraryPlugins()) {
-        address = dlsym(plugin.second, name);
-        if (address)
-          break;
-      }
-    }
-  }
-#endif
-  return address;
-}
-
 Type ASTContext::getNamedSwiftType(ModuleDecl *module, StringRef name) {
   if (!module)
     return Type();
@@ -6359,6 +6306,35 @@ Type ASTContext::getNamedSwiftType(ModuleDecl *module, StringRef name) {
   return decl->getDeclaredInterfaceType();
 }
 
+Optional<std::string>
+ASTContext::lookupLibraryPluginByModuleName(Identifier moduleName) {
+  auto fs = SourceMgr.getFileSystem();
+
+  // Look for 'lib${module name}(.dylib|.so)'.
+  SmallString<64> expectedBasename;
+  expectedBasename.append("lib");
+  expectedBasename.append(moduleName.str());
+  expectedBasename.append(LTDL_SHLIB_EXT);
+
+  // Try '-plugin-path'.
+  for (const auto &searchPath : SearchPathOpts.PluginSearchPaths) {
+    SmallString<128> fullPath(searchPath);
+    llvm::sys::path::append(fullPath, expectedBasename);
+    if (fs->exists(fullPath)) {
+      return std::string(fullPath);
+    }
+  }
+
+  // Try '-load-plugin-library'.
+  for (const auto &libPath : SearchPathOpts.getCompilerPluginLibraryPaths()) {
+    if (llvm::sys::path::filename(libPath) == expectedBasename) {
+      return libPath;
+    }
+  }
+
+  return None;
+}
+
 Optional<StringRef>
 ASTContext::lookupExecutablePluginByModuleName(Identifier moduleName) {
   auto &execPluginPaths = getImpl().ExecutablePluginPaths;
@@ -6401,6 +6377,25 @@ LoadedExecutablePlugin *ASTContext::loadExecutablePlugin(StringRef path) {
   return plugin.get();
 }
 
+void *ASTContext::loadLibraryPlugin(StringRef path) {
+  SmallString<128> resolvedPath;
+  auto fs = this->SourceMgr.getFileSystem();
+  if (auto err = fs->getRealPath(path, resolvedPath)) {
+    Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
+                   err.message());
+    return nullptr;
+  }
+
+  // Load the plugin.
+  auto plugin = getPluginRegistry()->loadLibraryPlugin(resolvedPath);
+  if (!plugin) {
+    Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
+                   llvm::toString(plugin.takeError()));
+  }
+
+  return plugin.get();
+}
+
 bool ASTContext::supportsMoveOnlyTypes() const {
   // currently the only thing holding back whether the types can appear is this.
   return SILOpts.LexicalLifetimes != LexicalLifetimesOption::Off;

lib/Frontend/CompilerInvocation.cpp

@@ -1564,14 +1564,28 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts,
   }
   Opts.setCompilerPluginLibraryPaths(CompilerPluginLibraryPaths);
 
-  std::vector<std::string> CompilerPluginExecutablePaths(
+  std::vector<PluginExecutablePathAndModuleNames> CompilerPluginExecutablePaths(
       Opts.getCompilerPluginExecutablePaths());
   for (const Arg *A : Args.filtered(OPT_load_plugin_executable)) {
-    // NOTE: The value has '#<module names>' after the path.
-    // But resolveSearchPath() works as long as the value starts with a path.
-    CompilerPluginExecutablePaths.push_back(resolveSearchPath(A->getValue()));
+    // 'A' is '<path to executable>#<module names>' where the module names are
+    // comma separated.
+    StringRef path;
+    StringRef modulesStr;
+    std::tie(path, modulesStr) = StringRef(A->getValue()).rsplit('#');
+    std::vector<std::string> moduleNames;
+    for (auto name : llvm::split(modulesStr, ',')) {
+      moduleNames.emplace_back(name);
+    }
+    if (path.empty() || moduleNames.empty()) {
+      Diags.diagnose(SourceLoc(), diag::error_load_plugin_executable,
+                     A->getValue());
+    } else {
+      CompilerPluginExecutablePaths.push_back(
+          {resolveSearchPath(path), std::move(moduleNames)});
+    }
   }
-  Opts.setCompilerPluginExecutablePaths(CompilerPluginExecutablePaths);
+  Opts.setCompilerPluginExecutablePaths(
+      std::move(CompilerPluginExecutablePaths));
 
   return false;
 }

lib/Sema/TypeCheckMacros.cpp

@@ -93,7 +93,11 @@ static void const *lookupMacroTypeMetadataByExternalName(
   for (auto typeKind : typeKinds) {
     auto symbolName = Demangle::mangledNameForTypeMetadataAccessor(
         moduleName, typeName, typeKind);
-    accessorAddr = ctx.getAddressOfSymbol(symbolName.c_str(), libraryHint);
+#if !defined(_WIN32)
+    /// FIXME: 'PluginRegistry' should vend a wrapper object of the library
+    /// handle (like llvm::sys::DynamicLibrary) and dlsym should be abstracted.
+    accessorAddr = dlsym(libraryHint, symbolName.c_str());
+#endif
     if (accessorAddr)
       break;
   }
@@ -286,15 +290,16 @@ MacroDefinition MacroDefinitionRequest::evaluate(
 }
 
 /// Load a plugin library based on a module name.
-static void *loadLibraryPluginByName(StringRef searchPath, StringRef moduleName,
-                                     llvm::vfs::FileSystem &fs,
-                                     PluginRegistry *registry) {
-  SmallString<128> fullPath(searchPath);
-  llvm::sys::path::append(fullPath, "lib" + moduleName + LTDL_SHLIB_EXT);
-  if (fs.getRealPath(fullPath, fullPath))
+static void *loadLibraryPluginByName(ASTContext &ctx, Identifier moduleName) {
+  std::string libraryPath;
+  if (auto found = ctx.lookupLibraryPluginByModuleName(moduleName)) {
+    libraryPath = *found;
+  } else {
     return nullptr;
-  auto loadResult = registry->loadLibraryPlugin(fullPath);
-  return loadResult ? *loadResult : nullptr;
+  }
+
+  // Load the plugin.
+  return ctx.loadLibraryPlugin(libraryPath);
 }
 
 static LoadedExecutablePlugin *
@@ -380,12 +385,10 @@ CompilerPluginLoadRequest::evaluate(Evaluator &evaluator, ASTContext *ctx,
   auto &searchPathOpts = ctx->SearchPathOpts;
   auto *registry = ctx->getPluginRegistry();
 
-  // First, check '-plugin-path' paths.
-  for (const auto &path : searchPathOpts.PluginSearchPaths) {
-    if (auto found =
-            loadLibraryPluginByName(path, moduleName.str(), *fs, registry))
-      return LoadedCompilerPlugin::inProcess(found);
-  }
+  // Check dynamic link library plugins.
+  // i.e. '-plugin-path', and '-load-plugin-library'.
+  if (auto found = loadLibraryPluginByName(*ctx, moduleName))
+    return LoadedCompilerPlugin::inProcess(found);
 
   // Fall back to executable plugins.
   // i.e. '-external-plugin-path', and '-load-plugin-executable'.
@@ -448,17 +451,13 @@ ExternalMacroDefinitionRequest::evaluate(Evaluator &evaluator, ASTContext *ctx,
   CompilerPluginLoadRequest loadRequest{ctx, moduleName};
   LoadedCompilerPlugin loaded =
       evaluateOrDefault(evaluator, loadRequest, nullptr);
+
   if (auto loadedLibrary = loaded.getAsInProcessPlugin()) {
     if (auto inProcess = resolveInProcessMacro(
             *ctx, moduleName, typeName, loadedLibrary))
       return *inProcess;
   }
 
-  // Try to resolve in-process.
-  if (auto inProcess = resolveInProcessMacro(*ctx, moduleName, typeName))
-    return *inProcess;
-
-  // Try executable plugins.
   if (auto *executablePlugin = loaded.getAsExecutablePlugin()) {
     if (auto executableMacro = resolveExecutableMacro(*ctx, executablePlugin,
                                                       moduleName, typeName)) {