Merge pull request #80474 from cachemeifyoucan/eng/PR-148465899

[Caching][Macro] Make macro plugin options cacheable

Author: cachemeifyoucan

include/swift/AST/DiagnosticsSema.def

@@ -1265,6 +1265,10 @@ REMARK(macro_loaded,none,
        "compiler plugin server '%2' with shared library '%3'}1",
        (Identifier, unsigned, StringRef, StringRef))
 
+ERROR(resolved_macro_changed,none,
+      "resolved macro library '%0' failed verification: %1",
+      (StringRef, StringRef))
+
 REMARK(transitive_dependency_behavior,none,
        "%1 has %select{a required|an optional|an ignored}2 "
        "transitive dependency on '%0'",

include/swift/AST/PluginLoader.h

@@ -52,10 +52,17 @@ class PluginLoader {
   /// Get or lazily create and populate 'PluginMap'.
   llvm::DenseMap<swift::Identifier, PluginEntry> &getPluginMap();
 
+  /// Resolved plugin path remappings.
+  std::vector<std::string> PathRemap;
+
 public:
   PluginLoader(ASTContext &Ctx, DependencyTracker *DepTracker,
+               std::optional<std::vector<std::string>> Remap = std::nullopt,
                bool disableSandbox = false)
-      : Ctx(Ctx), DepTracker(DepTracker), disableSandbox(disableSandbox) {}
+      : Ctx(Ctx), DepTracker(DepTracker), disableSandbox(disableSandbox) {
+    if (Remap)
+      PathRemap = std::move(*Remap);
+  }
 
   void setRegistry(PluginRegistry *newValue);
   PluginRegistry *getRegistry();

include/swift/AST/SearchPathOptions.h

@@ -510,6 +510,9 @@ class SearchPathOptions {
   /// Scanner Prefix Mapper.
   std::vector<std::string> ScannerPrefixMapper;
 
+  /// Verify resolved plugin is not changed.
+  bool ResolvedPluginVerification = false;
+
   /// When set, don't validate module system dependencies.
   ///
   /// If a system header is modified and this is not set, the compiler will

include/swift/Option/Options.td

@@ -2263,6 +2263,10 @@ def load_resolved_plugin:
            "and exectuable path can be empty if not used">,
   MetaVarName<"<library-path>#<executable-path>#<module-names>">;
 
+def resolved_plugin_verification : Flag<["-"], "resolved-plugin-verification">,
+                                   Flags<[FrontendOption, NoDriverOption]>,
+                                   HelpText<"verify resolved plugins">;
+
 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

@@ -13,11 +13,12 @@
 #include "swift/AST/PluginLoader.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/DiagnosticEngine.h"
-#include "swift/AST/DiagnosticsFrontend.h"
+#include "swift/AST/DiagnosticsSema.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Parse/Lexer.h"
 #include "llvm/Config/config.h"
+#include "llvm/Support/PrefixMapper.h"
 #include "llvm/Support/VirtualFileSystem.h"
 
 using namespace swift;
@@ -96,9 +97,45 @@ PluginLoader::getPluginMap() {
     map[moduleNameIdentifier] = {libPath, execPath};
   };
 
+  std::optional<llvm::PrefixMapper> mapper;
+  if (!PathRemap.empty()) {
+    SmallVector<llvm::MappedPrefix, 4> prefixes;
+    llvm::MappedPrefix::transformJoinedIfValid(PathRemap, prefixes);
+    mapper.emplace();
+    mapper->addRange(prefixes);
+    mapper->sort();
+  }
+  auto remapPath = [&mapper](StringRef path) {
+    if (!mapper)
+      return path.str();
+    return mapper->mapToString(path);
+  };
+
   auto fs = getPluginLoadingFS(Ctx);
   std::error_code ec;
 
+  auto validateLibrary = [&](StringRef path) -> llvm::Expected<std::string> {
+    auto remappedPath = remapPath(path);
+    if (!Ctx.SearchPathOpts.ResolvedPluginVerification || path.empty())
+      return remappedPath;
+
+    auto currentStat = fs->status(remappedPath);
+    if (!currentStat)
+      return llvm::createFileError(remappedPath, currentStat.getError());
+
+    auto goldStat = Ctx.SourceMgr.getFileSystem()->status(path);
+    if (!goldStat)
+      return llvm::createStringError(
+          "cannot open gold reference library to compare");
+
+    // Compare the size for difference for now.
+    if (currentStat->getSize() != goldStat->getSize())
+      return llvm::createStringError(
+          "plugin has changed since dependency scanning");
+
+    return remappedPath;
+  };
+
   for (auto &entry : Ctx.SearchPathOpts.PluginSearchOpts) {
     switch (entry.getKind()) {
 
@@ -156,9 +193,17 @@ PluginLoader::getPluginMap() {
       // 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,
+      for (auto &moduleName : val.ModuleNames) {
+        auto libPath = validateLibrary(val.LibraryPath);
+        if (!libPath) {
+          Ctx.Diags.diagnose(SourceLoc(), diag::resolved_macro_changed,
+                             remapPath(val.LibraryPath),
+                             toString(libPath.takeError()));
+          continue;
+        }
+        try_emplace(moduleName, *libPath, remapPath(val.ExecutablePath),
                     /*overwrite*/ true);
+      }
       continue;
     }
     }

lib/DependencyScan/ModuleDependencyScanner.cpp

@@ -201,6 +201,7 @@ ModuleDependencyScanningWorker::ModuleDependencyScanningWorker(
                       ScanASTContext.SourceMgr, Diagnostics));
   auto loader = std::make_unique<PluginLoader>(
       *workerASTContext, /*DepTracker=*/nullptr,
+      workerCompilerInvocation->getFrontendOptions().CacheReplayPrefixMap,
       workerCompilerInvocation->getFrontendOptions().DisableSandbox);
   workerASTContext->setPluginLoader(std::move(loader));
 

lib/DependencyScan/ScanDependencies.cpp

@@ -214,9 +214,23 @@ class ExplicitModuleDependencyResolver {
       depInfo.addMacroDependency(macro.first(), macro.second.LibraryPath,
                                  macro.second.ExecutablePath);
 
+    bool needPathRemapping = instance.getInvocation()
+                                 .getSearchPathOptions()
+                                 .ResolvedPluginVerification &&
+                             cache.getScanService().hasPathMapping();
+    auto mapPath = [&](StringRef path) {
+      if (!needPathRemapping)
+        return path.str();
+
+      return cache.getScanService().remapPath(path);
+    };
+    if (needPathRemapping)
+      commandline.push_back("-resolved-plugin-verification");
+
     for (auto &macro : depInfo.getMacroDependencies()) {
-      std::string arg = macro.second.LibraryPath + "#" +
-                        macro.second.ExecutablePath + "#" + macro.first;
+      std::string arg = mapPath(macro.second.LibraryPath) + "#" +
+                        mapPath(macro.second.ExecutablePath) + "#" +
+                        macro.first;
       commandline.push_back("-load-resolved-plugin");
       commandline.push_back(arg);
     }
@@ -480,9 +494,10 @@ class ExplicitModuleDependencyResolver {
       llvm::for_each(
           sourceDep->auxiliaryFiles,
           [this](const std::string &file) { tracker->trackFile(file); });
-      llvm::for_each(sourceDep->macroDependencies, [this](const auto &entry) {
-        tracker->trackFile(entry.second.LibraryPath);
-      });
+      llvm::for_each(dependencyInfoCopy.getMacroDependencies(),
+                     [this](const auto &entry) {
+                       tracker->trackFile(entry.second.LibraryPath);
+                     });
       auto root = tracker->createTreeFromDependencies();
       if (!root)
         return root.takeError();
@@ -496,7 +511,7 @@ class ExplicitModuleDependencyResolver {
       llvm::for_each(
           textualDep->auxiliaryFiles,
           [this](const std::string &file) { tracker->trackFile(file); });
-      llvm::for_each(textualDep->macroDependencies,
+      llvm::for_each(dependencyInfoCopy.getMacroDependencies(),
                      [this](const auto &entry) {
                        tracker->trackFile(entry.second.LibraryPath);
                      });

lib/Frontend/CachedDiagnostics.cpp

@@ -18,7 +18,9 @@
 
 #include "swift/AST/DiagnosticBridge.h"
 #include "swift/AST/DiagnosticConsumer.h"
+#include "swift/AST/DiagnosticsCommon.h"
 #include "swift/AST/DiagnosticsFrontend.h"
+#include "swift/AST/DiagnosticsSema.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Frontend/Frontend.h"
@@ -755,6 +757,13 @@ class CachingDiagnosticsProcessor::Implementation
     auto &Serializer = getSerializer();
     assert(SM.getFileSystem() == Serializer.getSourceMgr().getFileSystem() &&
            "Caching for a different file system");
+
+    // Bypass the caching.
+    if (BypassDiagIDs.count(Info.ID)) {
+      for (auto *Diag : OrigConsumers)
+        Diag->handleDiagnostic(Serializer.getSourceMgr(), Info);
+      return;
+    }
     Serializer.handleDiagnostic(SM, Info, [&](const DiagnosticInfo &Info) {
       for (auto *Diag : OrigConsumers)
         Diag->handleDiagnostic(Serializer.getSourceMgr(), Info);
@@ -809,6 +818,8 @@ class CachingDiagnosticsProcessor::Implementation
   // Processor/Serializer alive until then.
   std::unique_ptr<DiagnosticSerializer> Serializer;
 
+  const llvm::SmallDenseSet<DiagID> BypassDiagIDs = {diag::macro_loaded.ID};
+
   SourceManager &InstanceSourceMgr;
   const FrontendInputsAndOutputs &InAndOut;
   DiagnosticEngine &Diags;

lib/Frontend/CompilerInvocation.cpp

@@ -2417,6 +2417,9 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts, ArgList &Args,
     Opts.ScannerPrefixMapper.push_back(Opt.str());
   }
 
+  Opts.ResolvedPluginVerification |=
+      Args.hasArg(OPT_resolved_plugin_verification);
+
   // rdar://132340493 disable scanner-side validation for non-caching builds
   Opts.ScannerModuleValidation |= Args.hasFlag(OPT_scanner_module_validation,
                                                OPT_no_scanner_module_validation,

lib/Frontend/Frontend.cpp

@@ -878,6 +878,7 @@ bool CompilerInstance::setUpPluginLoader() {
   /// FIXME: If Invocation has 'PluginRegistry', we can set it. But should we?
   auto loader = std::make_unique<PluginLoader>(
       *Context, getDependencyTracker(),
+      Invocation.getFrontendOptions().CacheReplayPrefixMap,
       Invocation.getFrontendOptions().DisableSandbox);
   Context->setPluginLoader(std::move(loader));
   return false;