Support lifetime dependence inference on enum elements

Author: meg-gupta

include/swift/AST/Decl.h

@@ -8780,7 +8780,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
@@ -8792,6 +8793,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

@@ -328,8 +328,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

@@ -5095,8 +5095,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;
@@ -5105,7 +5104,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

@@ -286,6 +286,12 @@ 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()) {
@@ -351,6 +357,45 @@ class LifetimeDependenceChecker {
     return currentDependencies();
   }
 
+  std::optional<llvm::ArrayRef<LifetimeDependenceInfo>> checkEnumElementDecl() {
+    auto *eed = cast<EnumElementDecl>(decl);
+    auto enumType = eed->getParentEnum()->mapTypeIntoContext(
+        eed->getParentEnum()->getDeclaredInterfaceType());
+    // 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(
@@ -1334,8 +1379,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

lib/AST/TypeCheckRequests.cpp

@@ -2738,24 +2738,36 @@ void ExpandBodyMacroRequest::noteCycleStep(DiagnosticEngine &diags) const {
 
 std::optional<std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>>
 LifetimeDependenceInfoRequest::getCachedResult() const {
-  auto *func = std::get<0>(getStorage());
+  auto *decl = std::get<0>(getStorage());
+  bool noLifetimeDependenceInfo;
 
-  if (func->LazySemanticInfo.NoLifetimeDependenceInfo)
+  if (auto *func = dyn_cast<AbstractFunctionDecl>(decl)) {
+    noLifetimeDependenceInfo = func->LazySemanticInfo.NoLifetimeDependenceInfo;
+  } else {
+    auto *eed = cast<EnumElementDecl>(decl);
+    noLifetimeDependenceInfo = eed->LazySemanticInfo.NoLifetimeDependenceInfo;
+  }
+
+  if (noLifetimeDependenceInfo)
     return std::optional(std::optional<LifetimeDependenceInfo>());
 
-  return func->getASTContext().evaluator.getCachedNonEmptyOutput(*this);
+  return decl->getASTContext().evaluator.getCachedNonEmptyOutput(*this);
 }
 
 void LifetimeDependenceInfoRequest::cacheResult(
     std::optional<llvm::ArrayRef<LifetimeDependenceInfo>> result) const {
-  auto *func = std::get<0>(getStorage());
-  
+  auto *decl = std::get<0>(getStorage());
+
   if (!result) {
-    func->LazySemanticInfo.NoLifetimeDependenceInfo = 1;
-    return;
+    if (auto *func = dyn_cast<AbstractFunctionDecl>(decl)) {
+      func->LazySemanticInfo.NoLifetimeDependenceInfo = 1;
+      return;
+    }
+    auto *eed = cast<EnumElementDecl>(decl);
+    eed->LazySemanticInfo.NoLifetimeDependenceInfo = 1;
   }
 
-  func->getASTContext().evaluator.cacheNonEmptyOutput(*this, std::move(result));
+  decl->getASTContext().evaluator.cacheNonEmptyOutput(*this, std::move(result));
 }
 
 //----------------------------------------------------------------------------//

lib/Sema/TypeCheckDecl.cpp

@@ -2409,6 +2409,13 @@ static void maybeAddParameterIsolation(AnyFunctionType::ExtInfoBuilder &infoBuil
     infoBuilder = infoBuilder.withIsolation(FunctionTypeIsolation::forParameter());
 }
 
+std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>
+getLifetimeDependencies(ASTContext &context, EnumElementDecl *enumElemDecl) {
+  return evaluateOrDefault(context.evaluator,
+                           LifetimeDependenceInfoRequest{enumElemDecl},
+                           std::nullopt);
+}
+
 Type
 InterfaceTypeRequest::evaluate(Evaluator &eval, ValueDecl *D) const {
   auto &Context = D->getASTContext();
@@ -2698,13 +2705,25 @@ InterfaceTypeRequest::evaluate(Evaluator &eval, ValueDecl *D) const {
       resultTy = FunctionType::get(argTy, resultTy, info);
     }
 
+    auto lifetimeDependenceInfo = getLifetimeDependencies(Context, EED);
+
     // FIXME: Verify ExtInfo state is correct, not working by accident.
     if (auto genericSig = ED->getGenericSignature()) {
-      GenericFunctionType::ExtInfo info;
-      resultTy = GenericFunctionType::get(genericSig, {selfTy}, resultTy, info);
+      GenericFunctionType::ExtInfoBuilder infoBuilder;
+      if (lifetimeDependenceInfo.has_value()) {
+        infoBuilder =
+            infoBuilder.withLifetimeDependencies(*lifetimeDependenceInfo);
+      }
+      resultTy = GenericFunctionType::get(genericSig, {selfTy}, resultTy,
+                                          infoBuilder.build());
+
     } else {
-      FunctionType::ExtInfo info;
-      resultTy = FunctionType::get({selfTy}, resultTy, info);
+      FunctionType::ExtInfoBuilder infoBuilder;
+      if (lifetimeDependenceInfo.has_value()) {
+        infoBuilder =
+            infoBuilder.withLifetimeDependencies(*lifetimeDependenceInfo);
+      }
+      resultTy = FunctionType::get({selfTy}, resultTy, infoBuilder.build());
     }
 
     return resultTy;
@@ -3153,7 +3172,7 @@ ImplicitKnownProtocolConformanceRequest::evaluate(Evaluator &evaluator,
 
 std::optional<llvm::ArrayRef<LifetimeDependenceInfo>>
 LifetimeDependenceInfoRequest::evaluate(Evaluator &evaluator,
-                                        AbstractFunctionDecl *decl) const {
+                                        ValueDecl *decl) const {
   return LifetimeDependenceInfo::get(decl);
 }