Merge pull request swiftlang#81879 from meg-gupta/fixdeserialization

Fix deserialization of lifetime dependencies on ast function types

Author: meg-gupta

include/swift/AST/LifetimeDependence.h

@@ -227,6 +227,25 @@ class LifetimeDependenceInfo {
             || !conditionallyAddressableParamIndices
             || conditionallyAddressableParamIndices->isDisjointWith(
               addressableParamIndices)));
+
+    if (CONDITIONAL_ASSERT_enabled()) {
+      // Ensure inherit/scope/addressable param indices are of the same length
+      // or 0.
+      unsigned paramIndicesLength = 0;
+      if (inheritLifetimeParamIndices) {
+        paramIndicesLength = inheritLifetimeParamIndices->getCapacity();
+      }
+      if (scopeLifetimeParamIndices) {
+        assert(paramIndicesLength == 0 ||
+               paramIndicesLength == scopeLifetimeParamIndices->getCapacity());
+        paramIndicesLength = scopeLifetimeParamIndices->getCapacity();
+      }
+      if (addressableParamIndices) {
+        assert(paramIndicesLength == 0 ||
+               paramIndicesLength == addressableParamIndices->getCapacity());
+        paramIndicesLength = addressableParamIndices->getCapacity();
+      }
+    }
   }
 
   operator bool() const { return !empty(); }
@@ -250,6 +269,19 @@ class LifetimeDependenceInfo {
     return addressableParamIndicesAndImmortal.getPointer() != nullptr;
   }
 
+  unsigned getParamIndicesLength() const {
+    if (hasInheritLifetimeParamIndices()) {
+      return getInheritIndices()->getCapacity();
+    }
+    if (hasScopeLifetimeParamIndices()) {
+      return getScopeIndices()->getCapacity();
+    }
+    if (hasAddressableParamIndices()) {
+      return getAddressableIndices()->getCapacity();
+    }
+    return 0;
+  }
+
   IndexSubset *getInheritIndices() const { return inheritLifetimeParamIndices; }
 
   IndexSubset *getScopeIndices() const { return scopeLifetimeParamIndices; }

lib/Serialization/Deserialization.cpp

@@ -3891,7 +3891,7 @@ class DeclDeserializer {
     ctor->setParameters(bodyParams);
 
     SmallVector<LifetimeDependenceInfo, 1> lifetimeDependencies;
-    while (auto info = MF.maybeReadLifetimeDependence(bodyParams->size() + 1)) {
+    while (auto info = MF.maybeReadLifetimeDependence()) {
       assert(info.has_value());
       lifetimeDependencies.push_back(*info);
     }
@@ -4493,11 +4493,9 @@ class DeclDeserializer {
     ParameterList *paramList;
     SET_OR_RETURN_ERROR(paramList, MF.readParameterList());
     fn->setParameters(paramList);
-    auto numParams =
-        fn->hasImplicitSelfDecl() ? paramList->size() + 1 : paramList->size();
 
     SmallVector<LifetimeDependenceInfo, 1> lifetimeDependencies;
-    while (auto info = MF.maybeReadLifetimeDependence(numParams)) {
+    while (auto info = MF.maybeReadLifetimeDependence()) {
       assert(info.has_value());
       lifetimeDependencies.push_back(*info);
     }
@@ -7438,8 +7436,7 @@ detail::function_deserializer::deserialize(ModuleFile &MF,
 
   SmallVector<LifetimeDependenceInfo, 1> lifetimeDependencies;
 
-  while (auto lifetimeDependence =
-             MF.maybeReadLifetimeDependence(params.size())) {
+  while (auto lifetimeDependence = MF.maybeReadLifetimeDependence()) {
     lifetimeDependencies.push_back(*lifetimeDependence);
   }
   if (!lifetimeDependencies.empty()) {
@@ -8091,7 +8088,7 @@ Expected<Type> DESERIALIZE_TYPE(SIL_FUNCTION_TYPE)(
 
   SmallVector<LifetimeDependenceInfo, 1> lifetimeDependencies;
 
-  while (auto lifetimeDependence = MF.maybeReadLifetimeDependence(numParams)) {
+  while (auto lifetimeDependence = MF.maybeReadLifetimeDependence()) {
     lifetimeDependencies.push_back(*lifetimeDependence);
   }
 
@@ -9378,7 +9375,7 @@ bool ModuleFile::maybeReadLifetimeDependenceRecord(
 }
 
 std::optional<LifetimeDependenceInfo>
-ModuleFile::maybeReadLifetimeDependence(unsigned numParams) {
+ModuleFile::maybeReadLifetimeDependence() {
   using namespace decls_block;
 
   SmallVector<uint64_t, 8> scratch;
@@ -9387,28 +9384,29 @@ ModuleFile::maybeReadLifetimeDependence(unsigned numParams) {
   }
 
   unsigned targetIndex;
+  unsigned paramIndicesLength;
   bool isImmortal;
   bool hasInheritLifetimeParamIndices;
   bool hasScopeLifetimeParamIndices;
   bool hasAddressableParamIndices;
   ArrayRef<uint64_t> lifetimeDependenceData;
   LifetimeDependenceLayout::readRecord(
-      scratch, targetIndex, isImmortal, hasInheritLifetimeParamIndices,
-      hasScopeLifetimeParamIndices, hasAddressableParamIndices,
-      lifetimeDependenceData);
+      scratch, targetIndex, paramIndicesLength, isImmortal,
+      hasInheritLifetimeParamIndices, hasScopeLifetimeParamIndices,
+      hasAddressableParamIndices, lifetimeDependenceData);
 
-  SmallBitVector inheritLifetimeParamIndices(numParams, false);
-  SmallBitVector scopeLifetimeParamIndices(numParams, false);
-  SmallBitVector addressableParamIndices(numParams, false);
+  SmallBitVector inheritLifetimeParamIndices(paramIndicesLength, false);
+  SmallBitVector scopeLifetimeParamIndices(paramIndicesLength, false);
+  SmallBitVector addressableParamIndices(paramIndicesLength, false);
 
   unsigned startIndex = 0;
   auto pushData = [&](SmallBitVector &bits) {
-    for (unsigned i = 0; i < numParams; i++) {
+    for (unsigned i = 0; i < paramIndicesLength; i++) {
       if (lifetimeDependenceData[startIndex + i]) {
         bits.set(i);
       }
     }
-    startIndex += numParams;
+    startIndex += paramIndicesLength;
   };
 
   if (hasInheritLifetimeParamIndices) {

lib/Serialization/ModuleFile.h

@@ -1100,8 +1100,7 @@ class ModuleFile
   bool maybeReadLifetimeDependenceRecord(SmallVectorImpl<uint64_t> &scratch);
 
   // Reads lifetime dependence info from type if present.
-  std::optional<LifetimeDependenceInfo>
-  maybeReadLifetimeDependence(unsigned numParams);
+  std::optional<LifetimeDependenceInfo> maybeReadLifetimeDependence();
 
   // Reads lifetime dependence specifier from decl if present
   bool maybeReadLifetimeEntry(SmallVectorImpl<LifetimeEntry> &specifierList,

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,7 @@ 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 = 952; // Add @specialized
+const uint16_t SWIFTMODULE_VERSION_MINOR = 953; // update LifetimeDependence layout
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -2323,6 +2323,7 @@ namespace decls_block {
   using LifetimeDependenceLayout =
       BCRecordLayout<LIFETIME_DEPENDENCE,
                      BCVBR<4>,           // targetIndex
+                     BCVBR<4>,           // paramIndicesLength
                      BCFixed<1>,         // isImmortal
                      BCFixed<1>,         // hasInheritLifetimeParamIndices
                      BCFixed<1>,         // hasScopeLifetimeParamIndices

lib/Serialization/Serialization.cpp

@@ -2734,10 +2734,10 @@ void Serializer::writeLifetimeDependencies(
 
     auto abbrCode = DeclTypeAbbrCodes[LifetimeDependenceLayout::Code];
     LifetimeDependenceLayout::emitRecord(
-        Out, ScratchRecord, abbrCode, info.getTargetIndex(), info.isImmortal(),
+        Out, ScratchRecord, abbrCode, info.getTargetIndex(),
+        info.getParamIndicesLength(), info.isImmortal(),
         info.hasInheritLifetimeParamIndices(),
-        info.hasScopeLifetimeParamIndices(),
-        info.hasAddressableParamIndices(),
+        info.hasScopeLifetimeParamIndices(), info.hasAddressableParamIndices(),
         paramIndices);
     paramIndices.clear();
   }

validation-test/Serialization/rdar151768216.swift

@@ -0,0 +1,18 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module %s -o %t -enable-experimental-feature LifetimeDependence
+
+// RUN: %target-swift-frontend -emit-sil %t/rdar151768216.swiftmodule \
+// RUN: -enable-experimental-feature LifetimeDependence
+
+// REQUIRES: swift_feature_LifetimeDependence
+
+// Ensure no crash
+extension Result {   
+  @inlinable
+  func castError(i: Int, j: Int, k: Int) -> Result<Success, Failure> {
+    return self.mapError { error in
+      return error
+    }   
+  }
+}
+