[sourcekit] Use a shared_ptr for the SwiftASTManager

When the server shuts down we may still have outstanding async work to
build an AST, so use a shared_ptr + weak_ptr instead of unique_ptr +
unowned references.

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -283,7 +283,7 @@ public:
     return AST;
   }
 
-  void getASTUnitAsync(SwiftASTManager::Implementation &MgrImpl,
+  void getASTUnitAsync(std::shared_ptr<SwiftASTManager> Mgr,
                        ArrayRef<ImmutableTextSnapshotRef> Snapshots,
                 std::function<void(ASTUnitRef Unit, StringRef Error)> Receiver);
   bool shouldRebuild(SwiftASTManager::Implementation &MgrImpl,
@@ -592,7 +592,7 @@ void SwiftASTManager::processASTAsync(SwiftInvocationRef InvokRef,
     }
   };
 
-  Producer->getASTUnitAsync(Impl, Snapshots, std::move(handleAST));
+  Producer->getASTUnitAsync(shared_from_this(), Snapshots, std::move(handleAST));
 }
 
 void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
@@ -660,7 +660,7 @@ SwiftASTManager::Implementation::getMemoryBuffer(StringRef Filename,
   return nullptr;
 }
 
-void ASTProducer::getASTUnitAsync(SwiftASTManager::Implementation &MgrImpl,
+void ASTProducer::getASTUnitAsync(std::shared_ptr<SwiftASTManager> Mgr,
                                   ArrayRef<ImmutableTextSnapshotRef> Snaps,
                std::function<void(ASTUnitRef Unit, StringRef Error)> Receiver) {
 
@@ -668,9 +668,9 @@ void ASTProducer::getASTUnitAsync(SwiftASTManager::Implementation &MgrImpl,
   SmallVector<ImmutableTextSnapshotRef, 4> Snapshots;
   Snapshots.append(Snaps.begin(), Snaps.end());
 
-  MgrImpl.ASTBuildQueue.dispatch([ThisProducer, &MgrImpl, Snapshots, Receiver] {
+  Mgr->Impl.ASTBuildQueue.dispatch([ThisProducer, Mgr, Snapshots, Receiver] {
     std::string Error;
-    ASTUnitRef Unit = ThisProducer->getASTUnitImpl(MgrImpl, Snapshots, Error);
+    ASTUnitRef Unit = ThisProducer->getASTUnitImpl(Mgr->Impl, Snapshots, Error);
     Receiver(Unit, Error);
   }, /*isStackDeep=*/true);
 }

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -84,7 +84,7 @@ public:
 
 typedef std::shared_ptr<SwiftASTConsumer> SwiftASTConsumerRef;
 
-class SwiftASTManager {
+class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
 public:
   explicit SwiftASTManager(SwiftLangSupport &LangSupport);
   ~SwiftASTManager();
@@ -131,8 +131,6 @@ public:
   void removeCachedAST(SwiftInvocationRef Invok);
 
   struct Implementation;
-
-private:
   Implementation &Impl;
 };
 

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -595,7 +595,7 @@ class SwiftDocumentSemanticInfo :
     public ThreadSafeRefCountedBase<SwiftDocumentSemanticInfo> {
 
   const std::string Filename;
-  SwiftASTManager &ASTMgr;
+  std::weak_ptr<SwiftASTManager> ASTMgr;
   std::weak_ptr<NotificationCenter> NotificationCtr;
   ThreadSafeRefCntPtr<SwiftInvocation> InvokRef;
   std::string CompilerArgsError;
@@ -610,10 +610,10 @@ class SwiftDocumentSemanticInfo :
   mutable llvm::sys::Mutex Mtx;
 
 public:
-  SwiftDocumentSemanticInfo(StringRef Filename, SwiftLangSupport &LangSupport)
+  SwiftDocumentSemanticInfo(StringRef Filename,
-    : Filename(Filename),
+                            std::shared_ptr<SwiftASTManager> ASTMgr,
-      ASTMgr(LangSupport.getASTManager()),
+                            std::shared_ptr<NotificationCenter> NotificationCtr)
-      NotificationCtr(LangSupport.getContext().getNotificationCenter()) {}
+      : Filename(Filename), ASTMgr(ASTMgr), NotificationCtr(NotificationCtr) {}
 
   SwiftInvocationRef getInvocation() const {
     return InvokRef;
@@ -622,7 +622,9 @@ public:
   uint64_t getASTGeneration() const;
 
   void setCompilerArgs(ArrayRef<const char *> Args) {
-    InvokRef = ASTMgr.getInvocation(Args, Filename, CompilerArgsError);
+    if (auto ASTMgr = this->ASTMgr.lock()) {
+      InvokRef = ASTMgr->getInvocation(Args, Filename, CompilerArgsError);
+    }
   }
 
   void readSemanticInfo(ImmutableTextSnapshotRef NewSnapshot,
@@ -637,8 +639,11 @@ public:
                           ImmutableTextSnapshotRef Snapshot,
                           uint64_t ASTGeneration);
   void removeCachedAST() {
-    if (InvokRef)
+    if (InvokRef) {
-      ASTMgr.removeCachedAST(InvokRef);
+      if (auto ASTMgr = this->ASTMgr.lock()) {
+        ASTMgr->removeCachedAST(InvokRef);
+      }
+    }
   }
 
 private:
@@ -1018,7 +1023,9 @@ void SwiftDocumentSemanticInfo::processLatestSnapshotAsync(
   // previously queued queries for the same document. Each document has a
   // SwiftDocumentSemanticInfo pointer so use that for the token.
   const void *OncePerASTToken = SemaInfoRef.get();
-  ASTMgr.processASTAsync(Invok, std::move(Consumer), OncePerASTToken);
+  if (auto ASTMgr = this->ASTMgr.lock()) {
+    ASTMgr->processASTAsync(Invok, std::move(Consumer), OncePerASTToken);
+  }
 }
 
 struct SwiftEditorDocument::Implementation {
@@ -1050,8 +1057,10 @@ struct SwiftEditorDocument::Implementation {
 
   Implementation(StringRef FilePath, SwiftLangSupport &LangSupport,
                  CodeFormatOptions options)
-    : LangSupport(LangSupport), FilePath(FilePath), FormatOptions(options) {
+      : LangSupport(LangSupport), FilePath(FilePath), FormatOptions(options) {
-      SemanticInfo = new SwiftDocumentSemanticInfo(FilePath, LangSupport);
+    SemanticInfo = new SwiftDocumentSemanticInfo(
+        FilePath, LangSupport.getASTManager().shared_from_this(),
+        LangSupport.getContext().getNotificationCenter());
   }
 };
 
@@ -1626,8 +1635,9 @@ ImmutableTextSnapshotRef SwiftEditorDocument::initializeText(
   Impl.SyntaxMap.Tokens.clear();
   Impl.AffectedRange = {0, static_cast<unsigned>(Buf->getBufferSize())};
 
-  Impl.SemanticInfo =
+  Impl.SemanticInfo = new SwiftDocumentSemanticInfo(
-      new SwiftDocumentSemanticInfo(Impl.FilePath, Impl.LangSupport);
+      Impl.FilePath, Impl.LangSupport.getASTManager().shared_from_this(),
+      Impl.LangSupport.getContext().getNotificationCenter());
   Impl.SemanticInfo->setCompilerArgs(Args);
   return Impl.EditableBuffer->getSnapshot();
 }

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

@@ -178,7 +178,7 @@ SwiftLangSupport::SwiftLangSupport(SourceKit::Context &SKCtx)
   llvm::sys::path::append(LibPath, "swift");
   RuntimeResourcePath = LibPath.str();
 
-  ASTMgr.reset(new SwiftASTManager(*this));
+  ASTMgr = std::make_shared<SwiftASTManager>(*this);
   // By default, just use the in-memory cache.
   CCCache->inMemory = llvm::make_unique<ide::CodeCompletionCache>();
 }

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -272,7 +272,7 @@ struct SwiftStatistics {
 class SwiftLangSupport : public LangSupport {
   SourceKit::Context &SKCtx;
   std::string RuntimeResourcePath;
-  std::unique_ptr<SwiftASTManager> ASTMgr;
+  std::shared_ptr<SwiftASTManager> ASTMgr;
   SwiftEditorDocumentFileMap EditorDocuments;
   SwiftInterfaceGenMap IFaceGenContexts;
   ThreadSafeRefCntPtr<SwiftCompletionCache> CCCache;