[CodeCompleiton] Avoid walking the whole module to find the delayed func

Previously delayed parsing was performed by AST walker
'ParseDelayedFunctionBodies' by finding "delayed" function decl from the AST
of the whole module. This is not necessary. Optimize it by remembering the
"delayed" function decl in 'PersistentParserState'.

NOTE: 'SourceLoader' stopped using this delayed parsing mechanism in
d8b745d

rdar://problem/56819166

Author: rintaro

include/swift/Parse/Parser.h

@@ -167,14 +167,14 @@ class Parser {
   LocalContext *CurLocalContext = nullptr;
 
   bool isDelayedParsingEnabled() const {
-    return DelayBodyParsing || SourceMgr.getCodeCompletionLoc().isValid();
+    return DelayBodyParsing || isCodeCompletionFirstPass();
   }
 
   void setCodeCompletionCallbacks(CodeCompletionCallbacks *Callbacks) {
     CodeCompletion = Callbacks;
   }
 
-  bool isCodeCompletionFirstPass() {
+  bool isCodeCompletionFirstPass() const {
     return L->isCodeCompletion() && !CodeCompletion;
   }
 
@@ -1057,6 +1057,7 @@ class Parser {
                                        bool HasFuncKeyword = true);
   void parseAbstractFunctionBody(AbstractFunctionDecl *AFD);
   BraceStmt *parseAbstractFunctionBodyDelayed(AbstractFunctionDecl *AFD);
+  void parseAbstractFunctionBodyDelayed();
   ParserResult<ProtocolDecl> parseDeclProtocol(ParseDeclOptions Flags,
                                                DeclAttributes &Attributes);
 

include/swift/Parse/PersistentParserState.h

@@ -56,6 +56,7 @@ class PersistentParserState {
   enum class DelayedDeclKind {
     TopLevelCodeDecl,
     Decl,
+    FunctionBody,
   };
 
   class DelayedDeclState {

include/swift/Subsystems.h

@@ -132,11 +132,9 @@ namespace swift {
                                PersistentParserState *PersistentState = nullptr,
                                bool DelayBodyParsing = true);
 
-  /// Finish the parsing by going over the nodes that were delayed
-  /// during the first parsing pass.
-  void performDelayedParsing(DeclContext *DC,
-                             PersistentParserState &PersistentState,
-                             CodeCompletionCallbacksFactory *Factory);
+  /// Finish the code completion.
+  void performCodeCompletionSecondPass(PersistentParserState &PersistentState,
+                                       CodeCompletionCallbacksFactory &Factory);
 
   /// Lex and return a vector of tokens for the given buffer.
   std::vector<Token> tokenize(const LangOptions &LangOpts,

lib/Frontend/Frontend.cpp

@@ -869,8 +869,8 @@ void CompilerInstance::parseAndCheckTypesUpTo(
   });
 
   if (Invocation.isCodeCompletion()) {
-    performDelayedParsing(MainModule, *PersistentState.get(),
-                          Invocation.getCodeCompletionFactory());
+    performCodeCompletionSecondPass(*PersistentState.get(),
+                                    *Invocation.getCodeCompletionFactory());
   }
   finishTypeChecking(TypeCheckOptions);
 }

lib/IDE/REPLCodeCompletion.cpp

@@ -215,7 +215,7 @@ doCodeCompletion(SourceFile &SF, StringRef EnteredCode, unsigned *BufferID,
   performTypeChecking(SF, PersistentState.getTopLevelContext(), None,
                       OriginalDeclCount);
 
-  performDelayedParsing(&SF, PersistentState, CompletionCallbacksFactory);
+  performCodeCompletionSecondPass(PersistentState, *CompletionCallbacksFactory);
 
   // Now we are done with code completion.  Remove the declarations we
   // temporarily inserted.

lib/Parse/ParseDecl.cpp

@@ -5340,11 +5340,16 @@ void Parser::consumeAbstractFunctionBody(AbstractFunctionDecl *AFD,
 
   BodyRange.End = PreviousLoc;
 
-  if (SourceMgr.getCodeCompletionLoc().isInvalid() ||
-      SourceMgr.rangeContainsCodeCompletionLoc(BodyRange)) {
-    AFD->setBodyDelayed(BodyRange);
-  } else {
-    AFD->setBodySkipped(BodyRange);
+  AFD->setBodyDelayed(BodyRange);
+
+  if (isCodeCompletionFirstPass()) {
+    if (SourceMgr.rangeContainsCodeCompletionLoc(BodyRange)) {
+      State->delayDecl(PersistentParserState::DelayedDeclKind::FunctionBody,
+                       PD_Default, AFD, BodyRange,
+                       BeginParserPosition.PreviousLoc);
+    } else {
+      AFD->setBodySkipped(BodyRange);
+    }
   }
 }
 
@@ -5656,6 +5661,20 @@ BraceStmt *Parser::parseAbstractFunctionBodyDelayed(AbstractFunctionDecl *AFD) {
   return parseBraceItemList(diag::func_decl_without_brace).getPtrOrNull();
 }
 
+/// Parse a delayed function body from the 'PersistentParserState'.
+void Parser::parseAbstractFunctionBodyDelayed() {
+  auto DelayedState = State->takeDelayedDeclState();
+  assert(DelayedState.get() && "should have delayed state");
+  auto CD = DelayedState->ParentContext->getAsDecl();
+  auto AFD = cast<AbstractFunctionDecl>(CD);
+
+  // Eagarly parse local decls or nested function bodies inside the body.
+  llvm::SaveAndRestore<bool> DisableDelayedBody(DelayBodyParsing, false);
+
+  auto body = parseAbstractFunctionBodyDelayed(AFD);
+  AFD->setBodyParsed(body);
+}
+
 /// Parse a 'enum' declaration, returning true (and doing no token
 /// skipping) on error.
 ///

lib/Parse/Parser.cpp

@@ -110,75 +110,24 @@ using namespace swift::syntax;
 
 void SILParserTUStateBase::anchor() { }
 
-namespace {
-/// A visitor that does delayed parsing of function bodies.
-class ParseDelayedFunctionBodies : public ASTWalker {
-  PersistentParserState &ParserState;
-  CodeCompletionCallbacksFactory *CodeCompletionFactory;
-
-public:
-  ParseDelayedFunctionBodies(PersistentParserState &ParserState,
-                             CodeCompletionCallbacksFactory *Factory)
-    : ParserState(ParserState), CodeCompletionFactory(Factory) {}
-
-  bool walkToDeclPre(Decl *D) override {
-    if (auto AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      if (AFD->getBodyKind() != FuncDecl::BodyKind::Unparsed)
-        return false;
-      parseFunctionBody(AFD);
-      return true;
-    }
-    return true;
-  }
-
-private:
-  void parseFunctionBody(AbstractFunctionDecl *AFD) {
-    // FIXME: This duplicates the evaluation of
-    // ParseAbstractFunctionBodyRequest, but installs a code completion
-    // factory.
-    assert(AFD->getBodyKind() == FuncDecl::BodyKind::Unparsed);
-
-    SourceFile &SF = *AFD->getDeclContext()->getParentSourceFile();
-    SourceManager &SourceMgr = SF.getASTContext().SourceMgr;
-    unsigned BufferID = SourceMgr.findBufferContainingLoc(AFD->getLoc());
-    Parser TheParser(BufferID, SF, nullptr, &ParserState, nullptr,
-                     /*DelayBodyParsing=*/false);
-    TheParser.SyntaxContext->disable();
-    std::unique_ptr<CodeCompletionCallbacks> CodeCompletion;
-    if (CodeCompletionFactory) {
-      CodeCompletion.reset(
-          CodeCompletionFactory->createCodeCompletionCallbacks(TheParser));
-      TheParser.setCodeCompletionCallbacks(CodeCompletion.get());
-    }
-    auto body = TheParser.parseAbstractFunctionBodyDelayed(AFD);
-    AFD->setBodyParsed(body);
-
-    if (CodeCompletion)
-      CodeCompletion->doneParsing();
-  }
-};
-
-static void parseDelayedDecl(
-              PersistentParserState &ParserState,
-              CodeCompletionCallbacksFactory *CodeCompletionFactory) {
+void swift::performCodeCompletionSecondPass(
+    PersistentParserState &ParserState,
+    CodeCompletionCallbacksFactory &Factory) {
+  SharedTimer timer("CodeCompletionSecondPass");
   if (!ParserState.hasDelayedDecl())
     return;
 
-  SourceFile &SF = *ParserState.getDelayedDeclContext()->getParentSourceFile();
-  SourceManager &SourceMgr = SF.getASTContext().SourceMgr;
-  unsigned BufferID =
-    SourceMgr.findBufferContainingLoc(ParserState.getDelayedDeclLoc());
+  auto &SF = *ParserState.getDelayedDeclContext()->getParentSourceFile();
+  auto &SM = SF.getASTContext().SourceMgr;
+  auto BufferID = SM.findBufferContainingLoc(ParserState.getDelayedDeclLoc());
   Parser TheParser(BufferID, SF, nullptr, &ParserState, nullptr);
 
   // Disable libSyntax creation in the delayed parsing.
   TheParser.SyntaxContext->disable();
 
-  std::unique_ptr<CodeCompletionCallbacks> CodeCompletion;
-  if (CodeCompletionFactory) {
-    CodeCompletion.reset(
-      CodeCompletionFactory->createCodeCompletionCallbacks(TheParser));
-    TheParser.setCodeCompletionCallbacks(CodeCompletion.get());
-  }
+  std::unique_ptr<CodeCompletionCallbacks> CodeCompletion(
+      Factory.createCodeCompletionCallbacks(TheParser));
+  TheParser.setCodeCompletionCallbacks(CodeCompletion.get());
 
   switch (ParserState.getDelayedDeclKind()) {
   case PersistentParserState::DelayedDeclKind::TopLevelCodeDecl:
@@ -188,13 +137,16 @@ static void parseDelayedDecl(
   case PersistentParserState::DelayedDeclKind::Decl:
     TheParser.parseDeclDelayed();
     break;
+
+  case PersistentParserState::DelayedDeclKind::FunctionBody: {
+    TheParser.parseAbstractFunctionBodyDelayed();
+    break;
   }
+  }
+  assert(!ParserState.hasDelayedDecl());
 
-  if (CodeCompletion)
-    CodeCompletion->doneParsing();
+  CodeCompletion->doneParsing();
 }
-} // unnamed namespace
-
 
 swift::Parser::BacktrackingScope::~BacktrackingScope() {
   if (Backtrack) {
@@ -203,18 +155,6 @@ swift::Parser::BacktrackingScope::~BacktrackingScope() {
   }
 }
 
-void swift::performDelayedParsing(
-    DeclContext *DC, PersistentParserState &PersistentState,
-    CodeCompletionCallbacksFactory *CodeCompletionFactory) {
-  SharedTimer timer("Parsing");
-  ParseDelayedFunctionBodies Walker(PersistentState,
-                                    CodeCompletionFactory);
-  DC->walkContext(Walker);
-
-  if (CodeCompletionFactory)
-    parseDelayedDecl(PersistentState, CodeCompletionFactory);
-}
-
 /// Tokenizes a string literal, taking into account string interpolation.
 static void getStringPartTokens(const Token &Tok, const LangOptions &LangOpts,
                                 const SourceManager &SM,