Merge pull request #71295 from meg-gupta/lifetimedepserialization

meg-gupta · web-flow · commit 5aaa17f785f0 · 2024-02-01T10:54:38.000-08:00
Serialize/deserialize lifetime dependence when present
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -7797,6 +7797,9 @@ class FuncDecl : public AbstractFunctionDecl {
   SourceRange getSourceRange() const;
 
   TypeRepr *getResultTypeRepr() const { return FnRetType.getTypeRepr(); }
+
+  void setDeserializedResultTypeLoc(TypeLoc ResultTyR);
+
   SourceRange getResultTypeSourceRange() const {
     return FnRetType.getSourceRange();
   }
diff --git a/include/swift/AST/LifetimeDependence.h b/include/swift/AST/LifetimeDependence.h
@@ -164,8 +164,19 @@ class LifetimeDependenceInfo {
            mutateLifetimeParamIndices == nullptr;
   }
 
+  bool hasInheritLifetimeParamIndices() const {
+    return inheritLifetimeParamIndices != nullptr;
+  }
+  bool hasBorrowLifetimeParamIndices() const {
+    return borrowLifetimeParamIndices != nullptr;
+  }
+  bool hasMutateLifetimeParamIndices() const {
+    return mutateLifetimeParamIndices != nullptr;
+  }
+
   std::string getString() const;
   void Profile(llvm::FoldingSetNodeID &ID) const;
+  void getConcatenatedData(SmallVectorImpl<bool> &concatenatedData) const;
 
   static llvm::Optional<LifetimeDependenceInfo>
   get(AbstractFunctionDecl *decl, Type resultType, bool allowIndex = false);
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -9976,6 +9976,10 @@ void FuncDecl::setResultInterfaceType(Type type) {
                                         std::move(type));
 }
 
+void FuncDecl::setDeserializedResultTypeLoc(TypeLoc ResultTyR) {
+  FnRetType = ResultTyR;
+}
+
 FuncDecl *FuncDecl::createImpl(ASTContext &Context,
                                SourceLoc StaticLoc,
                                StaticSpellingKind StaticSpelling,
diff --git a/lib/Sema/LifetimeDependence.cpp b/lib/Sema/LifetimeDependence.cpp
@@ -85,6 +85,33 @@ LifetimeDependenceInfo LifetimeDependenceInfo::getForParamIndex(
                                 /*mutateLifetimeParamIndices*/ indexSubset};
 }
 
+void LifetimeDependenceInfo::getConcatenatedData(
+    SmallVectorImpl<bool> &concatenatedData) const {
+  auto pushData = [&](IndexSubset *paramIndices) {
+    if (paramIndices == nullptr) {
+      return;
+    }
+    assert(!paramIndices->isEmpty());
+
+    for (unsigned i = 0; i < paramIndices->getCapacity(); i++) {
+      if (paramIndices->contains(i)) {
+        concatenatedData.push_back(true);
+        continue;
+      }
+      concatenatedData.push_back(false);
+    }
+  };
+  if (hasInheritLifetimeParamIndices()) {
+    pushData(inheritLifetimeParamIndices);
+  }
+  if (hasBorrowLifetimeParamIndices()) {
+    pushData(borrowLifetimeParamIndices);
+  }
+  if (hasMutateLifetimeParamIndices()) {
+    pushData(mutateLifetimeParamIndices);
+  }
+}
+
 llvm::Optional<LifetimeDependenceInfo>
 LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
                                      bool allowIndex) {
@@ -181,9 +208,19 @@ LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
                        paramIndex);
         return llvm::None;
       }
+      if (paramIndex == 0) {
+        if (!afd->hasImplicitSelfDecl()) {
+          diags.diagnose(specifier.getLoc(),
+                         diag::lifetime_dependence_invalid_self);
+          return llvm::None;
+        }
+      }
+      auto ownership =
+          paramIndex == 0
+              ? afd->getImplicitSelfDecl()->getValueOwnership()
+              : afd->getParameters()->get(paramIndex - 1)->getValueOwnership();
       if (updateLifetimeDependenceInfo(
-              specifier, /*paramIndexToSet*/ specifier.getIndex() + 1,
-              afd->getParameters()->get(paramIndex)->getValueOwnership())) {
+              specifier, /*paramIndexToSet*/ specifier.getIndex(), ownership)) {
         return llvm::None;
       }
       break;
@@ -205,9 +242,15 @@ LifetimeDependenceInfo::fromTypeRepr(AbstractFunctionDecl *afd, Type resultType,
   }
 
   return LifetimeDependenceInfo(
-      IndexSubset::get(ctx, inheritLifetimeParamIndices),
-      IndexSubset::get(ctx, borrowLifetimeParamIndices),
-      IndexSubset::get(ctx, mutateLifetimeParamIndices));
+      inheritLifetimeParamIndices.any()
+          ? IndexSubset::get(ctx, inheritLifetimeParamIndices)
+          : nullptr,
+      borrowLifetimeParamIndices.any()
+          ? IndexSubset::get(ctx, borrowLifetimeParamIndices)
+          : nullptr,
+      mutateLifetimeParamIndices.any()
+          ? IndexSubset::get(ctx, mutateLifetimeParamIndices)
+          : nullptr);
 }
 
 llvm::Optional<LifetimeDependenceInfo>
diff --git a/lib/Serialization/DeclTypeRecordNodes.def b/lib/Serialization/DeclTypeRecordNodes.def
@@ -212,6 +212,8 @@ OTHER(ERROR_FLAG, 155)
 TRAILING_INFO(CONDITIONAL_SUBSTITUTION)
 TRAILING_INFO(CONDITIONAL_SUBSTITUTION_COND)
 
+OTHER(LIFETIME_DEPENDENCE, 158)
+
 #ifndef DECL_ATTR
 #define DECL_ATTR(NAME, CLASS, OPTIONS, CODE) RECORD_VAL(CLASS##_DECL_ATTR, 180+CODE)
 #endif
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
@@ -4105,6 +4105,13 @@ class DeclDeserializer {
 
     ParameterList *paramList = MF.readParameterList();
     fn->setParameters(paramList);
+    SmallVector<LifetimeDependenceSpecifier> specifierList;
+    if (MF.maybeReadLifetimeDependence(specifierList, paramList->size())) {
+      auto typeRepr = new (ctx) FixedTypeRepr(resultType, SourceLoc());
+      auto lifetimeTypeRepr =
+          LifetimeDependentReturnTypeRepr::create(ctx, typeRepr, specifierList);
+      fn->setDeserializedResultTypeLoc(TypeLoc(lifetimeTypeRepr, resultType));
+    }
 
     if (auto errorConvention = MF.maybeReadForeignErrorConvention())
       fn->setForeignErrorConvention(*errorConvention);
@@ -8548,3 +8555,58 @@ ModuleFile::maybeReadForeignAsyncConvention() {
       completionHandlerErrorFlagParamIndex,
       errorFlagPolarity);
 }
+
+bool ModuleFile::maybeReadLifetimeDependence(
+    SmallVectorImpl<LifetimeDependenceSpecifier> &specifierList,
+    unsigned numParams) {
+  using namespace decls_block;
+  SmallVector<uint64_t, 8> scratch;
+
+  BCOffsetRAII restoreOffset(DeclTypeCursor);
+
+  llvm::BitstreamEntry next =
+      fatalIfUnexpected(DeclTypeCursor.advance(AF_DontPopBlockAtEnd));
+  if (next.Kind != llvm::BitstreamEntry::Record)
+    return false;
+
+  unsigned recKind =
+      fatalIfUnexpected(DeclTypeCursor.readRecord(next.ID, scratch));
+  switch (recKind) {
+  case LIFETIME_DEPENDENCE:
+    restoreOffset.reset();
+    break;
+
+  default:
+    return false;
+  }
+
+  bool hasInheritLifetimeParamIndices, hasBorrowLifetimeParamIndices,
+      hasMutateLifetimeParamIndices;
+  ArrayRef<uint64_t> lifetimeDependenceData;
+  LifetimeDependenceLayout::readRecord(
+      scratch, hasInheritLifetimeParamIndices, hasBorrowLifetimeParamIndices,
+      hasMutateLifetimeParamIndices, lifetimeDependenceData);
+
+  unsigned startIndex = 0;
+  auto pushData = [&](LifetimeDependenceKind kind) {
+    for (unsigned i = 0; i < numParams + 1; i++) {
+      if (lifetimeDependenceData[startIndex + i]) {
+        specifierList.push_back(
+            LifetimeDependenceSpecifier::getOrderedLifetimeDependenceSpecifier(
+                SourceLoc(), kind, i));
+      }
+    }
+    startIndex += numParams + 1;
+  };
+
+  if (hasInheritLifetimeParamIndices) {
+    pushData(LifetimeDependenceKind::Consume);
+  }
+  if (hasBorrowLifetimeParamIndices) {
+    pushData(LifetimeDependenceKind::Borrow);
+  }
+  if (hasMutateLifetimeParamIndices) {
+    pushData(LifetimeDependenceKind::Mutate);
+  }
+  return true;
+}
diff --git a/lib/Serialization/ModuleFile.h b/lib/Serialization/ModuleFile.h
@@ -1060,6 +1060,10 @@ class ModuleFile
   /// Reads a foreign async convention from \c DeclTypeCursor, if present.
   llvm::Optional<ForeignAsyncConvention> maybeReadForeignAsyncConvention();
 
+  bool maybeReadLifetimeDependence(
+      SmallVectorImpl<LifetimeDependenceSpecifier> &specifierList,
+      unsigned numParams);
+
   /// Reads inlinable body text from \c DeclTypeCursor, if present.
   llvm::Optional<StringRef> maybeReadInlinableBodyText();
 
diff --git a/lib/Serialization/ModuleFormat.h b/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 = 846; // always trigger mismatch for NoncopyableGenerics
+const uint16_t SWIFTMODULE_VERSION_MINOR = 847; // add LifetimeDependence
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -2174,6 +2174,14 @@ namespace decls_block {
     BCFixed<1>   // completion handler error flag polarity
   >;
 
+  using LifetimeDependenceLayout =
+      BCRecordLayout<LIFETIME_DEPENDENCE,
+                     BCFixed<1>,         // hasInheritLifetimeParamIndices
+                     BCFixed<1>,         // hasBorrowLifetimeParamIndices
+                     BCFixed<1>,         // hasMutateLifetimeParamIndices
+                     BCArray<BCFixed<1>> // concatenated param indices
+                     >;
+
   using AbstractClosureExprLayout = BCRecordLayout<
     ABSTRACT_CLOSURE_EXPR_CONTEXT,
     TypeIDField, // type
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -3509,6 +3509,20 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
                                              fac.completionHandlerFlagIsErrorOnZero());
   }
 
+  void
+  writeLifetimeDependenceInfo(LifetimeDependenceInfo lifetimeDependenceInfo) {
+    using namespace decls_block;
+    SmallVector<bool> paramIndices;
+    lifetimeDependenceInfo.getConcatenatedData(paramIndices);
+
+    auto abbrCode = S.DeclTypeAbbrCodes[LifetimeDependenceLayout::Code];
+    LifetimeDependenceLayout::emitRecord(
+        S.Out, S.ScratchRecord, abbrCode,
+        lifetimeDependenceInfo.hasInheritLifetimeParamIndices(),
+        lifetimeDependenceInfo.hasBorrowLifetimeParamIndices(),
+        lifetimeDependenceInfo.hasMutateLifetimeParamIndices(), paramIndices);
+  }
+
   void writeGenericParams(const GenericParamList *genericParams) {
     using namespace decls_block;
 
@@ -4482,6 +4496,12 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     // Write the body parameters.
     writeParameterList(fn->getParameters());
 
+    auto fnType = ty->getAs<FunctionType>();
+    if (fnType && fnType->hasLifetimeDependenceInfo()) {
+      assert(!fnType->getLifetimeDependenceInfo().empty());
+      writeLifetimeDependenceInfo(fnType->getLifetimeDependenceInfo());
+    }
+
     if (auto errorConvention = fn->getForeignErrorConvention())
       writeForeignErrorConvention(*errorConvention);
     if (auto asyncConvention = fn->getForeignAsyncConvention())
@@ -5949,6 +5969,7 @@ void Serializer::writeAllDeclsAndTypes() {
 
   registerDeclTypeAbbr<ForeignErrorConventionLayout>();
   registerDeclTypeAbbr<ForeignAsyncConventionLayout>();
+  registerDeclTypeAbbr<LifetimeDependenceLayout>();
   registerDeclTypeAbbr<AbstractClosureExprLayout>();
   registerDeclTypeAbbr<PatternBindingInitializerLayout>();
   registerDeclTypeAbbr<DefaultArgumentInitializerLayout>();
diff --git a/test/Parse/explicit_lifetime_dependence_specifiers.swift b/test/Parse/explicit_lifetime_dependence_specifiers.swift
@@ -93,7 +93,8 @@ func invalidSpecifier3(_ x: borrowing BufferView) -> _borrow(*) BufferView { //
   return BufferView(x.ptr)
 } 
 
-func invalidSpecifier4(_ x: borrowing BufferView) -> _borrow(0) BufferView {
+// TODO: Diagnose using param indices on func decls in sema
+func invalidSpecifier4(_ x: borrowing BufferView) -> _borrow(0) BufferView { // expected-error{{invalid lifetime dependence specifier, self is valid in non-static methods only}}
   return BufferView(x.ptr)
 }
 
diff --git a/test/Serialization/Inputs/def_explicit_lifetime_dependence.swift b/test/Serialization/Inputs/def_explicit_lifetime_dependence.swift
@@ -0,0 +1,37 @@
+public struct BufferView : ~Escapable {
+  let ptr: UnsafeRawBufferPointer
+  @_unsafeNonescapableResult
+  public init(_ ptr: UnsafeRawBufferPointer) {
+    self.ptr = ptr
+  }
+}
+
+public struct MutableBufferView : ~Escapable, ~Copyable {
+  let ptr: UnsafeMutableRawBufferPointer
+  @_unsafeNonescapableResult
+  public init(_ ptr: UnsafeMutableRawBufferPointer) {
+    self.ptr = ptr
+  }
+}
+
+public func derive(_ x: borrowing BufferView) -> _borrow(x) BufferView {
+  return BufferView(x.ptr)
+}
+
+public func use(_ x: borrowing BufferView) {}
+
+public func borrowAndCreate(_ view: borrowing BufferView) -> _borrow(view) BufferView {
+  return BufferView(view.ptr)
+}
+
+public func consumeAndCreate(_ view: consuming BufferView) -> _consume(view) BufferView {
+  return BufferView(view.ptr)
+}
+
+public func deriveThisOrThat(_ this: borrowing BufferView, _ that: borrowing BufferView) -> _borrow(this, that) BufferView {
+  if (Int.random(in: 1..<100) == 0) {
+    return BufferView(this.ptr)
+  }
+  return BufferView(that.ptr)
+}
+
diff --git a/test/Serialization/Inputs/def_implicit_lifetime_dependence.swift b/test/Serialization/Inputs/def_implicit_lifetime_dependence.swift
@@ -0,0 +1,30 @@
+public struct BufferView : ~Escapable {
+  let ptr: UnsafeRawBufferPointer
+  @_unsafeNonescapableResult
+  public init(_ ptr: UnsafeRawBufferPointer) {
+    self.ptr = ptr
+  }
+}
+
+public struct MutableBufferView : ~Escapable, ~Copyable {
+  let ptr: UnsafeMutableRawBufferPointer
+  @_unsafeNonescapableResult
+  public init(_ ptr: UnsafeMutableRawBufferPointer) {
+    self.ptr = ptr
+  }
+}
+
+public func derive(_ x: borrowing BufferView) -> BufferView {
+  return BufferView(x.ptr)
+}
+
+public func use(_ x: borrowing BufferView) {}
+
+public func borrowAndCreate(_ view: borrowing BufferView) -> BufferView {
+  return BufferView(view.ptr)
+}
+
+public func consumeAndCreate(_ view: consuming BufferView) -> BufferView {
+  return BufferView(view.ptr)
+}
+
diff --git a/test/Serialization/explicit_lifetime_dependence.swift b/test/Serialization/explicit_lifetime_dependence.swift
@@ -0,0 +1,34 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -o %t  %S/Inputs/def_explicit_lifetime_dependence.swift \
+// RUN: -enable-experimental-feature NonescapableTypes \
+// RUN: -disable-experimental-parser-round-trip \
+// RUN: -enable-experimental-feature NoncopyableGenerics
+
+// RUN: llvm-bcanalyzer %t/def_explicit_lifetime_dependence.swiftmodule 
+
+// RUN: %target-swift-frontend -module-name lifetime-dependence -emit-silgen -I %t %s \
+// RUN: -enable-experimental-feature NonescapableTypes \
+// RUN: -disable-experimental-parser-round-trip \
+// RUN: -enable-experimental-feature NoncopyableGenerics | %FileCheck %s
+
+// REQUIRES: noncopyable_generics
+
+import def_explicit_lifetime_dependence
+
+func testBasic() {
+  let capacity = 4
+  let a = Array(0..<capacity)
+  a.withUnsafeBytes {
+    let view = BufferView($0)
+    let derivedView = derive(view)
+    let consumedView = consumeAndCreate(derivedView)
+    let borrowedView = borrowAndCreate(consumedView) 
+    let mysteryView = deriveThisOrThat(borrowedView, consumedView)
+    use(mysteryView)
+  }
+}
+
+// CHECK: sil @$s32def_explicit_lifetime_dependence6deriveyAA10BufferViewVADF : $@convention(thin) (@guaranteed BufferView) -> _borrow(1) @owned BufferView
+// CHECK: sil @$s32def_explicit_lifetime_dependence16consumeAndCreateyAA10BufferViewVADnF : $@convention(thin) (@owned BufferView) -> _inherit(1) @owned BufferView
+// CHECK: sil @$s32def_explicit_lifetime_dependence15borrowAndCreateyAA10BufferViewVADF : $@convention(thin) (@guaranteed BufferView) -> _borrow(1) @owned BufferView
+// CHECK: sil @$s32def_explicit_lifetime_dependence16deriveThisOrThatyAA10BufferViewVAD_ADtF : $@convention(thin) (@guaranteed BufferView, @guaranteed BufferView) -> _borrow(1, 2) @owned BufferView
diff --git a/test/Serialization/implicit_lifetime_dependence.swift b/test/Serialization/implicit_lifetime_dependence.swift

Original file line number	Diff line number	Diff line change
`@@ -7797,6 +7797,9 @@ class FuncDecl : public AbstractFunctionDecl {`
`7797`	`7797`	`SourceRange getSourceRange() const;`
`7798`	`7798`
`7799`	`7799`	`TypeRepr *getResultTypeRepr() const { return FnRetType.getTypeRepr(); }`
	`7800`	`+`
	`7801`	`+ void setDeserializedResultTypeLoc(TypeLoc ResultTyR);`
	`7802`	`+`
`7800`	`7803`	`SourceRange getResultTypeSourceRange() const {`
`7801`	`7804`	`return FnRetType.getSourceRange();`
`7802`	`7805`	`}`
Original file line number	Diff line number	Diff line change
`@@ -93,7 +93,8 @@ func invalidSpecifier3(_ x: borrowing BufferView) -> _borrow(*) BufferView { //`
`93`	`93`	`return BufferView(x.ptr)`
`94`	`94`	`}`
`95`	`95`
`96`		`-func invalidSpecifier4(_ x: borrowing BufferView) -> _borrow(0) BufferView {`
	`96`	`+// TODO: Diagnose using param indices on func decls in sema`
	`97`	`+func invalidSpecifier4(_ x: borrowing BufferView) -> _borrow(0) BufferView { // expected-error{{invalid lifetime dependence specifier, self is valid in non-static methods only}}`
`97`	`98`	`return BufferView(x.ptr)`
`98`	`99`	`}`
`99`	`100`