[Macro] Add a new macro loading option that do not involve searching

Add flag `-load-resolved-plugin` to load macro plugin, which provides a
pre-resolved entry into PluginLoader so the plugins can be loaded based
on module name without searching the file system. The option is mainly
intended to be used by explicitly module build and the flag is supplied
by dependency scanner.

Author: cachemeifyoucan

include/swift/AST/DiagnosticsFrontend.def

@@ -127,6 +127,9 @@ ERROR(error_no_source_location_scope_map,none,
 ERROR(error_load_plugin_executable,none,
       "invalid value '%0' in '-load-plugin-executable'; "
       "make sure to use format '<plugin path>#<module names>'", (StringRef))
+ERROR(error_load_resolved_plugin,none,
+      "invalid value '%0' in '-load-resolved-plugin'; "
+      "make sure to use format '<library path>#<plugin path>#<module names>' where library and plugin path can't both be empty", (StringRef))
 
 NOTE(note_valid_swift_versions, none,
      "valid arguments to '-swift-version' are %0", (StringRef))

include/swift/AST/SearchPathOptions.h

@@ -214,17 +214,24 @@ class PluginSearchOption {
     std::string SearchPath;
     std::string ServerPath;
   };
+  struct ResolvedPluginConfig {
+    std::string LibraryPath;
+    std::string ExecutablePath;
+    std::vector<std::string> ModuleNames;
+  };
 
   enum class Kind : uint8_t {
     LoadPluginLibrary,
     LoadPluginExecutable,
     PluginPath,
     ExternalPluginPath,
+    ResolvedPluginConfig,
   };
 
 private:
-  using Members = ExternalUnionMembers<LoadPluginLibrary, LoadPluginExecutable,
-                                       PluginPath, ExternalPluginPath>;
+  using Members =
+      ExternalUnionMembers<LoadPluginLibrary, LoadPluginExecutable, PluginPath,
+                           ExternalPluginPath, ResolvedPluginConfig>;
   static Members::Index getIndexForKind(Kind kind) {
     switch (kind) {
     case Kind::LoadPluginLibrary:
@@ -235,6 +242,8 @@ class PluginSearchOption {
       return Members::indexOf<PluginPath>();
     case Kind::ExternalPluginPath:
       return Members::indexOf<ExternalPluginPath>();
+    case Kind::ResolvedPluginConfig:
+      return Members::indexOf<ResolvedPluginConfig>();
     }
   };
   using Storage = ExternalUnion<Kind, Members, getIndexForKind>;
@@ -258,6 +267,10 @@ class PluginSearchOption {
       : kind(Kind::ExternalPluginPath) {
     storage.emplace<ExternalPluginPath>(kind, v);
   }
+  PluginSearchOption(const ResolvedPluginConfig &v)
+      : kind(Kind::ResolvedPluginConfig) {
+    storage.emplace<ResolvedPluginConfig>(kind, v);
+  }
   PluginSearchOption(const PluginSearchOption &o) : kind(o.kind) {
     storage.copyConstruct(o.kind, o.storage);
   }

include/swift/Option/Options.td

@@ -2095,6 +2095,14 @@ def load_plugin_executable:
            "of module names where the macro types are declared">,
   MetaVarName<"<path>#<module-names>">;
 
+def load_resolved_plugin:
+  Separate<["-"], "load-resolved-plugin">, Group<plugin_search_Group>,
+  Flags<[FrontendOption, DoesNotAffectIncrementalBuild, ArgumentIsPath]>,
+  HelpText<"Path to resolved plugin configuration and a comma-separated list "
+           "of module names where the macro types are declared. Library path "
+           "and exectuable path can be empty if not used">,
+  MetaVarName<"<library-path>#<executable-path>#<module-names>">;
+
 def in_process_plugin_server_path : Separate<["-"], "in-process-plugin-server-path">,
   Flags<[FrontendOption, ArgumentIsPath]>,
   HelpText<"Path to dynamic library plugin server">;

lib/AST/PluginLoader.cpp

@@ -83,18 +83,17 @@ PluginLoader::getPluginMap() {
   // 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) {
+                         StringRef execPath, bool overwrite = false) {
     auto moduleNameIdentifier = Ctx.getIdentifier(moduleName);
-    if (map.find(moduleNameIdentifier) != map.end()) {
-      // Specified module name is already in the map.
+    if (map.find(moduleNameIdentifier) != map.end() && !overwrite) {
+      // Specified module name is already in the map and no need to overwrite
+      // the current value.
       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;
+    map[moduleNameIdentifier] = {libPath, execPath};
   };
 
   auto fs = getPluginLoadingFS(Ctx);
@@ -150,6 +149,18 @@ PluginLoader::getPluginMap() {
       }
       continue;
     }
+
+    // '-load-resolved-plugin <library path>#<server path>#<module name>,...'.
+    case PluginSearchOption::Kind::ResolvedPluginConfig: {
+      auto &val = entry.get<PluginSearchOption::ResolvedPluginConfig>();
+      // Respect resolved plugin config above other search path, and it can
+      // overwrite plugins found by other options or previous resolved
+      // configuration.
+      for (auto &moduleName : val.ModuleNames)
+        try_emplace(moduleName, val.LibraryPath, val.ExecutablePath,
+                    /*overwrite*/ true);
+      continue;
+    }
     }
     llvm_unreachable("unhandled PluginSearchOption::Kind");
   }

lib/Frontend/CompilerInvocation.cpp

@@ -2169,6 +2169,27 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts, ArgList &Args,
           resolveSearchPath(dylibPath), resolveSearchPath(serverPath)});
       break;
     }
+    case OPT_load_resolved_plugin: {
+      StringRef libraryPath;
+      StringRef executablePath;
+      StringRef modulesStr;
+      std::tie(libraryPath, executablePath) =
+          StringRef(A->getValue()).split('#');
+      std::tie(executablePath, modulesStr) = executablePath.split('#');
+      if (modulesStr.empty() ||
+          (libraryPath.empty() && executablePath.empty())) {
+        Diags.diagnose(SourceLoc(), diag::error_load_resolved_plugin,
+                       A->getValue());
+      }
+      std::vector<std::string> moduleNames;
+      for (auto name : llvm::split(modulesStr, ',')) {
+        moduleNames.emplace_back(name);
+      }
+      Opts.PluginSearchOpts.emplace_back(
+          PluginSearchOption::ResolvedPluginConfig{
+              libraryPath.str(), executablePath.str(), std::move(moduleNames)});
+      break;
+    }
     default:
       llvm_unreachable("unhandled plugin search option");
     }

lib/Frontend/ModuleInterfaceLoader.cpp

@@ -1753,6 +1753,15 @@ void InterfaceSubContextDelegateImpl::inheritOptionsForBuildingInterface(
           ArgSaver.save(val.SearchPath + "#" + val.ServerPath));
       break;
     }
+    case PluginSearchOption::Kind::ResolvedPluginConfig: {
+      auto &val = entry.get<PluginSearchOption::ResolvedPluginConfig>();
+      for (auto &moduleName : val.ModuleNames) {
+        GenericArgs.push_back("-load-plugin-executable");
+        GenericArgs.push_back(ArgSaver.save(
+            val.LibraryPath + "#" + val.ExecutablePath + "#" + moduleName));
+      }
+      break;
+    }
     }
   }
 

lib/Serialization/ModuleFileSharedCore.cpp

@@ -148,6 +148,9 @@ static bool readOptionsBlock(llvm::BitstreamCursor &cursor,
       case PluginSearchOptionKind::LoadPluginExecutable:
         optKind = PluginSearchOption::Kind::LoadPluginExecutable;
         break;
+      case PluginSearchOptionKind::ResolvedPluginConfig:
+        optKind = PluginSearchOption::Kind::ResolvedPluginConfig;
+        break;
       }
       extendedInfo.addPluginSearchOption({optKind, blobData});
       break;

lib/Serialization/ModuleFormat.h

@@ -694,6 +694,7 @@ enum class PluginSearchOptionKind : uint8_t {
   ExternalPluginPath,
   LoadPluginLibrary,
   LoadPluginExecutable,
+  ResolvedPluginConfig,
 };
 using PluginSearchOptionKindField = BCFixed<3>;
 

lib/Serialization/Serialization.cpp

@@ -1245,6 +1245,18 @@ void Serializer::writeHeader() {
                 uint8_t(PluginSearchOptionKind::LoadPluginExecutable), optStr);
             continue;
           }
+          case PluginSearchOption::Kind::ResolvedPluginConfig: {
+            auto &opt = elem.get<PluginSearchOption::ResolvedPluginConfig>();
+            std::string optStr =
+                opt.LibraryPath + "#" + opt.ExecutablePath + "#";
+            llvm::interleave(
+                opt.ModuleNames, [&](auto &name) { optStr += name; },
+                [&]() { optStr += ","; });
+            PluginSearchOpt.emit(
+                ScratchRecord,
+                uint8_t(PluginSearchOptionKind::ResolvedPluginConfig), optStr);
+            continue;
+          }
           }
         }
       }

test/Macros/macro_plugin_server.swift

@@ -31,6 +31,26 @@
 
 // RUN: %FileCheck -strict-whitespace %s < %t/macro-expansions.txt
 
+/// Create file with matching name in alt directories and test
+/// `-load-resolved-plugin` takes the priority over other options.
+// RUN: %empty-directory(%t/alt)
+// RUN: touch %t/alt/%target-library-name(MacroDefinition)
+
+// RUN: env SWIFT_DUMP_PLUGIN_MESSAGING=1 %target-swift-frontend \
+// RUN:   -typecheck -verify \
+// RUN:   -swift-version 5 -enable-experimental-feature Macros \
+// RUN:   -external-plugin-path %t/alt#%swift-plugin-server \
+// RUN:   -load-resolved-plugin lib-do-not-exist.dylib##MacroDefinition \
+// RUN:   -load-resolved-plugin %t/plugins/%target-library-name(MacroDefinition)#%swift-plugin-server#MacroDefinition \
+// RUN:   -load-resolved-plugin %t/plugins/%target-library-name(EvilMacros)#%swift-plugin-server#EvilMacros \
+// RUN:   -external-plugin-path %t/alt#%swift-plugin-server \
+// RUN:   -Rmacro-loading -verify-ignore-unknown \
+// RUN:   -module-name MyApp \
+// RUN:   %s \
+// RUN:   2>&1 | tee %t/macro-expansions-2.txt
+
+// RUN: %FileCheck -strict-whitespace %s < %t/macro-expansions-2.txt
+
 // RUN: not %target-swift-frontend \
 // RUN:   -typecheck \
 // RUN:   -swift-version 5 \