Serialize lifetime dependence info on function types as well

Author: meg-gupta

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

@@ -3781,6 +3781,9 @@ class FunctionType final
 
   /// Returns nullptr for an empty dependence list.
   const LifetimeDependenceInfo *getLifetimeDependenceInfoOrNull() const {
+    if (!hasLifetimeDependenceInfo()) {
+      return nullptr;
+    }
     auto *info = getTrailingObjects<LifetimeDependenceInfo>();
     assert(!info->empty() && "If the LifetimeDependenceInfo was empty, we "
                              "shouldn't have stored it.");
@@ -3935,6 +3938,9 @@ class GenericFunctionType final
 
   /// Returns nullptr for an empty dependence list.
   const LifetimeDependenceInfo *getLifetimeDependenceInfoOrNull() const {
+    if (!hasLifetimeDependenceInfo()) {
+      return nullptr;
+    }
     auto *info = getTrailingObjects<LifetimeDependenceInfo>();
     assert(!info->empty() && "If the LifetimeDependenceInfo was empty, we "
                              "shouldn't have stored it.");

lib/Sema/LifetimeDependence.cpp

@@ -272,7 +272,8 @@ LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
           : nullptr,
       scopeLifetimeParamIndices.any()
           ? IndexSubset::get(ctx, scopeLifetimeParamIndices)
-          : nullptr);
+          : nullptr,
+      /*isExplicit*/ true);
 }
 
 std::optional<LifetimeDependenceInfo>

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) {

lib/Serialization/ModuleFile.h

@@ -1067,9 +1067,16 @@ class ModuleFile
   /// Reads a foreign async convention from \c DeclTypeCursor, if present.
   std::optional<ForeignAsyncConvention> maybeReadForeignAsyncConvention();
 
-  bool maybeReadLifetimeDependence(
+  bool maybeReadLifetimeDependenceRecord(SmallVectorImpl<uint64_t> &scratch);
+
+  // Reads lifetime dependence info from type if present.
+  std::optional<LifetimeDependenceInfo>
+  maybeReadLifetimeDependenceInfo(unsigned numParams);
+
+  // Reads lifetime dependence specifier from decl if present
+  bool maybeReadLifetimeDependenceSpecifier(
       SmallVectorImpl<LifetimeDependenceSpecifier> &specifierList,
-      unsigned numParams);
+      unsigned numDeclParams);
 
   /// Reads inlinable body text from \c DeclTypeCursor, if present.
   std::optional<StringRef> maybeReadInlinableBodyText();

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,8 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 861; // AllowNonResilientAccess
+const uint16_t SWIFTMODULE_VERSION_MINOR =
+    862; // add lifetime dependence info in type serialization as well.
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -1245,6 +1246,7 @@ namespace decls_block {
     FunctionTypeIsolationField,      // isolation
     BCFixed<1>                       // has transferring result
     // trailed by parameters
+    // Optionally lifetime dependence info
   );
 
   using FunctionParamLayout =
@@ -1346,6 +1348,7 @@ namespace decls_block {
     GenericSignatureIDField          // generic signature
 
     // trailed by parameters
+    // Optionally lifetime dependence info
   );
 
   TYPE_LAYOUT(SILFunctionTypeLayout,
@@ -1374,6 +1377,7 @@ namespace decls_block {
                           // followed by error result type/convention
     // Optionally a protocol conformance (for witness_methods)
     // Optionally a substitution map (for substituted function types)
+    // Optionally lifetime dependence info
   );
 
   TYPE_LAYOUT(SILBlockStorageTypeLayout,