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

Load '-load-plugin-library' libraries lazily until resolving macros.

Author: rintaro

include/swift/AST/ASTContext.h

@@ -1461,10 +1461,6 @@ 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 an executable plugin that is declared to handle \p moduleName
@@ -1506,7 +1502,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,13 @@ class ModuleSearchPathLookup {
                             llvm::vfs::FileSystem *FS, bool IsOSDarwin);
 };
 
+/// Pair of a plugin search and the corresponding plugin server executable
+/// path.
+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 +249,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 +367,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;
   }
 

include/swift/AST/TypeCheckRequests.h

@@ -4054,6 +4054,27 @@ class CompilerPluginLoadRequest
   bool isCached() const { return true; }
 };
 
+/// Load a plugin library specified by '-load-plugin-library'.
+class LoadPluginLibraryRequest
+    : public SimpleRequest<LoadPluginLibraryRequest,
+                           LoadedCompilerPlugin(ASTContext *, std::string path),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  LoadedCompilerPlugin evaluate(Evaluator &evaluator, ASTContext *ctx,
+                                std::string path) const;
+
+public:
+  // Source location
+  SourceLoc getNearestLoc() const { return SourceLoc(); }
+
+  bool isCached() const { return true; }
+};
+
 /// Expand peer macros attached to the given declaration.
 class ExpandPeerMacroRequest
     : public SimpleRequest<ExpandPeerMacroRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -433,6 +433,9 @@ SWIFT_REQUEST(TypeChecker, MacroDefinitionRequest,
 SWIFT_REQUEST(TypeChecker, CompilerPluginLoadRequest,
               LoadedCompilerPlugin(ASTContext *, Identifier),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, LoadPluginLibraryRequest,
+              LoadedCompilerPlugin(ASTContext *, std::string),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExternalMacroDefinitionRequest,
               Optional<ExternalMacroDefinition>(ASTContext *, Identifier, Identifier),
               Cached, NoLocationInfo)

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();

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

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTMangler.h"
 #include "swift/AST/ASTNode.h"
+#include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/MacroDefinition.h"
 #include "swift/AST/NameLookupRequests.h"
@@ -93,7 +94,9 @@ 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)
+    accessorAddr = dlsym(libraryHint, symbolName.c_str());
+#endif
     if (accessorAddr)
       break;
   }
@@ -396,6 +399,26 @@ CompilerPluginLoadRequest::evaluate(Evaluator &evaluator, ASTContext *ctx,
   return nullptr;
 }
 
+LoadedCompilerPlugin
+LoadPluginLibraryRequest::evaluate(Evaluator &evaluator, ASTContext *ctx,
+                                   std::string path) const {
+  auto fs = ctx->SourceMgr.getFileSystem();
+  auto *registry = ctx->getPluginRegistry();
+  auto &diags = ctx->Diags;
+
+  SmallString<128> resolvedPath;
+  if (fs->getRealPath(path, resolvedPath))
+    return nullptr;
+  auto loaded = registry->loadLibraryPlugin(resolvedPath);
+  if (!loaded) {
+    diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
+                   llvm::toString(loaded.takeError()));
+    return nullptr;
+  }
+
+  return LoadedCompilerPlugin::inProcess(loaded.get());
+}
+
 static Optional<ExternalMacroDefinition>
 resolveInProcessMacro(
     ASTContext &ctx, Identifier moduleName, Identifier typeName,
@@ -448,15 +471,25 @@ ExternalMacroDefinitionRequest::evaluate(Evaluator &evaluator, ASTContext *ctx,
   CompilerPluginLoadRequest loadRequest{ctx, moduleName};
   LoadedCompilerPlugin loaded =
       evaluateOrDefault(evaluator, loadRequest, nullptr);
+
+  // Loaded in '-plugin-path'?
   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 to resolve in '-load-plugin-library' plugins.
+  for (auto &path : ctx->SearchPathOpts.getCompilerPluginLibraryPaths()) {
+    LoadPluginLibraryRequest loadRequest{ctx, path};
+    LoadedCompilerPlugin loaded =
+        evaluateOrDefault(evaluator, loadRequest, nullptr);
+    if (auto loadedLibrary = loaded.getAsInProcessPlugin()) {
+      if (auto inProcess =
+              resolveInProcessMacro(*ctx, moduleName, typeName, loadedLibrary))
+        return *inProcess;
+    }
+  }
 
   // Try executable plugins.
   if (auto *executablePlugin = loaded.getAsExecutablePlugin()) {