Merge pull request #80535 from hamishknight/cap-req

[Sema] Requestify PatternBindingDecl capture computation

Author: hamishknight

include/swift/AST/Decl.h

@@ -2284,6 +2284,7 @@ class PatternBindingEntry {
   // Flags::Checked.
   friend class PatternBindingEntryRequest;
   friend class PatternBindingCheckedAndContextualizedInitRequest;
+  friend class PatternBindingCaptureInfoRequest;
 
   bool isFullyValidated() const {
     return InitContextAndFlags.getInt().contains(
@@ -2432,8 +2433,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;
@@ -2457,6 +2470,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.
@@ -2636,13 +2650,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

@@ -2291,6 +2291,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

@@ -2744,6 +2744,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

@@ -856,43 +856,30 @@ CaptureInfo ParamCaptureInfoRequest::evaluate(Evaluator &evaluator,
   return finder.getCaptureInfo();
 }
 
-static bool isLazy(PatternBindingDecl *PBD) {
-  if (auto var = PBD->getSingleVar())
-    return var->getAttrs().hasAttribute<LazyAttr>();
-  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()) {
+    ctx.Diags.diagnose(finder.getDynamicSelfCaptureLoc(),
+                       diag::dynamic_self_stored_property_init);
   }
+
+  return finder.getCaptureInfo();
 }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -3233,8 +3233,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.
@@ -3304,8 +3302,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
       visit(Member);
     }
 
-    TypeChecker::checkPatternBindingCaptures(SD);
-
     checkInheritanceClause(SD);
     diagnoseMissingExplicitSendable(SD);
 
@@ -3487,8 +3483,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())
@@ -4071,8 +4065,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,

test/Macros/macro_expand.swift

@@ -63,6 +63,9 @@ macro freestandingWithClosure<T>(_ value: T, body: (T) -> T) = #externalMacro(mo
 
 #endif
 
+@freestanding(declaration, names: named(storedProperty))
+public macro AddStoredProperty<T>(_ x: T) = #externalMacro(module: "MacroDefinition", type: "StoredPropertyMacro")
+
 #if TEST_DIAGNOSTICS
 @freestanding(declaration)
 macro NotCovered() = #externalMacro(module: "MacroDefinition", type: "InvalidMacro")
@@ -72,7 +75,7 @@ struct MemberNotCovered {
   // expected-note@-1 {{in expansion of macro 'NotCovered' here}}
 
   // CHECK-DIAGS: error: declaration name 'value' is not covered by macro 'NotCovered'
-  // CHECK-DIAGS: CONTENTS OF FILE @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX70_2_33_4361AD9339943F52AE6186DD51E04E91Ll10NotCoveredfMf_.swift
+  // CHECK-DIAGS: CONTENTS OF FILE @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX[[@LINE-5]]_2_33_4361AD9339943F52AE6186DD51E04E91Ll10NotCoveredfMf_.swift
   // CHECK-DIAGS: var value: Int
   // CHECK-DIAGS: END CONTENTS OF FILE
 }
@@ -126,6 +129,11 @@ struct Bad {}
 // CHECK-DIAGS: typealias _ImageLiteralType = Void
 // CHECK-DIAGS: typealias _FileReferenceLiteralType = Void
 // CHECK-DIAGS: END CONTENTS OF FILE
+
+class HasStoredPropertyClassInvalid {
+  #AddStoredProperty((Self.self, 0).1) // expected-note {{in expansion of macro 'AddStoredProperty' here}}
+  // CHECK-DIAGS: @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX[[@LINE-2]]_2_33_{{.*}}AddStoredPropertyfMf_.swift:1:22: error: covariant 'Self' type cannot be referenced from a stored property initializer
+}
 #endif
 
 @freestanding(declaration)
@@ -138,7 +146,7 @@ macro AccidentalCodeItem() = #externalMacro(module: "MacroDefinition", type: "Fa
 func invalidDeclarationMacro() {
   #accidentalCodeItem
   // expected-note@-1 {{in expansion of macro 'accidentalCodeItem' here}}
-  // CHECK-DIAGS: @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX138_2_18accidentalCodeItemfMf_.swift:1:1: error: expected macro expansion to produce a declaration
+  // CHECK-DIAGS: @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX[[@LINE-3]]_2_18accidentalCodeItemfMf_.swift:1:1: error: expected macro expansion to produce a declaration
 
   @AccidentalCodeItem struct S {}
   // expected-note@-1 {{in expansion of macro 'AccidentalCodeItem' on struct 'S' here}}
@@ -681,3 +689,11 @@ func invalidDeclarationMacro2() {
   }
 }
 #endif
+
+// Make sure we compute captures for the introduced stored property here.
+struct HasStoredProperty {
+  #AddStoredProperty(0)
+}
+class HasStoredPropertyClass {
+  #AddStoredProperty(0)
+}