Fix serialization of lifetime dependence

include/swift/AST/ExtInfo.h

@@ -1335,6 +1335,10 @@ class SILExtInfo {
     return builder.withTransferringResult(hasTransferringResult).build();
   }
 
+  SILExtInfo withLifetimeDependenceInfo(LifetimeDependenceInfo info) const {
+    return builder.withLifetimeDependenceInfo(info);
+  }
+
   void Profile(llvm::FoldingSetNodeID &ID) const { builder.Profile(ID); }
 
   bool isEqualTo(SILExtInfo other, bool useClangTypes) const {

include/swift/AST/LifetimeDependence.h

@@ -132,6 +132,7 @@ class LifetimeDependenceSpecifier {
 class LifetimeDependenceInfo {
   IndexSubset *inheritLifetimeParamIndices;
   IndexSubset *scopeLifetimeParamIndices;
+  bool isExplicit;
 
   static LifetimeDependenceInfo getForParamIndex(AbstractFunctionDecl *afd,
                                                  unsigned index,
@@ -146,11 +147,18 @@ class LifetimeDependenceInfo {
 public:
   LifetimeDependenceInfo()
       : inheritLifetimeParamIndices(nullptr),
-        scopeLifetimeParamIndices(nullptr) {}
+        scopeLifetimeParamIndices(nullptr), isExplicit(false) {}
   LifetimeDependenceInfo(IndexSubset *inheritLifetimeParamIndices,
-                         IndexSubset *scopeLifetimeParamIndices)
+                         IndexSubset *scopeLifetimeParamIndices,
+                         bool isExplicit = false)
       : inheritLifetimeParamIndices(inheritLifetimeParamIndices),
-        scopeLifetimeParamIndices(scopeLifetimeParamIndices) {}
+        scopeLifetimeParamIndices(scopeLifetimeParamIndices),
+        isExplicit(isExplicit) {
+    assert(!empty());
+    assert(!inheritLifetimeParamIndices ||
+           !inheritLifetimeParamIndices->isEmpty());
+    assert(!scopeLifetimeParamIndices || !scopeLifetimeParamIndices->isEmpty());
+  }
 
   operator bool() const { return !empty(); }
 
@@ -159,6 +167,8 @@ class LifetimeDependenceInfo {
            scopeLifetimeParamIndices == nullptr;
   }
 
+  bool isExplicitlySpecified() const { return isExplicit; }
+
   bool hasInheritLifetimeParamIndices() const {
     return inheritLifetimeParamIndices != nullptr;
   }

include/swift/AST/Types.h

@@ -3433,7 +3433,14 @@ class AnyFunctionType : public TypeBase {
   Type getGlobalActor() const;
   Type getThrownError() const;
 
-  LifetimeDependenceInfo getLifetimeDependenceInfo() const;
+  const LifetimeDependenceInfo *getLifetimeDependenceInfoOrNull() const;
+
+  LifetimeDependenceInfo getLifetimeDependenceInfo() const {
+    if (auto *depInfo = getLifetimeDependenceInfoOrNull()) {
+      return *depInfo;
+    }
+    return LifetimeDependenceInfo();
+  }
 
   FunctionTypeIsolation getIsolation() const {
     if (hasExtInfo())
@@ -3765,14 +3772,22 @@ class FunctionType final
     return getTrailingObjects<Type>()[hasGlobalActor()];
   }
 
-  LifetimeDependenceInfo getLifetimeDependenceInfo() const {
+  inline LifetimeDependenceInfo getLifetimeDependenceInfo() const {
+    if (auto *depInfo = getLifetimeDependenceInfoOrNull()) {
+      return *depInfo;
+    }
+    return LifetimeDependenceInfo();
+  }
+
+  /// Returns nullptr for an empty dependence list.
+  const LifetimeDependenceInfo *getLifetimeDependenceInfoOrNull() const {
     if (!hasLifetimeDependenceInfo()) {
-      return LifetimeDependenceInfo();
+      return nullptr;
     }
     auto *info = getTrailingObjects<LifetimeDependenceInfo>();
     assert(!info->empty() && "If the LifetimeDependenceInfo was empty, we "
                              "shouldn't have stored it.");
-    return *info;
+    return info;
   }
 
   void Profile(llvm::FoldingSetNodeID &ID) {
@@ -3914,14 +3929,22 @@ class GenericFunctionType final
     return getTrailingObjects<Type>()[hasGlobalActor()];
   }
 
-  LifetimeDependenceInfo getLifetimeDependenceInfo() const {
+  inline LifetimeDependenceInfo getLifetimeDependenceInfo() const {
+    if (auto *depInfo = getLifetimeDependenceInfoOrNull()) {
+      return *depInfo;
+    }
+    return LifetimeDependenceInfo();
+  }
+
+  /// Returns nullptr for an empty dependence list.
+  const LifetimeDependenceInfo *getLifetimeDependenceInfoOrNull() const {
     if (!hasLifetimeDependenceInfo()) {
-      return LifetimeDependenceInfo();
+      return nullptr;
     }
     auto *info = getTrailingObjects<LifetimeDependenceInfo>();
     assert(!info->empty() && "If the LifetimeDependenceInfo was empty, we "
                              "shouldn't have stored it.");
-    return *info;
+    return info;
   }
 
   /// Retrieve the generic signature of this function type.

lib/AST/Type.cpp

@@ -3896,12 +3896,13 @@ Type AnyFunctionType::getGlobalActor() const {
   }
 }
 
-LifetimeDependenceInfo AnyFunctionType::getLifetimeDependenceInfo() const {
+const LifetimeDependenceInfo *
+AnyFunctionType::getLifetimeDependenceInfoOrNull() const {
   switch (getKind()) {
   case TypeKind::Function:
-    return cast<FunctionType>(this)->getLifetimeDependenceInfo();
+    return cast<FunctionType>(this)->getLifetimeDependenceInfoOrNull();
   case TypeKind::GenericFunction:
-    return cast<GenericFunctionType>(this)->getLifetimeDependenceInfo();
+    return cast<GenericFunctionType>(this)->getLifetimeDependenceInfoOrNull();
 
   default:
     llvm_unreachable("Illegal type kind for AnyFunctionType.");

lib/Sema/LifetimeDependence.cpp

@@ -103,6 +103,24 @@ void LifetimeDependenceInfo::getConcatenatedData(
   }
 }
 
+static bool hasEscapableResultOrYield(AbstractFunctionDecl *afd,
+                                      Type resultType) {
+  Type resultTypeInContext = afd->mapTypeIntoContext(resultType);
+  std::optional<Type> yieldTyInContext;
+
+  if (auto *accessor = dyn_cast<AccessorDecl>(afd)) {
+    if (accessor->isCoroutine()) {
+      yieldTyInContext = accessor->getStorage()->getValueInterfaceType();
+      yieldTyInContext = accessor->mapTypeIntoContext(*yieldTyInContext);
+    }
+  }
+  if (resultTypeInContext->isEscapable() &&
+      (!yieldTyInContext.has_value() || (*yieldTyInContext)->isEscapable())) {
+    return true;
+  }
+  return false;
+}
+
 std::optional<LifetimeDependenceInfo>
 LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
                                      bool allowIndex) {
@@ -114,6 +132,12 @@ LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
   auto lifetimeDependentRepr =
       cast<LifetimeDependentReturnTypeRepr>(afd->getResultTypeRepr());
 
+  if (hasEscapableResultOrYield(afd, resultType)) {
+    diags.diagnose(lifetimeDependentRepr->getLoc(),
+                   diag::lifetime_dependence_invalid_return_type);
+    return std::nullopt;
+  }
+
   SmallBitVector inheritLifetimeParamIndices(capacity);
   SmallBitVector scopeLifetimeParamIndices(capacity);
 
@@ -122,11 +146,6 @@ LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
                                           ValueOwnership ownership) {
     auto loc = specifier.getLoc();
     auto kind = specifier.getLifetimeDependenceKind();
-    Type resultTypeInContext = afd->mapTypeIntoContext(resultType);
-    if (resultTypeInContext->isEscapable()) {
-      diags.diagnose(loc, diag::lifetime_dependence_invalid_return_type);
-      return true;
-    }
 
     if (ownership == ValueOwnership::Default) {
       diags.diagnose(loc, diag::lifetime_dependence_missing_ownership_modifier);
@@ -253,7 +272,8 @@ LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
           : nullptr,
       scopeLifetimeParamIndices.any()
           ? IndexSubset::get(ctx, scopeLifetimeParamIndices)
-          : nullptr);
+          : nullptr,
+      /*isExplicit*/ true);
 }
 
 std::optional<LifetimeDependenceInfo>
@@ -274,19 +294,11 @@ LifetimeDependenceInfo::infer(AbstractFunctionDecl *afd, Type resultType) {
   auto &diags = ctx.Diags;
   auto returnTypeRepr = afd->getResultTypeRepr();
   auto returnLoc = returnTypeRepr ? returnTypeRepr->getLoc() : afd->getLoc();
-  Type returnTyInContext = afd->mapTypeIntoContext(resultType);
-  std::optional<Type> yieldTyInContext;
 
-  if (auto *accessor = dyn_cast<AccessorDecl>(afd)) {
-    if (accessor->isCoroutine()) {
-      yieldTyInContext = accessor->getStorage()->getValueInterfaceType();
-      yieldTyInContext = accessor->mapTypeIntoContext(*yieldTyInContext);
-    }
-  }
-  if (returnTyInContext->isEscapable() &&
-      (!yieldTyInContext.has_value() || (*yieldTyInContext)->isEscapable())) {
+  if (hasEscapableResultOrYield(afd, resultType)) {
     return std::nullopt;
   }
+
   if (afd->getAttrs().hasAttribute<UnsafeNonEscapableResultAttr>()) {
     return std::nullopt;
   }

lib/Serialization/Deserialization.cpp

@@ -3590,7 +3590,8 @@ class DeclDeserializer {
     ctor->setParameters(bodyParams);
 
     SmallVector<LifetimeDependenceSpecifier> specifierList;
-    if (MF.maybeReadLifetimeDependence(specifierList, bodyParams->size())) {
+    if (MF.maybeReadLifetimeDependenceSpecifier(specifierList,
+                                                bodyParams->size())) {
       auto SelfType = ctor->getDeclaredInterfaceType();
       auto typeRepr = new (ctx) FixedTypeRepr(SelfType, SourceLoc());
       auto lifetimeTypeRepr =
@@ -4159,7 +4160,8 @@ class DeclDeserializer {
     ParameterList *paramList = MF.readParameterList();
     fn->setParameters(paramList);
     SmallVector<LifetimeDependenceSpecifier> specifierList;
-    if (MF.maybeReadLifetimeDependence(specifierList, paramList->size())) {
+    if (MF.maybeReadLifetimeDependenceSpecifier(specifierList,
+                                                paramList->size())) {
       auto typeRepr = new (ctx) FixedTypeRepr(resultType, SourceLoc());
       auto lifetimeTypeRepr =
           LifetimeDependentReturnTypeRepr::create(ctx, typeRepr, specifierList);
@@ -6851,7 +6853,6 @@ detail::function_deserializer::deserialize(ModuleFile &MF,
     isolation = swift::FunctionTypeIsolation::forGlobalActor(globalActorTy.get());
   }
 
-  // TODO: Handle LifetimeDependenceInfo here.
   auto info = FunctionType::ExtInfoBuilder(
                   *representation, noescape, throws, thrownError, *diffKind,
                   clangFunctionType, isolation, LifetimeDependenceInfo(),
@@ -6866,6 +6867,7 @@ detail::function_deserializer::deserialize(ModuleFile &MF,
 
   SmallVector<AnyFunctionType::Param, 8> params;
   while (true) {
+    BCOffsetRAII restoreOffset(MF.DeclTypeCursor);
     llvm::BitstreamEntry entry =
         MF.fatalIfUnexpected(MF.DeclTypeCursor.advance(AF_DontPopBlockAtEnd));
     if (entry.Kind != llvm::BitstreamEntry::Record)
@@ -6877,6 +6879,8 @@ detail::function_deserializer::deserialize(ModuleFile &MF,
     if (recordID != decls_block::FUNCTION_PARAM)
       break;
 
+    restoreOffset.reset();
+
     IdentifierID labelID;
     IdentifierID internalLabelID;
     TypeID typeID;
@@ -6907,6 +6911,13 @@ detail::function_deserializer::deserialize(ModuleFile &MF,
                         MF.getIdentifier(internalLabelID));
   }
 
+  auto lifetimeDependenceInfo =
+      MF.maybeReadLifetimeDependenceInfo(params.size());
+
+  if (lifetimeDependenceInfo.has_value()) {
+    info = info.withLifetimeDependenceInfo(*lifetimeDependenceInfo);
+  }
+
   if (!isGeneric) {
     assert(genericSig.isNull());
     return FunctionType::get(params, resultTy.get(), info);
@@ -7384,7 +7395,6 @@ Expected<Type> DESERIALIZE_TYPE(SIL_FUNCTION_TYPE)(
   if (erasedIsolation)
     isolation = SILFunctionTypeIsolation::Erased;
 
-  // Handle LifetimeDependenceInfo here.
   auto extInfo =
       SILFunctionType::ExtInfoBuilder(
           *representation, pseudogeneric, noescape, concurrent, async,
@@ -7527,6 +7537,13 @@ Expected<Type> DESERIALIZE_TYPE(SIL_FUNCTION_TYPE)(
   if (!patternSubsOrErr)
     return patternSubsOrErr.takeError();
 
+  auto lifetimeDependenceInfo = MF.maybeReadLifetimeDependenceInfo(
+      extInfo.hasSelfParam() ? numParams : numParams + 1);
+
+  if (lifetimeDependenceInfo.has_value()) {
+    extInfo = extInfo.withLifetimeDependenceInfo(*lifetimeDependenceInfo);
+  }
+
   return SILFunctionType::get(invocationSig, extInfo, coroutineKind.value(),
                               calleeConvention.value(), allParams, allYields,
                               allResults, errorResult,
@@ -8687,11 +8704,9 @@ ModuleFile::maybeReadForeignAsyncConvention() {
       errorFlagPolarity);
 }
 
-bool ModuleFile::maybeReadLifetimeDependence(
-    SmallVectorImpl<LifetimeDependenceSpecifier> &specifierList,
-    unsigned numParams) {
+bool ModuleFile::maybeReadLifetimeDependenceRecord(
+    SmallVectorImpl<uint64_t> &scratch) {
   using namespace decls_block;
-  SmallVector<uint64_t, 8> scratch;
 
   BCOffsetRAII restoreOffset(DeclTypeCursor);
 
@@ -8711,22 +8726,82 @@ bool ModuleFile::maybeReadLifetimeDependence(
     return false;
   }
 
-  bool hasInheritLifetimeParamIndices, hasScopeLifetimeParamIndices;
+  return true;
+}
+
+std::optional<LifetimeDependenceInfo>
+ModuleFile::maybeReadLifetimeDependenceInfo(unsigned numParams) {
+  using namespace decls_block;
+
+  SmallVector<uint64_t, 8> scratch;
+  if (!maybeReadLifetimeDependenceRecord(scratch)) {
+    return std::nullopt;
+  }
+
+  bool hasInheritLifetimeParamIndices;
+  bool hasScopeLifetimeParamIndices;
+  ArrayRef<uint64_t> lifetimeDependenceData;
+  LifetimeDependenceLayout::readRecord(scratch, hasInheritLifetimeParamIndices,
+                                       hasScopeLifetimeParamIndices,
+                                       lifetimeDependenceData);
+
+  SmallBitVector inheritLifetimeParamIndices(numParams, false);
+  SmallBitVector scopeLifetimeParamIndices(numParams, false);
+
+  unsigned startIndex = 0;
+  auto pushData = [&](SmallBitVector &bits) {
+    for (unsigned i = 0; i < numParams; i++) {
+      if (lifetimeDependenceData[startIndex + i]) {
+        bits.set(i);
+      }
+    }
+    startIndex += numParams;
+  };
+
+  if (hasInheritLifetimeParamIndices) {
+    pushData(inheritLifetimeParamIndices);
+  }
+  if (hasScopeLifetimeParamIndices) {
+    pushData(scopeLifetimeParamIndices);
+  }
+
+  ASTContext &ctx = getContext();
+  return LifetimeDependenceInfo(
+      hasInheritLifetimeParamIndices
+          ? IndexSubset::get(ctx, inheritLifetimeParamIndices)
+          : nullptr,
+      hasScopeLifetimeParamIndices
+          ? IndexSubset::get(ctx, scopeLifetimeParamIndices)
+          : nullptr);
+}
+
+bool ModuleFile::maybeReadLifetimeDependenceSpecifier(
+    SmallVectorImpl<LifetimeDependenceSpecifier> &specifierList,
+    unsigned numDeclParams) {
+  using namespace decls_block;
+
+  SmallVector<uint64_t, 8> scratch;
+  if (!maybeReadLifetimeDependenceRecord(scratch)) {
+    return false;
+  }
+
+  bool hasInheritLifetimeParamIndices;
+  bool hasScopeLifetimeParamIndices;
   ArrayRef<uint64_t> lifetimeDependenceData;
   LifetimeDependenceLayout::readRecord(scratch, hasInheritLifetimeParamIndices,
                                        hasScopeLifetimeParamIndices,
                                        lifetimeDependenceData);
 
   unsigned startIndex = 0;
   auto pushData = [&](LifetimeDependenceKind kind) {
-    for (unsigned i = 0; i < numParams + 1; i++) {
+    for (unsigned i = 0; i < numDeclParams + 1; i++) {
       if (lifetimeDependenceData[startIndex + i]) {
         specifierList.push_back(
             LifetimeDependenceSpecifier::getOrderedLifetimeDependenceSpecifier(
                 SourceLoc(), kind, i));
       }
     }
-    startIndex += numParams + 1;
+    startIndex += numDeclParams + 1;
   };
 
   if (hasInheritLifetimeParamIndices) {