Merge pull request swiftlang#82354 from meg-gupta/lifetimeenum

Add lifetime dependencies on function types representing ~Escapable enum elements with ~Escapable payloads

Author: meg-gupta

include/swift/AST/Decl.h

@@ -8779,7 +8779,8 @@ class EnumCaseDecl final : public Decl,
 /// EnumCaseDecl.
 class EnumElementDecl : public DeclContext, public ValueDecl {
   friend class EnumRawValuesRequest;
-  
+  friend class LifetimeDependenceInfoRequest;
+
   /// This is the type specified with the enum element, for
   /// example 'Int' in 'case Y(Int)'.  This is null if there is no type
   /// associated with this element, as in 'case Z' or in all elements of enum
@@ -8791,6 +8792,11 @@ class EnumElementDecl : public DeclContext, public ValueDecl {
   /// The raw value literal for the enum element, or null.
   LiteralExpr *RawValueExpr;
 
+protected:
+  struct {
+    unsigned NoLifetimeDependenceInfo : 1;
+  } LazySemanticInfo = {};
+
 public:
   EnumElementDecl(SourceLoc IdentifierLoc, DeclName Name,
                   ParameterList *Params,

include/swift/AST/LifetimeDependence.h

@@ -324,8 +324,7 @@ class LifetimeDependenceInfo {
   /// Builds LifetimeDependenceInfo from a swift decl, either from the explicit
   /// lifetime dependence specifiers or by inference based on types and
   /// ownership modifiers.
-  static std::optional<ArrayRef<LifetimeDependenceInfo>>
-  get(AbstractFunctionDecl *decl);
+  static std::optional<ArrayRef<LifetimeDependenceInfo>> get(ValueDecl *decl);
 
   /// Builds LifetimeDependenceInfo from SIL
   static std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>

include/swift/AST/TypeCheckRequests.h

@@ -5114,8 +5114,7 @@ class ImportDeclRequest
 class LifetimeDependenceInfoRequest
     : public SimpleRequest<
           LifetimeDependenceInfoRequest,
-          std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>(
-              AbstractFunctionDecl *),
+          std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>(ValueDecl *),
           RequestFlags::SeparatelyCached | RequestFlags::SplitCached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -5124,7 +5123,7 @@ class LifetimeDependenceInfoRequest
   friend SimpleRequest;
 
   std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>
-  evaluate(Evaluator &evaluator, AbstractFunctionDecl *AFD) const;
+  evaluate(Evaluator &evaluator, ValueDecl *AFD) const;
 
 public:
   // Separate caching.

lib/AST/LifetimeDependence.cpp

@@ -252,7 +252,7 @@ void LifetimeDependenceInfo::getConcatenatedData(
 }
 
 class LifetimeDependenceChecker {
-  AbstractFunctionDecl *afd;
+  ValueDecl *decl;
 
   DeclContext *dc;
   ASTContext &ctx;
@@ -273,9 +273,8 @@ class LifetimeDependenceChecker {
   bool performedDiagnostics = false;
 
 public:
-  LifetimeDependenceChecker(AbstractFunctionDecl *afd):
-    afd(afd), dc(afd->getDeclContext()), ctx(dc->getASTContext())
-  {
+  LifetimeDependenceChecker(AbstractFunctionDecl *afd)
+      : decl(afd), dc(afd->getDeclContext()), ctx(dc->getASTContext()) {
     auto resultTypeRepr = afd->getResultTypeRepr();
     returnLoc = resultTypeRepr ? resultTypeRepr->getLoc() : afd->getLoc();
 
@@ -287,18 +286,25 @@ class LifetimeDependenceChecker {
     }
   }
 
+  LifetimeDependenceChecker(EnumElementDecl *eed)
+      : decl(eed), dc(eed->getDeclContext()), ctx(dc->getASTContext()) {
+    auto *paramList = eed->getParameterList();
+    resultIndex = paramList ? eed->getParameterList()->size() + 1 : 1;
+  }
+
   std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>
   currentDependencies() const {
     if (lifetimeDependencies.empty()) {
       return std::nullopt;
     }
-    return afd->getASTContext().AllocateCopy(lifetimeDependencies);
+    return decl->getASTContext().AllocateCopy(lifetimeDependencies);
   }
 
   std::optional<llvm::ArrayRef<LifetimeDependenceInfo>> checkFuncDecl() {
-    assert(isa<FuncDecl>(afd) || isa<ConstructorDecl>(afd));
+    assert(isa<FuncDecl>(decl) || isa<ConstructorDecl>(decl));
     assert(lifetimeDependencies.empty());
 
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     // Handle Builtins first because, even though Builtins require
     // LifetimeDependence, we don't force the experimental feature
     // to be enabled when importing the Builtin module.
@@ -351,6 +357,52 @@ class LifetimeDependenceChecker {
     return currentDependencies();
   }
 
+  std::optional<llvm::ArrayRef<LifetimeDependenceInfo>> checkEnumElementDecl() {
+    auto *eed = cast<EnumElementDecl>(decl);
+    auto *parentEnum = eed->getParentEnum();
+    auto enumType =
+        parentEnum->mapTypeIntoContext(parentEnum->getDeclaredInterfaceType());
+
+    // Add early bailout for imported enums.
+    if (parentEnum->hasClangNode()) {
+      return std::nullopt;
+    }
+
+    // Escapable enum, bailout.
+    if (!isDiagnosedNonEscapable(enumType)) {
+      return std::nullopt;
+    }
+    auto *params = eed->getParameterList();
+    // No payload, bailout.
+    if (!params) {
+      return std::nullopt;
+    }
+
+    auto resultIndex = params->size() + /*selfType*/ 1;
+    auto capacity = resultIndex + 1;
+    SmallBitVector inheritIndices(capacity);
+    SmallVector<LifetimeDependenceInfo, 1> lifetimeDependencies;
+
+    // Add all indices of ~Escapable parameters as lifetime dependence sources.
+    for (size_t i = 0; i < params->size(); i++) {
+      auto paramType = params->get(i)->getTypeInContext();
+      if (!isDiagnosedNonEscapable(paramType)) {
+        continue;
+      }
+      inheritIndices.set(i);
+    }
+    if (inheritIndices.none()) {
+      return std::nullopt;
+    }
+    auto lifetimeDependenceInfo = LifetimeDependenceInfo(
+        IndexSubset::get(eed->getASTContext(), inheritIndices), nullptr,
+        resultIndex,
+        /*isImmortal*/ false);
+    lifetimeDependencies.push_back(lifetimeDependenceInfo);
+
+    return eed->getASTContext().AllocateCopy(lifetimeDependencies);
+  }
+
 protected:
   template<typename ...ArgTypes>
   InFlightDiagnostic diagnose(
@@ -367,9 +419,7 @@ class LifetimeDependenceChecker {
     return ctx.Diags.diagnose(decl, Diagnostic(id, std::move(args)...));
   }
 
-  bool isInit() const {
-    return isa<ConstructorDecl>(afd);
-  }
+  bool isInit() const { return isa<ConstructorDecl>(decl); }
 
   // For initializers, the implicit self parameter is ignored and instead shows
   // up as the result type.
@@ -381,11 +431,13 @@ class LifetimeDependenceChecker {
   // the extra formal self parameter, a dependency targeting the formal result
   // index would incorrectly target the SIL metatype parameter.
   bool hasImplicitSelfParam() const {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     return !isInit() && afd->hasImplicitSelfDecl();
   }
 
   // In SIL, implicit initializers and accessors become explicit.
   bool isImplicitOrSIL() const {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (afd->isImplicit()) {
       return true;
     }
@@ -404,7 +456,7 @@ class LifetimeDependenceChecker {
   bool isInterfaceFile() const {
     // TODO: remove this check once all compilers that are rev-locked to the
     // stdlib print the 'copy' dependence kind in the interface (Aug '25)
-    if (auto *sf = afd->getParentSourceFile()) {
+    if (auto *sf = decl->getDeclContext()->getParentSourceFile()) {
       if (sf->Kind == SourceFileKind::Interface) {
         return true;
       }
@@ -419,6 +471,7 @@ class LifetimeDependenceChecker {
   }
 
   std::string diagnosticQualifier() const {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (afd->isImplicit()) {
       if (isInit()) {
         return "an implicit initializer";
@@ -463,6 +516,7 @@ class LifetimeDependenceChecker {
   // initializers, the inout self parameter is actually considered the result
   // type so is not handled here.
   void diagnoseMissingSelfDependencies(DiagID diagID) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (!hasImplicitSelfParam()) {
       return;
     }
@@ -483,6 +537,7 @@ class LifetimeDependenceChecker {
   }
 
   void diagnoseMissingInoutDependencies(DiagID diagID) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     unsigned paramIndex = 0;
     for (auto *param : *afd->getParameters()) {
       SWIFT_DEFER { paramIndex++; };
@@ -529,6 +584,7 @@ class LifetimeDependenceChecker {
 
   bool isCompatibleWithOwnership(LifetimeDependenceKind kind, Type type,
                                  ValueOwnership ownership) const {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (kind == LifetimeDependenceKind::Inherit) {
       return true;
     }
@@ -569,6 +625,7 @@ class LifetimeDependenceChecker {
   };
 
   TargetDeps createDeps(unsigned targetIndex) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     unsigned capacity = afd->hasImplicitSelfDecl()
       ? (afd->getParameters()->size() + 1)
       : afd->getParameters()->size();
@@ -599,6 +656,7 @@ class LifetimeDependenceChecker {
   }
 
   Type getResultOrYield() const {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (auto *accessor = dyn_cast<AccessorDecl>(afd)) {
       if (accessor->isCoroutine()) {
         auto yieldTyInContext = accessor->mapTypeIntoContext(
@@ -618,11 +676,12 @@ class LifetimeDependenceChecker {
 
   std::optional<LifetimeDependenceKind>
   getDependenceKindFromDescriptor(LifetimeDescriptor descriptor,
-                                  ParamDecl *decl) {
+                                  ParamDecl *paramDecl) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     auto loc = descriptor.getLoc();
-    auto type = decl->getTypeInContext();
+    auto type = paramDecl->getTypeInContext();
     auto parsedLifetimeKind = descriptor.getParsedLifetimeDependenceKind();
-    auto ownership = decl->getValueOwnership();
+    auto ownership = paramDecl->getValueOwnership();
     auto loweredOwnership = ownership != ValueOwnership::Default
                                 ? ownership
                                 : getLoweredOwnership(afd);
@@ -704,6 +763,7 @@ class LifetimeDependenceChecker {
   // Finds the ParamDecl* and its index from a LifetimeDescriptor
   std::optional<std::pair<ParamDecl *, unsigned>>
   getParamDeclFromDescriptor(LifetimeDescriptor descriptor) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     switch (descriptor.getDescriptorKind()) {
     case LifetimeDescriptor::DescriptorKind::Named: {
       unsigned paramIndex = 0;
@@ -752,6 +812,7 @@ class LifetimeDependenceChecker {
   }
 
   std::optional<ArrayRef<LifetimeDependenceInfo>> checkAttribute() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     SmallVector<LifetimeDependenceInfo, 1> lifetimeDependencies;
     llvm::SmallSet<unsigned, 1> lifetimeDependentTargets;
     auto lifetimeAttrs = afd->getAttrs().getAttributes<LifetimeAttr>();
@@ -776,6 +837,7 @@ class LifetimeDependenceChecker {
 
   std::optional<LifetimeDependenceInfo>
   checkAttributeEntry(LifetimeEntry *entry) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     auto capacity = afd->hasImplicitSelfDecl()
       ? (afd->getParameters()->size() + 1)
       : afd->getParameters()->size();
@@ -897,6 +959,7 @@ class LifetimeDependenceChecker {
   /// If the current function is a mutating method and 'self' is non-Escapable,
   /// return 'self's ParamDecl.
   bool isMutatingNonEscapableSelf() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (!hasImplicitSelfParam())
       return false;
 
@@ -914,6 +977,7 @@ class LifetimeDependenceChecker {
   // Infer method dependence of result on self for
   // methods, getters, and _modify accessors.
   void inferNonEscapableResultOnSelf() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     Type selfTypeInContext = dc->getSelfTypeInContext();
     if (selfTypeInContext->hasError()) {
       return;
@@ -980,6 +1044,7 @@ class LifetimeDependenceChecker {
   // stored property (for getters or initializers).
   std::optional<LifetimeDependenceKind>
   inferLifetimeDependenceKind(Type sourceType, ValueOwnership ownership) {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (!sourceType->isEscapable()) {
       return LifetimeDependenceKind::Inherit;
     }
@@ -1003,6 +1068,7 @@ class LifetimeDependenceChecker {
   // to an implicit setter, because the implementation is simply an assignment
   // to stored property.
   void inferImplicitInit() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (afd->getParameters()->size() == 0) {
       // Empty ~Escapable types can be implicitly initialized without any
       // dependencies. In SIL, implicit initializers become explicit. Set
@@ -1042,6 +1108,7 @@ class LifetimeDependenceChecker {
   // inference if any exist, infer scoped dependency, or infer no
   // dependency. Implicit setters for Escapable properties are not inferred.
   void inferNonEscapableResultOnParam() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     // This is only called when there is no 'self' argument that can be the
     // source of a dependence.
     assert(!hasImplicitSelfParam());
@@ -1091,6 +1158,7 @@ class LifetimeDependenceChecker {
   // Lazy inference for .swiftinterface backward compatibility and
   // experimentation. Inference cases can be added but not removed.
   void lazillyInferNonEscapableResultOnParam() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     std::optional<unsigned> candidateParamIndex;
     std::optional<LifetimeDependenceKind> candidateLifetimeKind;
     unsigned paramIndex = 0;
@@ -1137,6 +1205,7 @@ class LifetimeDependenceChecker {
   // Infer a mutating 'self' dependency when 'self' is non-Escapable and the
   // result is 'void'.
   void inferMutatingSelf() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (!isMutatingNonEscapableSelf()) {
       return;
     }
@@ -1166,7 +1235,7 @@ class LifetimeDependenceChecker {
   // Infer dependence for an accessor whose non-escapable result depends on
   // self. This includes _read and _modify.
   void inferAccessor(AccessorDecl *accessor, Type selfTypeInContext) {
-    // Explicit accessors require explicit lifetime dependencies.
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (!isImplicitOrSIL() && !useLazyInference()) {
       return;
     }
@@ -1205,6 +1274,7 @@ class LifetimeDependenceChecker {
       break;
     case AccessorKind::Set: {
       const unsigned newValIdx = 0;
+      auto *afd = cast<AbstractFunctionDecl>(decl);
       auto *param = afd->getParameters()->get(newValIdx);
       Type paramTypeInContext =
         afd->mapTypeIntoContext(param->getInterfaceType());
@@ -1285,6 +1355,7 @@ class LifetimeDependenceChecker {
         }
       }
     }
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     // Either a Get or Modify without any wrapped accessor. Handle these like a
     // read of the stored property.
     return inferLifetimeDependenceKind(
@@ -1315,6 +1386,7 @@ class LifetimeDependenceChecker {
   // Do not issue any diagnostics. This inference is triggered even when the
   // feature is disabled!
   void inferInoutParams() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (isMutatingNonEscapableSelf()) {
       return;
     }
@@ -1347,6 +1419,7 @@ class LifetimeDependenceChecker {
   }
 
   void inferUnambiguousInoutParams() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     if (afd->getParameters()->size() != 1) {
       return;
     }
@@ -1364,6 +1437,7 @@ class LifetimeDependenceChecker {
   }
 
   void inferBuiltin() {
+    auto *afd = cast<AbstractFunctionDecl>(decl);
     // Normal inout parameter inference works for most generic Builtins.
     inferUnambiguousInoutParams();
     if (!lifetimeDependencies.empty()) {
@@ -1397,8 +1471,12 @@ class LifetimeDependenceChecker {
 };
 
 std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>
-LifetimeDependenceInfo::get(AbstractFunctionDecl *afd) {
-  return LifetimeDependenceChecker(afd).checkFuncDecl();
+LifetimeDependenceInfo::get(ValueDecl *decl) {
+  if (auto *afd = dyn_cast<AbstractFunctionDecl>(decl)) {
+    return LifetimeDependenceChecker(afd).checkFuncDecl();
+  }
+  auto *eed = cast<EnumElementDecl>(decl);
+  return LifetimeDependenceChecker(eed).checkEnumElementDecl();
 }
 
 // This implements the logic for SIL type descriptors similar to source-level