[Sema] Requestify PatternBindingDecl capture computation

Introduce `PatternBindingCaptureInfoRequest`, and kick it after
contextualizing a property initializer. This ensures it gets run
for stored properties added by macro expansions.

rdar://143429551

Author: hamishknight

include/swift/AST/Decl.h

@@ -2287,6 +2287,7 @@ class PatternBindingEntry {
   // Flags::Checked.
   friend class PatternBindingEntryRequest;
   friend class PatternBindingCheckedAndContextualizedInitRequest;
+  friend class PatternBindingCaptureInfoRequest;
 
   bool isFullyValidated() const {
     return InitContextAndFlags.getInt().contains(
@@ -2435,8 +2436,20 @@ class PatternBindingEntry {
   /// from the source range.
   SourceRange getSourceRange(bool omitAccessors = false) const;
 
-  CaptureInfo getCaptureInfo() const { return Captures; }
-  void setCaptureInfo(CaptureInfo captures) { Captures = captures; }
+  /// Retrieve the computed capture info, or \c nullopt if it hasn't been
+  /// computed yet.
+  std::optional<CaptureInfo> getCachedCaptureInfo() const {
+    if (!Captures.hasBeenComputed())
+      return std::nullopt;
+
+    return Captures;
+  }
+
+  void setCaptureInfo(CaptureInfo captures) {
+    ASSERT(!Captures.hasBeenComputed());
+    ASSERT(captures.hasBeenComputed());
+    Captures = captures;
+  }
 
 private:
   SourceLoc getLastAccessorEndLoc() const;
@@ -2460,6 +2473,7 @@ class PatternBindingDecl final : public Decl,
   friend class Decl;
   friend class PatternBindingEntryRequest;
   friend class PatternBindingCheckedAndContextualizedInitRequest;
+  friend class PatternBindingCaptureInfoRequest;
 
   SourceLoc StaticLoc; ///< Location of the 'static/class' keyword, if present.
   SourceLoc VarLoc;    ///< Location of the 'var' keyword.
@@ -2639,13 +2653,9 @@ class PatternBindingDecl final : public Decl,
     getMutablePatternList()[i].setInitContext(init);
   }
 
-  CaptureInfo getCaptureInfo(unsigned i) const {
-    return getPatternList()[i].getCaptureInfo();
-  }
-
-  void setCaptureInfo(unsigned i, CaptureInfo captures) {
-    getMutablePatternList()[i].setCaptureInfo(captures);
-  }
+  /// Retrieve the capture info for the initializer at the given index,
+  /// computing if needed.
+  CaptureInfo getCaptureInfo(unsigned i) const;
 
   /// Given that this PBD is the parent pattern for the specified VarDecl,
   /// return the entry of the VarDecl in our PatternList.  For example, in:

include/swift/AST/TypeCheckRequests.h

@@ -5134,6 +5134,26 @@ class ParamCaptureInfoRequest :
   void cacheResult(CaptureInfo value) const;
 };
 
+class PatternBindingCaptureInfoRequest
+    : public SimpleRequest<PatternBindingCaptureInfoRequest,
+                           CaptureInfo(PatternBindingDecl *, unsigned),
+                           RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  CaptureInfo evaluate(Evaluator &evaluator, PatternBindingDecl *PBD,
+                       unsigned idx) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  std::optional<CaptureInfo> getCachedResult() const;
+  void cacheResult(CaptureInfo info) const;
+};
+
 class SuppressesConformanceRequest
     : public SimpleRequest<SuppressesConformanceRequest,
                            bool(NominalTypeDecl *decl, KnownProtocolKind kp),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -608,6 +608,9 @@ SWIFT_REQUEST(TypeChecker, CaptureInfoRequest,
 SWIFT_REQUEST(TypeChecker, ParamCaptureInfoRequest,
               CaptureInfo(ParamDecl *),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PatternBindingCaptureInfoRequest,
+              CaptureInfo(PatternBindingDecl *, unsigned),
+              SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, CustomDerivativesRequest,
               CustomDerivativesResult(SourceFile *),
               Cached, NoLocationInfo)

lib/AST/Decl.cpp

@@ -2287,6 +2287,13 @@ PatternBindingDecl *PatternBindingDecl::createDeserialized(
   return PBD;
 }
 
+CaptureInfo PatternBindingDecl::getCaptureInfo(unsigned i) const {
+  auto *mutableThis = const_cast<PatternBindingDecl *>(this);
+  PatternBindingCaptureInfoRequest req(mutableThis, i);
+  return evaluateOrDefault(getASTContext().evaluator, req,
+                           CaptureInfo::empty());
+}
+
 PatternBindingInitializer *
 PatternBindingInitializer::createDeserialized(PatternBindingDecl *PBD,
                                               unsigned index) {

lib/AST/TypeCheckRequests.cpp

@@ -2751,6 +2751,23 @@ void ParamCaptureInfoRequest::cacheResult(CaptureInfo info) const {
   param->setDefaultArgumentCaptureInfo(info);
 }
 
+//----------------------------------------------------------------------------//
+// PatternBindingCaptureInfoRequest caching.
+//----------------------------------------------------------------------------//
+
+std::optional<CaptureInfo>
+PatternBindingCaptureInfoRequest::getCachedResult() const {
+  auto *PBD = std::get<0>(getStorage());
+  auto idx = std::get<1>(getStorage());
+  return PBD->getPatternList()[idx].getCachedCaptureInfo();
+}
+
+void PatternBindingCaptureInfoRequest::cacheResult(CaptureInfo info) const {
+  auto *PBD = std::get<0>(getStorage());
+  auto idx = std::get<1>(getStorage());
+  PBD->getMutablePatternList()[idx].setCaptureInfo(info);
+}
+
 //----------------------------------------------------------------------------//
 // SemanticAvailableAttrRequest computation.
 //----------------------------------------------------------------------------//

lib/Sema/TypeCheckCaptures.cpp

@@ -863,37 +863,30 @@ static bool isLazy(PatternBindingDecl *PBD) {
   return false;
 }
 
-void TypeChecker::checkPatternBindingCaptures(IterableDeclContext *DC) {
-  for (auto member : DC->getMembers()) {
-    // Ignore everything other than PBDs.
-    auto *PBD = dyn_cast<PatternBindingDecl>(member);
-    if (!PBD) continue;
-    // Walk the initializers for all properties declared in the type with
-    // an initializer.
-    for (unsigned i : range(PBD->getNumPatternEntries())) {
-      if (PBD->isInitializerSubsumed(i))
-        continue;
+CaptureInfo PatternBindingCaptureInfoRequest::evaluate(Evaluator &evaluator,
+                                                       PatternBindingDecl *PBD,
+                                                       unsigned int idx) const {
+  auto *init = PBD->getExecutableInit(idx);
+  if (!init)
+    return CaptureInfo::empty();
 
-      auto *init = PBD->getInit(i);
-      if (init == nullptr)
-        continue;
+  // Only have captures when we have a PatternBindingInitializer context, i.e
+  // local variables don't have captures.
+  auto *DC = PBD->getInitContext(idx);
+  if (!DC)
+    return CaptureInfo::empty();
 
-      auto *DC = PBD->getInitContext(i);
-      FindCapturedVars finder(init->getLoc(),
-                              DC,
-                              /*NoEscape=*/false,
-                              /*ObjC=*/false,
-                              /*IsGenericFunction*/false);
-      init->walk(finder);
-
-      auto &ctx = DC->getASTContext();
-      if (finder.getDynamicSelfCaptureLoc().isValid() && !isLazy(PBD)) {
-        ctx.Diags.diagnose(finder.getDynamicSelfCaptureLoc(),
-                           diag::dynamic_self_stored_property_init);
-      }
+  FindCapturedVars finder(init->getLoc(), DC,
+                          /*NoEscape=*/false,
+                          /*ObjC=*/false,
+                          /*IsGenericFunction*/ false);
+  init->walk(finder);
 
-      auto captures = finder.getCaptureInfo();
-      PBD->setCaptureInfo(i, captures);
-    }
+  auto &ctx = DC->getASTContext();
+  if (finder.getDynamicSelfCaptureLoc().isValid() && !isLazy(PBD)) {
+    ctx.Diags.diagnose(finder.getDynamicSelfCaptureLoc(),
+                       diag::dynamic_self_stored_property_init);
   }
+
+  return finder.getCaptureInfo();
 }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -3236,8 +3236,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     diagnoseMissingExplicitSendable(ED);
     checkAccessControl(ED);
 
-    TypeChecker::checkPatternBindingCaptures(ED);
-
     auto &DE = Ctx.Diags;
     if (auto rawTy = ED->getRawType()) {
       // The raw type must be one of the blessed literal convertible types.
@@ -3307,8 +3305,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
       visit(Member);
     }
 
-    TypeChecker::checkPatternBindingCaptures(SD);
-
     checkInheritanceClause(SD);
     diagnoseMissingExplicitSendable(SD);
 
@@ -3490,8 +3486,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     for (Decl *Member : CD->getABIMembers())
       visit(Member);
 
-    TypeChecker::checkPatternBindingCaptures(CD);
-
     // If this class requires all of its stored properties to have
     // in-class initializers, diagnose this now.
     if (CD->requiresStoredPropertyInits())
@@ -4074,8 +4068,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     for (Decl *Member : ED->getMembers())
       visit(Member);
 
-    TypeChecker::checkPatternBindingCaptures(ED);
-
     TypeChecker::checkConformancesInContext(ED);
 
     checkAccessControl(ED);

lib/Sema/TypeCheckStorage.cpp

@@ -619,6 +619,9 @@ static void checkAndContextualizePatternBindingInit(PatternBindingDecl *binding,
     auto *init = binding->getInit(i);
     TypeChecker::contextualizeExpr(init, initContext);
   }
+
+  // Compute captures in case diagnostics are emitted.
+  (void)binding->getCaptureInfo(i);
 }
 
 Expr *PatternBindingCheckedAndContextualizedInitRequest::evaluate(

lib/Sema/TypeChecker.h

@@ -798,9 +798,6 @@ bool typeCheckForEachPreamble(DeclContext *dc, ForEachStmt *stmt);
 /// Compute the set of captures for the given closure.
 void computeCaptures(AbstractClosureExpr *ACE);
 
-/// Check for invalid captures from stored property initializers.
-void checkPatternBindingCaptures(IterableDeclContext *DC);
-
 /// Update the DeclContexts for AST nodes in a given DeclContext. This is
 /// necessary after type-checking since autoclosures may have been introduced.
 void contextualizeExpr(Expr *E, DeclContext *DC);

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2134,6 +2134,20 @@ public struct VarValueMacro: DeclarationMacro, PeerMacro {
   }
 }
 
+struct StoredPropertyMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard let argument = node.arguments.first?.expression else {
+      fatalError("boom")
+    }
+    return [
+      "var storedProperty = \(argument)"
+    ]
+  }
+}
+
 public struct GenericToVoidMacro: ExpressionMacro {
   public static func expansion(
     of node: some FreestandingMacroExpansionSyntax,