[SourceKit] Use a single PluginRegistry in multiple ASTContexts

Make a single 'PluginRegistry' and share it between SwiftASTManager,
IDEInspectionInstance, and CompileInstance. And inject the plugin
registry to ASTContext right after 'CompilerInstance.setup()'

That way, all sema-capable ASTContext in SourceKit share a single
PluginRegistry.

(cherry picked from commit a4d1fc6)

Author: rintaro

include/swift/AST/PluginRegistry.h

@@ -139,6 +139,8 @@ class PluginRegistry {
   /// Flag to dump plugin messagings.
   bool dumpMessaging = false;
 
+  std::mutex mtx;
+
 public:
   PluginRegistry();
 

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