[SourceKit] Use a single PluginRegistry in multiple ASTContexts

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

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -95,6 +95,7 @@ namespace swift {
   class CompilerInstance;
   class CompilerInvocation;
   class DiagnosticEngine;
+  class PluginRegistry;
   class SourceFile;
   class SourceManager;
 }
@@ -235,6 +236,7 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
                            std::shared_ptr<GlobalConfig> Config,
                            std::shared_ptr<SwiftStatistics> Stats,
                            std::shared_ptr<RequestTracker> ReqTracker,
+                           std::shared_ptr<swift::PluginRegistry> Plugins,
                            StringRef SwiftExecutablePath,
                            StringRef RuntimeResourcePath,
                            StringRef DiagnosticDocumentationPath);

tools/SourceKit/lib/SwiftLang/SwiftCompile.cpp

@@ -31,7 +31,9 @@ compile::SessionManager::getSession(StringRef name) {
   }
 
   bool inserted = false;
-  std::tie(i, inserted) = sessions.try_emplace(name, std::make_shared<compile::Session>(RuntimeResourcePath, DiagnosticDocumentationPath));
+  std::tie(i, inserted) = sessions.try_emplace(
+      name, std::make_shared<compile::Session>(
+                RuntimeResourcePath, DiagnosticDocumentationPath, Plugins));
   assert(inserted);
   return i->second;
 }
@@ -141,10 +143,10 @@ void SwiftLangSupport::performCompile(
     CancellationFlag->store(true, std::memory_order_relaxed);
   });
 
-  CompileManager.performCompileAsync(Name, Args, std::move(fileSystem),
-                                     CancellationFlag, Receiver);
+  CompileManager->performCompileAsync(Name, Args, std::move(fileSystem),
+                                      CancellationFlag, Receiver);
 }
 
 void SwiftLangSupport::closeCompile(StringRef Name) {
-  CompileManager.clearSession(Name);
+  CompileManager->clearSession(Name);
 }

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

@@ -276,24 +276,29 @@ static void configureIDEInspectionInstance(
 SwiftLangSupport::SwiftLangSupport(SourceKit::Context &SKCtx)
     : NotificationCtr(SKCtx.getNotificationCenter()),
       SwiftExecutablePath(SKCtx.getSwiftExecutablePath()),
-      ReqTracker(SKCtx.getRequestTracker()), CCCache(new SwiftCompletionCache),
-      CompileManager(RuntimeResourcePath, DiagnosticDocumentationPath) {
+      ReqTracker(SKCtx.getRequestTracker()), CCCache(new SwiftCompletionCache) {
   llvm::SmallString<128> LibPath(SKCtx.getRuntimeLibPath());
   llvm::sys::path::append(LibPath, "swift");
   RuntimeResourcePath = std::string(LibPath.str());
   DiagnosticDocumentationPath = SKCtx.getDiagnosticDocumentationPath().str();
 
   Stats = std::make_shared<SwiftStatistics>();
   EditorDocuments = std::make_shared<SwiftEditorDocumentFileMap>();
+
+  std::shared_ptr<PluginRegistry> Plugins = std::make_shared<PluginRegistry>();
+
   ASTMgr = std::make_shared<SwiftASTManager>(
       EditorDocuments, SKCtx.getGlobalConfiguration(), Stats, ReqTracker,
-      SwiftExecutablePath, RuntimeResourcePath, DiagnosticDocumentationPath);
-
-  IDEInspectionInst = std::make_shared<IDEInspectionInstance>();
+      Plugins, SwiftExecutablePath, RuntimeResourcePath,
+      DiagnosticDocumentationPath);
 
+  IDEInspectionInst = std::make_shared<IDEInspectionInstance>(Plugins);
   configureIDEInspectionInstance(IDEInspectionInst,
                                  SKCtx.getGlobalConfiguration());
 
+  CompileManager = std::make_shared<compile::SessionManager>(
+      RuntimeResourcePath, DiagnosticDocumentationPath, Plugins);
+
   // By default, just use the in-memory cache.
   CCCache->inMemory = std::make_unique<ide::CodeCompletionCache>();
 

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -293,8 +293,9 @@ class Session {
 
 public:
   Session(const std::string &RuntimeResourcePath,
-          const std::string &DiagnosticDocumentationPath)
-  : Compiler(RuntimeResourcePath, DiagnosticDocumentationPath) {}
+          const std::string &DiagnosticDocumentationPath,
+          std::shared_ptr<swift::PluginRegistry> Plugins)
+      : Compiler(RuntimeResourcePath, DiagnosticDocumentationPath, Plugins) {}
 
   bool
   performCompile(llvm::ArrayRef<const char *> Args,
@@ -308,6 +309,7 @@ class Session {
 class SessionManager {
   const std::string &RuntimeResourcePath;
   const std::string &DiagnosticDocumentationPath;
+  const std::shared_ptr<swift::PluginRegistry> Plugins;
 
   llvm::StringMap<std::shared_ptr<Session>> sessions;
   WorkQueue compileQueue{WorkQueue::Dequeuing::Concurrent,
@@ -316,9 +318,11 @@ class SessionManager {
 
 public:
   SessionManager(std::string &RuntimeResourcePath,
-                 std::string &DiagnosticDocumentationPath)
+                 std::string &DiagnosticDocumentationPath,
+                 std::shared_ptr<swift::PluginRegistry> Plugins)
       : RuntimeResourcePath(RuntimeResourcePath),
-        DiagnosticDocumentationPath(DiagnosticDocumentationPath) {}
+        DiagnosticDocumentationPath(DiagnosticDocumentationPath),
+        Plugins(Plugins) {}
 
   std::shared_ptr<Session> getSession(StringRef name);
 
@@ -356,7 +360,7 @@ class SwiftLangSupport : public LangSupport {
   std::shared_ptr<SwiftStatistics> Stats;
   llvm::StringMap<std::unique_ptr<FileSystemProvider>> FileSystemProviders;
   std::shared_ptr<swift::ide::IDEInspectionInstance> IDEInspectionInst;
-  compile::SessionManager CompileManager;
+  std::shared_ptr<compile::SessionManager> CompileManager;
 
 public:
   explicit SwiftLangSupport(SourceKit::Context &SKCtx);

tools/swift-ide-test/swift-ide-test.cpp

@@ -23,6 +23,7 @@
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/NameLookupRequests.h"
+#include "swift/AST/PluginRegistry.h"
 #include "swift/AST/PrintOptions.h"
 #include "swift/AST/RawComment.h"
 #include "swift/AST/USRGeneration.h"
@@ -1188,7 +1189,8 @@ static int doTypeContextInfo(const CompilerInvocation &InitInvok,
       InitInvok, SourceFilename, SecondSourceFileName, CodeCompletionToken,
       CodeCompletionDiagnostics,
       [&](CompletionLikeOperationParams Params) -> bool {
-        IDEInspectionInstance IDEInspectionInst;
+        IDEInspectionInstance IDEInspectionInst(
+            std::make_shared<PluginRegistry>());
         int ExitCode = 2;
         IDEInspectionInst.typeContextInfo(
             Params.Invocation, Params.Args, Params.FileSystem,
@@ -1260,7 +1262,8 @@ doConformingMethodList(const CompilerInvocation &InitInvok,
       InitInvok, SourceFilename, SecondSourceFileName, CodeCompletionToken,
       CodeCompletionDiagnostics,
       [&](CompletionLikeOperationParams Params) -> bool {
-        IDEInspectionInstance IDEInspectionInst;
+        IDEInspectionInstance IDEInspectionInst(
+            std::make_shared<PluginRegistry>());
         int ExitCode = 2;
         IDEInspectionInst.conformingMethodList(
             Params.Invocation, Params.Args, Params.FileSystem,
@@ -1410,7 +1413,7 @@ doCodeCompletion(const CompilerInvocation &InitInvok, StringRef SourceFilename,
       InitInvok, SourceFilename, SecondSourceFileName, CodeCompletionToken,
       CodeCompletionDiagnostics,
       [&](CompletionLikeOperationParams Params) -> bool {
-        IDEInspectionInstance Inst;
+        IDEInspectionInstance Inst(std::make_shared<PluginRegistry>());
         int ExitCode = 2;
         Inst.codeComplete(
             Params.Invocation, Params.Args, Params.FileSystem,
@@ -1504,7 +1507,7 @@ static int doBatchCodeCompletion(const CompilerInvocation &InitInvok,
   CompilerInvocation Invocation(InitInvok);
   auto FileSystem = llvm::vfs::getRealFileSystem();
 
-  IDEInspectionInstance IDEInspectionInst;
+  IDEInspectionInstance IDEInspectionInst(std::make_shared<PluginRegistry>());
 
   std::unique_ptr<ide::OnDiskCodeCompletionCache> OnDiskCache;
   if (!options::CompletionCachePath.empty())