[Sema] [NFC] Don't accumulate lazily typechecked functions in lazy immediate mode

Author: Zak Kent

include/swift/AST/SourceFile.h

@@ -170,6 +170,14 @@ class SourceFile final : public FileUnit {
   /// located in the source file.
   llvm::SetVector<ValueDecl *> DeclsWithRuntimeDiscoverableAttrs;
 
+  /// The list of functions defined in this file whose bodies have yet to be
+  /// typechecked. They must be held in this list instead of eagerly validated
+  /// because their bodies may force us to perform semantic checks of arbitrary
+  /// complexity, and we currently cannot handle those checks in isolation. E.g.
+  /// we cannot, in general, perform witness matching on singular requirements
+  /// unless the entire conformance has been evaluated.
+  std::vector<AbstractFunctionDecl *> DelayedFunctions;
+
   /// The list of top-level items in the source file. This is \c None if
   /// they have not yet been parsed.
   /// FIXME: Once addTopLevelDecl/prependTopLevelDecl
@@ -271,13 +279,11 @@ class SourceFile final : public FileUnit {
   /// The list of local type declarations in the source file.
   llvm::SetVector<TypeDecl *> LocalTypeDecls;
 
-  /// The list of functions defined in this file whose bodies have yet to be
-  /// typechecked. They must be held in this list instead of eagerly validated
-  /// because their bodies may force us to perform semantic checks of arbitrary
-  /// complexity, and we currently cannot handle those checks in isolation. E.g.
-  /// we cannot, in general, perform witness matching on singular requirements
-  /// unless the entire conformance has been evaluated.
-  std::vector<AbstractFunctionDecl *> DelayedFunctions;
+  /// Defer type checking of `AFD` to the end of `Sema`
+  void addDelayedFunction(AbstractFunctionDecl *AFD);
+
+  /// Typecheck the bodies of all lazily checked functions
+  void typeCheckDelayedFunctions();
 
   /// A mapping from Objective-C selectors to the methods that have
   /// those selectors.

lib/AST/Module.cpp

@@ -3802,6 +3802,27 @@ void SourceFile::prependTopLevelDecl(Decl *d) {
   ctx.evaluator.clearCachedOutput(ParseTopLevelDeclsRequest{mutableThis});
 }
 
+void SourceFile::addDelayedFunction(AbstractFunctionDecl *AFD) {
+  // If we defer type checking to runtime, we won't
+  // have to type check `AFD` ahead of time
+  auto &Ctx = getASTContext();
+  if (Ctx.TypeCheckerOpts.DeferToRuntime &&
+      Ctx.LangOpts.hasFeature(Feature::LazyImmediate))
+    return;
+  DelayedFunctions.push_back(AFD);
+}
+
+void SourceFile::typeCheckDelayedFunctions() {
+
+  for (unsigned i = 0; i < DelayedFunctions.size(); i++) {
+    auto *AFD = DelayedFunctions[i];
+    assert(!AFD->getDeclContext()->isLocalContext());
+    AFD->getTypecheckedBody();
+  }
+
+  DelayedFunctions.clear();
+}
+
 ArrayRef<ASTNode> SourceFile::getTopLevelItems() const {
   auto &ctx = getASTContext();
   auto *mutableThis = const_cast<SourceFile *>(this);

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -1844,12 +1844,13 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
 
   ASTContext &getASTContext() const { return Ctx; }
   void addDelayedFunction(AbstractFunctionDecl *AFD) {
-    if (!SF) return;
+    if (!SF)
+      return;
 
-    while (auto enclosingSF = SF->getEnclosingSourceFile())
-      SF = enclosingSF;
+    while (auto *ESF = SF->getEnclosingSourceFile())
+      SF = ESF;
 
-    SF->DelayedFunctions.push_back(AFD);
+    SF->addDelayedFunction(AFD);
   }
 
   void visit(Decl *decl) {

lib/Sema/TypeCheckDistributed.cpp

@@ -717,7 +717,7 @@ void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal
         continue;
 
       if (auto thunk = var->getDistributedThunk())
-        SF->DelayedFunctions.push_back(thunk);
+        SF->addDelayedFunction(thunk);
 
       continue;
     }
@@ -780,7 +780,7 @@ void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal
       }
 
       if (auto thunk = func->getDistributedThunk()) {
-        SF->DelayedFunctions.push_back(thunk);
+        SF->addDelayedFunction(thunk);
       }
     }
   }

lib/Sema/TypeChecker.cpp

@@ -241,19 +241,6 @@ void swift::bindExtensions(ModuleDecl &mod) {
   // typeCheckDecl().
 }
 
-static void typeCheckDelayedFunctions(SourceFile &SF) {
-  unsigned currentFunctionIdx = 0;
-
-  while (currentFunctionIdx < SF.DelayedFunctions.size()) {
-    auto *AFD = SF.DelayedFunctions[currentFunctionIdx];
-    assert(!AFD->getDeclContext()->isLocalContext());
-    (void) AFD->getTypecheckedBody();
-    ++currentFunctionIdx;
-  }
-
-  SF.DelayedFunctions.clear();
-}
-
 void swift::performTypeChecking(SourceFile &SF) {
   return (void)evaluateOrDefault(SF.getASTContext().evaluator,
                                  TypeCheckSourceFileRequest{&SF}, {});
@@ -306,10 +293,7 @@ TypeCheckSourceFileRequest::evaluate(Evaluator &eval, SourceFile *SF) const {
         }
       }
     }
-    if (!Ctx.TypeCheckerOpts.DeferToRuntime ||
-        !Ctx.LangOpts.hasFeature(Feature::LazyImmediate)) {
-      typeCheckDelayedFunctions(*SF);
-    }
+    SF->typeCheckDelayedFunctions();
   }
 
   diagnoseUnnecessaryPreconcurrencyImports(*SF);