Merge pull request #72936 from slavapestov/requestify-compute-captures

slavapestov · web-flow · commit ae9eeb099346 · 2024-04-21T09:56:23.000-04:00
Request-ify TypeChecker::computeCaptures()
diff --git a/include/swift/AST/CaptureInfo.h b/include/swift/AST/CaptureInfo.h
@@ -185,24 +185,19 @@ class CaptureInfo {
   }
 
   bool isTrivial() const {
+    assert(hasBeenComputed());
     return getCaptures().empty() && !hasGenericParamCaptures() &&
            !hasDynamicSelfCapture() && !hasOpaqueValueCapture();
   }
 
   ArrayRef<CapturedValue> getCaptures() const {
-    // FIXME: Ideally, everywhere that synthesizes a function should include
-    // its capture info.
-    if (!hasBeenComputed())
-      return std::nullopt;
+    assert(hasBeenComputed());
     return StorageAndFlags.getPointer()->getCaptures();
   }
 
   /// \returns true if the function captures any generic type parameters.
   bool hasGenericParamCaptures() const {
-    // FIXME: Ideally, everywhere that synthesizes a function should include
-    // its capture info.
-    if (!hasBeenComputed())
-      return false;
+    assert(hasBeenComputed());
     return StorageAndFlags.getInt().contains(Flags::HasGenericParamCaptures);
   }
 
@@ -213,22 +208,17 @@ class CaptureInfo {
 
   /// \returns the captured dynamic Self type, if any.
   DynamicSelfType *getDynamicSelfType() const {
-    // FIXME: Ideally, everywhere that synthesizes a function should include
-    // its capture info.
-    if (!hasBeenComputed())
-      return nullptr;
+    assert(hasBeenComputed());
     return StorageAndFlags.getPointer()->getDynamicSelfType();
   }
 
   bool hasOpaqueValueCapture() const {
+    assert(hasBeenComputed());
     return getOpaqueValue() != nullptr;
   }
 
   OpaqueValueExpr *getOpaqueValue() const {
-    // FIXME: Ideally, everywhere that synthesizes a function should include
-    // its capture info.
-    if (!hasBeenComputed())
-      return nullptr;
+    assert(hasBeenComputed());
     return StorageAndFlags.getPointer()->getOpaqueValue();
   }
 
diff --git a/include/swift/AST/Decl.h b/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,8 +7759,18 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Retrieve the source range of the function declaration name + patterns.
   SourceRange getSignatureSourceRange() const;
 
-  CaptureInfo getCaptureInfo() const { return Captures; }
-  void setCaptureInfo(CaptureInfo captures) { Captures = captures; }
+  CaptureInfo getCaptureInfo() const;
+
+  std::optional<CaptureInfo> getCachedCaptureInfo() const {
+    if (!Captures.hasBeenComputed())
+      return std::nullopt;
+    return Captures;
+  }
+
+  void setCaptureInfo(CaptureInfo captures) {
+    assert(captures.hasBeenComputed());
+    Captures = captures;
+  }
 
   /// Retrieve the Objective-C selector that names this method.
   ObjCSelector getObjCSelector(DeclName preferredName = DeclName(),
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -3912,8 +3912,21 @@ class AbstractClosureExpr : public DeclContext, public Expr {
           std::tuple<CapturedValue, unsigned, ApplyIsolationCrossing>>
           &foundIsolationCrossings);
 
-  CaptureInfo getCaptureInfo() const { return Captures; }
-  void setCaptureInfo(CaptureInfo captures) { Captures = 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());
+    Captures = captures;
+  }
 
   /// Retrieve the parameters of this closure.
   ParameterList *getParameters() { return parameterList; }
diff --git a/include/swift/AST/TypeCheckRequests.h b/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
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/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)
diff --git a/lib/AST/ASTDumper.cpp b/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.
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -8716,8 +8716,19 @@ 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());
   DefaultValueAndFlags.getPointer()->Captures = captures;
 }
 
@@ -9170,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();
diff --git a/lib/AST/Module.cpp b/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();
diff --git a/lib/AST/TypeCheckRequests.cpp b/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);
+}
diff --git a/lib/SIL/IR/TypeLowering.cpp b/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())
diff --git a/lib/SILGen/SILGenFunction.cpp b/lib/SILGen/SILGenFunction.cpp
@@ -1002,7 +1002,7 @@ SILGenFunction::emitClosureValue(SILLocation loc, SILDeclRef constant,
   // If we're in top-level code, we don't need to physically capture script
   // globals, but we still need to mark them as escaping so that DI can flag
   // uninitialized uses.
-  if (isEmittingTopLevelCode()) {
+  if (isEmittingTopLevelCode() && dc->getParentSourceFile()) {
     auto captureInfo = closure.getCaptureInfo();
     emitMarkFunctionEscapeForTopLevelCodeGlobals(loc, captureInfo);
   }
diff --git a/lib/SILGen/SILGenProlog.cpp b/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());
 
diff --git a/lib/SILGen/SILGenTopLevel.cpp b/lib/SILGen/SILGenTopLevel.cpp
@@ -335,6 +335,9 @@ void SILGenFunction::emitCallToMain(FuncDecl *mainFunc) {
 }
 
 void SILGenModule::emitEntryPoint(SourceFile *SF) {
+  if (getASTContext().SILOpts.SkipFunctionBodies != FunctionBodySkipping::None)
+    return;
+
   assert(!M.lookUpFunction(getASTContext().getEntryPointFunctionName()) &&
          "already emitted toplevel?!");
 
@@ -366,9 +369,16 @@ void SILGenFunction::emitMarkFunctionEscapeForTopLevelCodeGlobals(
 /// uninitialized global variable
 static void emitMarkFunctionEscape(SILGenFunction &SGF,
                                    AbstractFunctionDecl *AFD) {
+  auto &Ctx = SGF.getASTContext();
+  if (Ctx.TypeCheckerOpts.DeferToRuntime &&
+      Ctx.LangOpts.hasFeature(Feature::LazyImmediate))
+    return;
+
   if (AFD->getDeclContext()->isLocalContext())
     return;
   auto CaptureInfo = AFD->getCaptureInfo();
+  if (!CaptureInfo.hasBeenComputed())
+    return;
   SGF.emitMarkFunctionEscapeForTopLevelCodeGlobals(AFD, std::move(CaptureInfo));
 }
 
diff --git a/lib/Sema/TypeCheckCaptures.cpp b/lib/Sema/TypeCheckCaptures.cpp
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
diff --git a/lib/Sema/TypeCheckStmt.cpp b/lib/Sema/TypeCheckStmt.cpp
diff --git a/lib/Sema/TypeChecker.h b/lib/Sema/TypeChecker.h
diff --git a/test/SILGen/skip_function_bodies_guard.swift b/test/SILGen/skip_function_bodies_guard.swift
diff --git a/test/decl/func/dynamic_self.swift b/test/decl/func/dynamic_self.swift

Original file line number	Diff line number	Diff line change
`@@ -3468,6 +3468,7 @@ void SourceFile::typeCheckDelayedFunctions() {`
`3468`	`3468`	`auto *AFD = DelayedFunctions[i];`
`3469`	`3469`	`assert(!AFD->getDeclContext()->isLocalContext());`
`3470`	`3470`	`AFD->getTypecheckedBody();`
	`3471`	`+ (void) AFD->getCaptureInfo();`
`3471`	`3472`	`}`
`3472`	`3473`
`3473`	`3474`	`DelayedFunctions.clear();`