Move PersistentParserState onto SourceFile

Move the global PersistentParserState from
the CompilerInstance to the source file that code
completion is operating on, only hooking up the
state when it's needed. This will help make it
easier to requestify source file parsing.

Author: hamishknight

include/swift/AST/SourceFile.h

@@ -18,6 +18,8 @@
 
 namespace swift {
 
+class PersistentParserState;
+
 /// A file containing Swift source code.
 ///
 /// This is a .swift or .sil file (or a virtual file, such as the contents of
@@ -154,6 +156,16 @@ class SourceFile final : public FileUnit {
   /// mechanism which is not SourceFile-dependent.)
   SeparatelyImportedOverlayMap separatelyImportedOverlays;
 
+  /// A pointer to PersistentParserState with a function reference to its
+  /// deleter to handle the fact that it's forward declared.
+  using ParserStatePtr =
+      std::unique_ptr<PersistentParserState, void (*)(PersistentParserState *)>;
+
+  /// Stores delayed parser state that code completion needs to be able to
+  /// resume parsing at the code completion token in the file.
+  ParserStatePtr DelayedParserState =
+      ParserStatePtr(/*ptr*/ nullptr, /*deleter*/ nullptr);
+
   friend ASTContext;
   friend Impl;
 
@@ -405,6 +417,20 @@ class SourceFile final : public FileUnit {
   /// Retrieve the scope that describes this source file.
   ASTScope &getScope();
 
+  /// Retrieves the previously set delayed parser state, asserting that it
+  /// exists.
+  PersistentParserState *getDelayedParserState() {
+    auto *state = DelayedParserState.get();
+    assert(state && "Didn't set any delayed parser state!");
+    return state;
+  }
+
+  /// Record delayed parser state for the source file. This is needed for code
+  /// completion's second pass.
+  void setDelayedParserState(ParserStatePtr &&state) {
+    DelayedParserState = std::move(state);
+  }
+
   SWIFT_DEBUG_DUMP;
   void dump(raw_ostream &os) const;
 

include/swift/Frontend/Frontend.h

@@ -400,8 +400,6 @@ class CompilerInstance {
   std::unique_ptr<Lowering::TypeConverter> TheSILTypes;
   std::unique_ptr<SILModule> TheSILModule;
 
-  std::unique_ptr<PersistentParserState> PersistentState;
-
   /// Null if no tracker.
   std::unique_ptr<DependencyTracker> DepTracker;
   /// If there is no stats output directory by the time the
@@ -431,6 +429,9 @@ class CompilerInstance {
   /// buffer will also have its buffer ID in PrimaryBufferIDs.
   std::vector<SourceFile *> PrimarySourceFiles;
 
+  /// The file that has been registered for code completion.
+  NullablePtr<SourceFile> CodeCompletionFile;
+
   /// Return whether there is an entry in PrimaryInputs for buffer \p BufID.
   bool isPrimaryInput(unsigned BufID) const {
     return PrimaryBufferIDs.count(BufID) != 0;
@@ -550,12 +551,13 @@ class CompilerInstance {
     return Invocation;
   }
 
-  bool hasPersistentParserState() const {
-    return bool(PersistentState);
-  }
+  /// If a code completion buffer has been set, returns the corresponding source
+  /// file.
+  NullablePtr<SourceFile> getCodeCompletionFile() { return CodeCompletionFile; }
 
-  PersistentParserState &getPersistentParserState() {
-    return *PersistentState.get();
+  /// Set a new file that we're performing code completion on.
+  void setCodeCompletionFile(SourceFile *file) {
+    CodeCompletionFile = file;
   }
 
 private:

include/swift/Parse/PersistentParserState.h

@@ -79,13 +79,17 @@ class PersistentParserState {
   void restoreCodeCompletionDelayedDeclState(
       const CodeCompletionDelayedDeclState &other);
 
-  bool hasCodeCompletionDelayedDeclState() {
+  bool hasCodeCompletionDelayedDeclState() const {
     return CodeCompletionDelayedDeclStat.get() != nullptr;
   }
 
   CodeCompletionDelayedDeclState &getCodeCompletionDelayedDeclState() {
     return *CodeCompletionDelayedDeclStat.get();
   }
+  const CodeCompletionDelayedDeclState &
+  getCodeCompletionDelayedDeclState() const {
+    return *CodeCompletionDelayedDeclStat.get();
+  }
 
   std::unique_ptr<CodeCompletionDelayedDeclState>
   takeCodeCompletionDelayedDeclState() {

include/swift/Subsystems.h

@@ -55,7 +55,6 @@ namespace swift {
   class ModuleDecl;
   typedef void *OpaqueSyntaxNode;
   class Parser;
-  class PersistentParserState;
   class SerializationOptions;
   class SILOptions;
   class SILModule;
@@ -110,21 +109,17 @@ namespace swift {
   ///
   /// \param BufferID The buffer to parse from.
   ///
-  /// \param PersistentState If non-null the same PersistentState object can be
-  /// used to save parser state for code completion.
-  ///
   /// \param DelayBodyParsing Whether parsing of type and function bodies can be
   /// delayed.
   void parseIntoSourceFile(SourceFile &SF, unsigned BufferID,
-                           PersistentParserState *PersistentState = nullptr,
                            bool DelayBodyParsing = true,
                            bool EvaluateConditionals = true);
 
   /// Parse a source file's SIL declarations into a given SIL module.
   void parseSourceFileSIL(SourceFile &SF, SILParserState *sil);
 
   /// Finish the code completion.
-  void performCodeCompletionSecondPass(PersistentParserState &PersistentState,
+  void performCodeCompletionSecondPass(SourceFile &SF,
                                        CodeCompletionCallbacksFactory &Factory);
 
   /// Lex and return a vector of tokens for the given buffer.

lib/Frontend/Frontend.cpp

@@ -906,8 +906,6 @@ void CompilerInstance::parseAndCheckTypesUpTo(
     const ImplicitImports &implicitImports, SourceFile::ASTStage_t limitStage) {
   FrontendStatsTracer tracer(getStatsReporter(), "parse-and-check-types");
 
-  PersistentState = std::make_unique<PersistentParserState>();
-
   bool hadLoadError = parsePartialModulesAndLibraryFiles(implicitImports);
   if (Invocation.isCodeCompletion()) {
     // When we are doing code completion, make sure to emit at least one
@@ -978,8 +976,7 @@ void CompilerInstance::parseLibraryFile(
   auto DidSuppressWarnings = Diags.getSuppressWarnings();
   Diags.setSuppressWarnings(DidSuppressWarnings || !IsPrimary);
 
-  parseIntoSourceFile(*NextInput, BufferID, PersistentState.get(),
-                      /*DelayedBodyParsing=*/!IsPrimary);
+  parseIntoSourceFile(*NextInput, BufferID, /*DelayedBodyParsing=*/!IsPrimary);
 
   Diags.setSuppressWarnings(DidSuppressWarnings);
 
@@ -1026,7 +1023,7 @@ void CompilerInstance::parseAndTypeCheckMainFileUpTo(
   Diags.setSuppressWarnings(DidSuppressWarnings || !mainIsPrimary);
 
   // Parse the Swift decls into the source file.
-  parseIntoSourceFile(MainFile, MainBufferID, PersistentState.get(),
+  parseIntoSourceFile(MainFile, MainBufferID,
                       /*delayBodyParsing*/ !mainIsPrimary);
 
   // For a primary, also perform type checking if needed. Otherwise, just do
@@ -1096,6 +1093,11 @@ SourceFile *CompilerInstance::createSourceFileForMainModule(
     recordPrimarySourceFile(inputFile);
   }
 
+  if (bufferID == SourceMgr.getCodeCompletionBufferID()) {
+    assert(!CodeCompletionFile && "Multiple code completion files?");
+    CodeCompletionFile = inputFile;
+  }
+
   return inputFile;
 }
 
@@ -1121,8 +1123,6 @@ void CompilerInstance::performParseOnly(bool EvaluateConditionals,
                                   MainBufferID);
   }
 
-  PersistentState = std::make_unique<PersistentParserState>();
-
   auto shouldDelayBodies = [&](unsigned bufferID) -> bool {
     if (!CanDelayBodies)
       return false;
@@ -1140,8 +1140,8 @@ void CompilerInstance::performParseOnly(bool EvaluateConditionals,
         SourceFileKind::Library, SourceFile::ImplicitModuleImportKind::None,
         BufferID);
 
-    parseIntoSourceFile(*NextInput, BufferID, PersistentState.get(),
-                        shouldDelayBodies(BufferID), EvaluateConditionals);
+    parseIntoSourceFile(*NextInput, BufferID, shouldDelayBodies(BufferID),
+                        EvaluateConditionals);
   }
 
   // Now parse the main file.
@@ -1151,16 +1151,15 @@ void CompilerInstance::performParseOnly(bool EvaluateConditionals,
     MainFile.SyntaxParsingCache = Invocation.getMainFileSyntaxParsingCache();
     assert(MainBufferID == MainFile.getBufferID());
 
-    parseIntoSourceFile(MainFile, MainBufferID, PersistentState.get(),
-                        shouldDelayBodies(MainBufferID), EvaluateConditionals);
+    parseIntoSourceFile(MainFile, MainBufferID, shouldDelayBodies(MainBufferID),
+                        EvaluateConditionals);
   }
 
   assert(Context->LoadedModules.size() == 1 &&
          "Loaded a module during parse-only");
 }
 
 void CompilerInstance::freeASTContext() {
-  PersistentState.reset();
   TheSILTypes.reset();
   Context.reset();
   MainModule = nullptr;

lib/IDE/CompletionInstance.cpp

@@ -176,21 +176,17 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
     return false;
 
   auto &CI = *CachedCI;
+  auto *oldSF = CI.getCodeCompletionFile().get();
 
-  if (!CI.hasPersistentParserState())
-    return false;
-  auto &oldState = CI.getPersistentParserState();
-  if (!oldState.hasCodeCompletionDelayedDeclState())
-    return false;
+  auto *oldState = oldSF->getDelayedParserState();
+  assert(oldState->hasCodeCompletionDelayedDeclState());
+  auto &oldInfo = oldState->getCodeCompletionDelayedDeclState();
 
   auto &SM = CI.getSourceMgr();
   if (SM.getIdentifierForBuffer(SM.getCodeCompletionBufferID()) !=
       completionBuffer->getBufferIdentifier())
     return false;
 
-  auto &oldInfo = oldState.getCodeCompletionDelayedDeclState();
-  auto *oldSF = oldInfo.ParentContext->getParentSourceFile();
-
   // Parse the new buffer into temporary SourceFile.
   SourceManager tmpSM;
   auto tmpBufferID = tmpSM.addMemBufferCopy(completionBuffer);
@@ -207,19 +203,20 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
   registerTypeCheckerRequestFunctions(tmpCtx->evaluator);
   registerSILGenRequestFunctions(tmpCtx->evaluator);
   ModuleDecl *tmpM = ModuleDecl::create(Identifier(), *tmpCtx);
-  PersistentParserState newState;
   SourceFile *tmpSF =
       new (*tmpCtx) SourceFile(*tmpM, oldSF->Kind, tmpBufferID,
                                SourceFile::ImplicitModuleImportKind::None);
   tmpSF->enableInterfaceHash();
   // Ensure all non-function-body tokens are hashed into the interface hash
   tmpCtx->LangOpts.EnableTypeFingerprints = false;
-  parseIntoSourceFile(*tmpSF, tmpBufferID, &newState);
+  parseIntoSourceFile(*tmpSF, tmpBufferID);
+
   // Couldn't find any completion token?
-  if (!newState.hasCodeCompletionDelayedDeclState())
+  auto *newState = tmpSF->getDelayedParserState();
+  if (!newState->hasCodeCompletionDelayedDeclState())
     return false;
 
-  auto &newInfo = newState.getCodeCompletionDelayedDeclState();
+  auto &newInfo = newState->getCodeCompletionDelayedDeclState();
   unsigned newBufferID;
 
   switch (newInfo.Kind) {
@@ -273,7 +270,7 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
 
     // Construct dummy scopes. We don't need to restore the original scope
     // because they are probably not 'isResolvable()' anyway.
-    auto &SI = oldState.getScopeInfo();
+    auto &SI = oldState->getScopeInfo();
     assert(SI.getCurrentScope() == nullptr);
     Scope Top(SI, ScopeKind::TopLevel);
     Scope Body(SI, ScopeKind::FunctionBody);
@@ -282,7 +279,7 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
     oldInfo.StartOffset = newInfo.StartOffset;
     oldInfo.EndOffset = newInfo.EndOffset;
     oldInfo.PrevOffset = newInfo.PrevOffset;
-    oldState.restoreCodeCompletionDelayedDeclState(oldInfo);
+    oldState->restoreCodeCompletionDelayedDeclState(oldInfo);
 
     auto *AFD = cast<AbstractFunctionDecl>(DC);
     if (AFD->isBodySkipped())
@@ -324,14 +321,22 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
     CompilerInstance::addAdditionalInitialImportsTo(newSF, implicitImport);
     newSF->enableInterfaceHash();
 
+    // Tell the compiler instance we've replaced the code completion file.
+    CI.setCodeCompletionFile(newSF);
+
     // Re-parse the whole file. Still re-use imported modules.
-    (void)oldState.takeCodeCompletionDelayedDeclState();
-    parseIntoSourceFile(*newSF, newBufferID, &oldState);
+    parseIntoSourceFile(*newSF, newBufferID);
     performNameBinding(*newSF);
     bindExtensions(*newSF);
 
-    assert(oldState.hasCodeCompletionDelayedDeclState() &&
-           oldState.getCodeCompletionDelayedDeclState().Kind == newInfo.Kind);
+#ifndef NDEBUG
+    const auto *reparsedState = newSF->getDelayedParserState();
+    assert(reparsedState->hasCodeCompletionDelayedDeclState() &&
+           "Didn't find completion token?");
+
+    auto &reparsedInfo = reparsedState->getCodeCompletionDelayedDeclState();
+    assert(reparsedInfo.Kind == newInfo.Kind);
+#endif
     break;
   }
   }
@@ -360,27 +365,42 @@ bool CompletionInstance::performNewOperation(
     llvm::function_ref<void(CompilerInstance &)> Callback) {
 
   auto TheInstance = std::make_unique<CompilerInstance>();
-  auto &CI = *TheInstance;
-  if (DiagC)
-    CI.addDiagnosticConsumer(DiagC);
+  {
+    auto &CI = *TheInstance;
+    if (DiagC)
+      CI.addDiagnosticConsumer(DiagC);
 
-  if (FileSystem != llvm::vfs::getRealFileSystem())
-    CI.getSourceMgr().setFileSystem(FileSystem);
+    SWIFT_DEFER {
+      if (DiagC)
+        CI.removeDiagnosticConsumer(DiagC);
+    };
 
-  Invocation.setCodeCompletionPoint(completionBuffer, Offset);
+    if (FileSystem != llvm::vfs::getRealFileSystem())
+      CI.getSourceMgr().setFileSystem(FileSystem);
 
-  if (CI.setup(Invocation)) {
-    Error = "failed to setup compiler instance";
-    return false;
-  }
-  registerIDERequestFunctions(CI.getASTContext().evaluator);
+    Invocation.setCodeCompletionPoint(completionBuffer, Offset);
 
-  CI.performParseAndResolveImportsOnly();
-  if (CI.hasPersistentParserState())
-    Callback(CI);
+    if (CI.setup(Invocation)) {
+      Error = "failed to setup compiler instance";
+      return false;
+    }
+    registerIDERequestFunctions(CI.getASTContext().evaluator);
 
-  if (DiagC)
-    CI.removeDiagnosticConsumer(DiagC);
+    CI.performParseAndResolveImportsOnly();
+
+    // If we didn't create a source file for completion, bail. This can happen
+    // if for example we fail to load the stdlib.
+    auto completionFile = CI.getCodeCompletionFile();
+    if (!completionFile)
+      return true;
+
+    // If we didn't find a code completion token, bail.
+    auto *state = completionFile.get()->getDelayedParserState();
+    if (!state->hasCodeCompletionDelayedDeclState())
+      return true;
+
+    Callback(CI);
+  }
 
   if (ArgsHash.hasValue()) {
     CachedCI = std::move(TheInstance);

lib/IDE/REPLCodeCompletion.cpp

@@ -232,11 +232,10 @@ doCodeCompletion(SourceFile &SF, StringRef EnteredCode, unsigned *BufferID,
     newSF.addImports(importsWithOptions);
   }
 
-  PersistentParserState PersistentState;
-  parseIntoSourceFile(newSF, *BufferID, &PersistentState);
+  parseIntoSourceFile(newSF, *BufferID);
   performTypeChecking(newSF);
 
-  performCodeCompletionSecondPass(PersistentState, *CompletionCallbacksFactory);
+  performCodeCompletionSecondPass(newSF, *CompletionCallbacksFactory);
 
   // Reset the error state because it's only relevant to the code that we just
   // processed, which now gets thrown away.