Sema: Convert TypeChecker::computeCaptures() into two requests

We now compute captures of functions and default arguments
lazily, instead of as a side effect of primary file checking.

Captures of closures are computed as part of the enclosing
context, not lazily, because the type checking of a single
closure body is not lazy.

This fixes a specific issue with the `-experimental-skip-*` flags,
where functions declared after a top-level `guard` statement are
considered to have local captures, but nothing was forcing these
captures to be computed.

Fixes rdar://problem/125981663.

Author: slavapestov

include/swift/AST/Decl.h

@@ -6814,9 +6814,14 @@ class ParamDecl : public VarDecl {
 
   void setDefaultArgumentInitContext(Initializer *initContext);
 
-  CaptureInfo getDefaultArgumentCaptureInfo() const {
+  CaptureInfo getDefaultArgumentCaptureInfo() const;
+
+  std::optional<CaptureInfo> getCachedDefaultArgumentCaptureInfo() const {
     assert(DefaultValueAndFlags.getPointer());
-    return DefaultValueAndFlags.getPointer()->Captures;
+    const auto &captures = DefaultValueAndFlags.getPointer()->Captures;
+    if (!captures.hasBeenComputed())
+      return std::nullopt;
+    return captures;
   }
 
   void setDefaultArgumentCaptureInfo(CaptureInfo captures);
@@ -7754,7 +7759,13 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Retrieve the source range of the function declaration name + patterns.
   SourceRange getSignatureSourceRange() const;
 
-  CaptureInfo getCaptureInfo() const { return Captures; }
+  CaptureInfo getCaptureInfo() const;
+
+  std::optional<CaptureInfo> getCachedCaptureInfo() const {
+    if (!Captures.hasBeenComputed())
+      return std::nullopt;
+    return Captures;
+  }
 
   void setCaptureInfo(CaptureInfo captures) {
     assert(captures.hasBeenComputed());

include/swift/AST/Expr.h

@@ -3912,7 +3912,16 @@ class AbstractClosureExpr : public DeclContext, public Expr {
           std::tuple<CapturedValue, unsigned, ApplyIsolationCrossing>>
           &foundIsolationCrossings);
 
-  CaptureInfo getCaptureInfo() const { return Captures; }
+  CaptureInfo getCaptureInfo() const {
+    assert(Captures.hasBeenComputed());
+    return Captures;
+  }
+
+  std::optional<CaptureInfo> getCachedCaptureInfo() const {
+    if (!Captures.hasBeenComputed())
+      return std::nullopt;
+    return Captures;
+  }
 
   void setCaptureInfo(CaptureInfo captures) {
     assert(captures.hasBeenComputed());

include/swift/AST/TypeCheckRequests.h

@@ -4949,6 +4949,47 @@ class LifetimeDependenceInfoRequest
 public:
   // Caching.
   bool isCached() const { return true; }
+
+};
+
+class CaptureInfoRequest :
+    public SimpleRequest<CaptureInfoRequest,
+                         CaptureInfo(AbstractFunctionDecl *),
+                         RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  CaptureInfo evaluate(Evaluator &evaluator, AbstractFunctionDecl *func) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  std::optional<CaptureInfo> getCachedResult() const;
+  void cacheResult(CaptureInfo value) const;
+};
+
+class ParamCaptureInfoRequest :
+    public SimpleRequest<ParamCaptureInfoRequest,
+                         CaptureInfo(ParamDecl *),
+                         RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  CaptureInfo evaluate(Evaluator &evaluator, ParamDecl *param) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  std::optional<CaptureInfo> getCachedResult() const;
+  void cacheResult(CaptureInfo value) const;
 };
 
 class SuppressesConformanceRequest

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -571,3 +571,10 @@ SWIFT_REQUEST(TypeChecker, LifetimeDependenceInfoRequest,
 SWIFT_REQUEST(TypeChecker, SuppressesConformanceRequest,
               bool(NominalTypeDecl *decl, KnownProtocolKind kp),
               SeparatelyCached, NoLocationInfo)
+
+SWIFT_REQUEST(TypeChecker, CaptureInfoRequest,
+              CaptureInfo(AbstractFunctionDecl *),
+              SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ParamCaptureInfoRequest,
+              CaptureInfo(ParamDecl *),
+              SeparatelyCached, NoLocationInfo)

lib/AST/ASTDumper.cpp

@@ -1402,10 +1402,10 @@ namespace {
         printField(PD->getDefaultArgumentKind(), "default_arg");
       }
       if (PD->hasDefaultExpr() &&
-          PD->getDefaultArgumentCaptureInfo().hasBeenComputed() &&
-          !PD->getDefaultArgumentCaptureInfo().isTrivial()) {
+          PD->getCachedDefaultArgumentCaptureInfo() &&
+          !PD->getCachedDefaultArgumentCaptureInfo()->isTrivial()) {
         printFieldRaw([&](raw_ostream &OS) {
-          PD->getDefaultArgumentCaptureInfo().print(OS);
+          PD->getCachedDefaultArgumentCaptureInfo()->print(OS);
         }, "", CapturesColor);
       }
 
@@ -1489,11 +1489,12 @@ namespace {
 
     void printCommonAFD(AbstractFunctionDecl *D, const char *Type, StringRef Label) {
       printCommon(D, Type, Label, FuncColor);
-      if (D->getCaptureInfo().hasBeenComputed() &&
-          !D->getCaptureInfo().isTrivial()) {
-        printFlagRaw([&](raw_ostream &OS) {
-          D->getCaptureInfo().print(OS);
-        });
+      if (auto captureInfo = D->getCachedCaptureInfo()) {
+        if (!captureInfo->isTrivial()) {
+          printFlagRaw([&](raw_ostream &OS) {
+            captureInfo->print(OS);
+          });
+        }
       }
 
       if (auto *attr = D->getAttrs().getAttribute<NonisolatedAttr>()) {
@@ -2828,11 +2829,12 @@ class PrintExpr : public ExprVisitor<PrintExpr, void, StringRef>,
       break;
     }
 
-    if (E->getCaptureInfo().hasBeenComputed() &&
-        !E->getCaptureInfo().isTrivial()) {
-      printFieldRaw([&](raw_ostream &OS) {
-        E->getCaptureInfo().print(OS);
-      }, "", CapturesColor);
+    if (auto captureInfo = E->getCachedCaptureInfo()) {
+      if (!captureInfo->isTrivial()) {
+        printFieldRaw([&](raw_ostream &OS) {
+          captureInfo->print(OS);
+        }, "", CapturesColor);
+      }
     }
     // Printing a function type doesn't indicate whether it's escaping because it doesn't 
     // matter in 99% of contexts. AbstractClosureExpr nodes are one of the only exceptions.

lib/AST/Decl.cpp

@@ -8716,6 +8716,16 @@ void ParamDecl::setDefaultArgumentInitContext(Initializer *initContext) {
   defaultInfo->InitContextAndIsTypeChecked.setPointer(initContext);
 }
 
+CaptureInfo ParamDecl::getDefaultArgumentCaptureInfo() const {
+  if (!DefaultValueAndFlags.getPointer())
+    return CaptureInfo::empty();
+
+  auto &ctx = getASTContext();
+  return evaluateOrDefault(ctx.evaluator,
+                           ParamCaptureInfoRequest{const_cast<ParamDecl *>(this)},
+                           CaptureInfo::empty());
+}
+
 void ParamDecl::setDefaultArgumentCaptureInfo(CaptureInfo captures) {
   assert(DefaultValueAndFlags.getPointer());
   assert(captures.hasBeenComputed());
@@ -9171,6 +9181,13 @@ const ParamDecl *swift::getParameterAt(const DeclContext *source,
   return nullptr;
 }
 
+CaptureInfo AbstractFunctionDecl::getCaptureInfo() const {
+  auto &ctx = getASTContext();
+  return evaluateOrDefault(ctx.evaluator,
+                           CaptureInfoRequest{const_cast<AbstractFunctionDecl *>(this)},
+                           CaptureInfo::empty());
+}
+
 Type AbstractFunctionDecl::getMethodInterfaceType() const {
   assert(getDeclContext()->isTypeContext());
   auto Ty = getInterfaceType();

lib/AST/Module.cpp

@@ -3468,6 +3468,7 @@ void SourceFile::typeCheckDelayedFunctions() {
     auto *AFD = DelayedFunctions[i];
     assert(!AFD->getDeclContext()->isLocalContext());
     AFD->getTypecheckedBody();
+    (void) AFD->getCaptureInfo();
   }
 
   DelayedFunctions.clear();

lib/AST/TypeCheckRequests.cpp

@@ -2258,3 +2258,25 @@ void ExpandBodyMacroRequest::noteCycleStep(DiagnosticEngine &diags) const {
                  "body",
                  decl->getName());
 }
+
+std::optional<CaptureInfo>
+CaptureInfoRequest::getCachedResult() const {
+  auto *func = std::get<0>(getStorage());
+  return func->getCachedCaptureInfo();
+}
+
+void CaptureInfoRequest::cacheResult(CaptureInfo info) const {
+  auto *func = std::get<0>(getStorage());
+  return func->setCaptureInfo(info);
+}
+
+std::optional<CaptureInfo>
+ParamCaptureInfoRequest::getCachedResult() const {
+  auto *param = std::get<0>(getStorage());
+  return param->getCachedDefaultArgumentCaptureInfo();
+}
+
+void ParamCaptureInfoRequest::cacheResult(CaptureInfo info) const {
+  auto *param = std::get<0>(getStorage());
+  param->setDefaultArgumentCaptureInfo(info);
+}

lib/SIL/IR/TypeLowering.cpp

@@ -4169,8 +4169,6 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   };
 
   collectCaptures = [&](CaptureInfo captureInfo) {
-    assert(captureInfo.hasBeenComputed());
-
     if (captureInfo.hasGenericParamCaptures())
       capturesGenericParams = true;
     if (captureInfo.hasDynamicSelfCapture())

lib/SILGen/SILGenProlog.cpp

@@ -1235,9 +1235,6 @@ void SILGenFunction::emitProlog(
     SourceLoc throwsLoc) {
   // Emit the capture argument variables. These are placed last because they
   // become the first curry level of the SIL function.
-  assert(captureInfo.hasBeenComputed() &&
-         "can't emit prolog of function with uncomputed captures");
-
   bool hasErasedIsolation =
     (TypeContext && TypeContext->ExpectedLoweredType->hasErasedIsolation());