Merge pull request #64655 from rintaro/sourcekit-single-pluginregistry

[SourceKit] Use a single PluginRegistry in multiple ASTContexts

Author: rintaro

include/swift/AST/ASTContext.h

@@ -1503,6 +1503,8 @@ class ASTContext final {
   /// This should be called before any plugin is loaded.
   void setPluginRegistry(PluginRegistry *newValue);
 
+  const llvm::StringSet<> &getLoadedPluginLibraryPaths() const;
+
 private:
   friend Decl;
 

include/swift/AST/PluginRegistry.h

@@ -139,6 +139,8 @@ class PluginRegistry {
   /// Flag to dump plugin messagings.
   bool dumpMessaging = false;
 
+  std::mutex mtx;
+
 public:
   PluginRegistry();
 
@@ -150,10 +152,6 @@ class PluginRegistry {
   /// If \p path plugin is already loaded, this returns the cached object.
   llvm::Expected<LoadedExecutablePlugin *>
   loadExecutablePlugin(llvm::StringRef path);
-
-  const llvm::StringMap<void *> &getLoadedLibraryPlugins() const {
-    return LoadedPluginLibraries;
-  }
 };
 
 } // namespace swift

include/swift/IDETool/CompileInstance.h

@@ -25,13 +25,15 @@ namespace swift {
 
 class CompilerInstance;
 class DiagnosticConsumer;
+class PluginRegistry;
 
 namespace ide {
 
 /// Manages \c CompilerInstance for completion like operations.
 class CompileInstance {
   const std::string &RuntimeResourcePath;
   const std::string &DiagnosticDocumentationPath;
+  const std::shared_ptr<swift::PluginRegistry> Plugins;
 
   struct Options {
     unsigned MaxASTReuseCount = 100;
@@ -66,10 +68,11 @@ class CompileInstance {
 
 public:
   CompileInstance(const std::string &RuntimeResourcePath,
-                  const std::string &DiagnosticDocumentationPath)
+                  const std::string &DiagnosticDocumentationPath,
+                  std::shared_ptr<swift::PluginRegistry> Plugins = nullptr)
       : RuntimeResourcePath(RuntimeResourcePath),
         DiagnosticDocumentationPath(DiagnosticDocumentationPath),
-        CachedCIInvalidated(false), CachedReuseCount(0) {}
+        Plugins(Plugins), CachedCIInvalidated(false), CachedReuseCount(0) {}
 
   /// NOTE: \p Args is only used for checking the equaity of the invocation.
   /// Since this function assumes that it is already normalized, exact the same

include/swift/IDETool/IDEInspectionInstance.h

@@ -35,6 +35,7 @@ namespace swift {
 class CompilerInstance;
 class CompilerInvocation;
 class DiagnosticConsumer;
+class PluginRegistry;
 
 namespace ide {
 
@@ -96,6 +97,8 @@ class IDEInspectionInstance {
 
   std::mutex mtx;
 
+  std::shared_ptr<PluginRegistry> Plugins;
+
   std::shared_ptr<CompilerInstance> CachedCI;
   llvm::hash_code CachedArgHash;
   llvm::sys::TimePoint<> DependencyCheckedTimestamp;
@@ -167,7 +170,8 @@ class IDEInspectionInstance {
           Callback);
 
 public:
-  IDEInspectionInstance() : CachedCIShouldBeInvalidated(false) {}
+  IDEInspectionInstance(std::shared_ptr<PluginRegistry> Plugins = nullptr)
+      : Plugins(Plugins), CachedCIShouldBeInvalidated(false) {}
 
   // Mark the cached compiler instance "should be invalidated". In the next
   // completion, new compiler instance will be used. (Thread safe.)

lib/AST/ASTContext.cpp

@@ -534,6 +534,8 @@ struct ASTContext::Implementation {
   /// Map a module name to an executable plugin path that provides the module.
   llvm::DenseMap<Identifier, StringRef> ExecutablePluginPaths;
 
+  llvm::StringSet<> LoadedPluginLibraryPaths;
+
   /// The permanent arena.
   Arena Permanent;
 
@@ -6395,6 +6397,9 @@ LoadedExecutablePlugin *ASTContext::loadExecutablePlugin(StringRef path) {
 }
 
 void *ASTContext::loadLibraryPlugin(StringRef path) {
+  // Remember the path (even if it fails to load.)
+  getImpl().LoadedPluginLibraryPaths.insert(path);
+
   SmallString<128> resolvedPath;
   auto fs = this->SourceMgr.getFileSystem();
   if (auto err = fs->getRealPath(path, resolvedPath)) {
@@ -6413,6 +6418,10 @@ void *ASTContext::loadLibraryPlugin(StringRef path) {
   return plugin.get();
 }
 
+const llvm::StringSet<> &ASTContext::getLoadedPluginLibraryPaths() const {
+  return getImpl().LoadedPluginLibraryPaths;
+}
+
 bool ASTContext::supportsMoveOnlyTypes() const {
   // currently the only thing holding back whether the types can appear is this.
   return SILOpts.LexicalLifetimes != LexicalLifetimesOption::Off;

lib/AST/PluginRegistry.cpp

@@ -45,6 +45,7 @@ PluginRegistry::PluginRegistry() {
 }
 
 llvm::Expected<void *> PluginRegistry::loadLibraryPlugin(StringRef path) {
+  std::lock_guard<std::mutex> lock(mtx);
   auto found = LoadedPluginLibraries.find(path);
   if (found != LoadedPluginLibraries.end()) {
     // Already loaded.
@@ -74,6 +75,8 @@ PluginRegistry::loadExecutablePlugin(StringRef path) {
     return llvm::errorCodeToError(err);
   }
 
+  std::lock_guard<std::mutex> lock(mtx);
+
   // See if the plugin is already loaded.
   auto &storage = LoadedPluginExecutables[path];
   if (storage) {

lib/FrontendTool/LoadedModuleTrace.cpp

@@ -760,8 +760,7 @@ bool swift::emitLoadedModuleTraceIfNeeded(ModuleDecl *mainModule,
   }
 
   // Add compiler plugin libraries as dependencies.
-  auto *pluginRegistry = ctxt.getPluginRegistry();
-  for (auto &pluginEntry : pluginRegistry->getLoadedLibraryPlugins())
+  for (auto &pluginEntry : ctxt.getLoadedPluginLibraryPaths())
     depTracker->addDependency(pluginEntry.getKey(), /*IsSystem*/ false);
 
   std::vector<SwiftModuleTraceInfo> swiftModules;

lib/IDETool/CompileInstance.cpp

@@ -299,6 +299,7 @@ bool CompileInstance::setupCI(
     assert(Diags.hadAnyError());
     return false;
   }
+  CI->getASTContext().setPluginRegistry(Plugins.get());
 
   return true;
 }

lib/IDETool/IDEInspectionInstance.cpp

@@ -482,6 +482,7 @@ void IDEInspectionInstance::performNewOperation(
           InstanceSetupError));
       return;
     }
+    CI->getASTContext().setPluginRegistry(Plugins.get());
     CI->getASTContext().CancellationFlag = CancellationFlag;
     registerIDERequestFunctions(CI->getASTContext().evaluator);
 

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -553,10 +553,12 @@ struct SwiftASTManager::Implementation {
       std::shared_ptr<SwiftEditorDocumentFileMap> EditorDocs,
       std::shared_ptr<GlobalConfig> Config,
       std::shared_ptr<SwiftStatistics> Stats,
-      std::shared_ptr<RequestTracker> ReqTracker, StringRef SwiftExecutablePath,
+      std::shared_ptr<RequestTracker> ReqTracker,
+      std::shared_ptr<PluginRegistry> Plugins, StringRef SwiftExecutablePath,
       StringRef RuntimeResourcePath, StringRef DiagnosticDocumentationPath)
       : EditorDocs(EditorDocs), Config(Config), Stats(Stats),
-        ReqTracker(ReqTracker), SwiftExecutablePath(SwiftExecutablePath),
+        ReqTracker(ReqTracker), Plugins(Plugins),
+        SwiftExecutablePath(SwiftExecutablePath),
         RuntimeResourcePath(RuntimeResourcePath),
         DiagnosticDocumentationPath(DiagnosticDocumentationPath),
         SessionTimestamp(llvm::sys::toTimeT(std::chrono::system_clock::now())) {
@@ -566,6 +568,7 @@ struct SwiftASTManager::Implementation {
   std::shared_ptr<GlobalConfig> Config;
   std::shared_ptr<SwiftStatistics> Stats;
   std::shared_ptr<RequestTracker> ReqTracker;
+  std::shared_ptr<PluginRegistry> Plugins;
   /// The path of the swift-frontend executable.
   /// Used to find clang relative to it.
   std::string SwiftExecutablePath;
@@ -638,9 +641,10 @@ SwiftASTManager::SwiftASTManager(
     std::shared_ptr<SwiftEditorDocumentFileMap> EditorDocs,
     std::shared_ptr<GlobalConfig> Config,
     std::shared_ptr<SwiftStatistics> Stats,
-    std::shared_ptr<RequestTracker> ReqTracker, StringRef SwiftExecutablePath,
+    std::shared_ptr<RequestTracker> ReqTracker,
+    std::shared_ptr<PluginRegistry> Plugins, StringRef SwiftExecutablePath,
     StringRef RuntimeResourcePath, StringRef DiagnosticDocumentationPath)
-    : Impl(*new Implementation(EditorDocs, Config, Stats, ReqTracker,
+    : Impl(*new Implementation(EditorDocs, Config, Stats, ReqTracker, Plugins,
                                SwiftExecutablePath, RuntimeResourcePath,
                                DiagnosticDocumentationPath)) {}
 
@@ -1073,6 +1077,7 @@ ASTUnitRef ASTBuildOperation::buildASTUnit(std::string &Error) {
     }
     return nullptr;
   }
+  CompIns.getASTContext().setPluginRegistry(ASTManager->Impl.Plugins.get());
   CompIns.getASTContext().CancellationFlag = CancellationFlag;
   registerIDERequestFunctions(CompIns.getASTContext().evaluator);
   if (TracedOp.enabled()) {