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

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,