Store parsed @Lifetime attribute on the function type

meg-gupta · meg-gupta · commit 04076a5473b3 · 2024-09-04T00:10:51.000-07:00
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -2650,6 +2650,10 @@ class LifetimeAttr final
                               SourceRange baseRange, bool implicit,
                               ArrayRef<LifetimeEntry> entries);
 
+  ArrayRef<LifetimeEntry> getLifetimeEntries() const {
+    return {getTrailingObjects<LifetimeEntry>(), NumEntries};
+  }
+
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DeclAttrKind::Lifetime;
   }
diff --git a/include/swift/AST/LifetimeDependence.h b/include/swift/AST/LifetimeDependence.h
@@ -153,6 +153,10 @@ class LifetimeDependenceInfo {
                                             unsigned sourceIndex,
                                             LifetimeDependenceKind kind);
 
+  /// Builds LifetimeDependenceInfo from @lifetime attribute
+  static std::optional<ArrayRef<LifetimeDependenceInfo>>
+  fromLifetimeAttribute(AbstractFunctionDecl *afd);
+
   /// Builds LifetimeDependenceInfo from dependsOn type modifier
   static std::optional<LifetimeDependenceInfo>
   fromDependsOn(AbstractFunctionDecl *afd, TypeRepr *targetRepr,
diff --git a/lib/AST/Attr.cpp b/lib/AST/Attr.cpp
@@ -1812,6 +1812,20 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     break;
   }
 
+  case DeclAttrKind::Lifetime: {
+    auto *attr = cast<LifetimeAttr>(this);
+    bool firstElem = true;
+    Printer << "@lifetime(";
+    for (auto entry : attr->getLifetimeEntries()) {
+      if (!firstElem) {
+        Printer << ", ";
+      }
+      Printer << entry.getParamString();
+    }
+    Printer << ")";
+    break;
+  }
+
 #define SIMPLE_DECL_ATTR(X, CLASS, ...) case DeclAttrKind::CLASS:
 #include "swift/AST/DeclAttr.def"
     llvm_unreachable("handled above");
diff --git a/lib/AST/LifetimeDependence.cpp b/lib/AST/LifetimeDependence.cpp
@@ -218,6 +218,22 @@ getLifetimeDependenceKind(LifetimeEntry specifier, AbstractFunctionDecl *afd,
   auto ownership = decl->getValueOwnership();
   auto type = decl->getTypeInContext();
 
+  // For @lifetime attribute, we determine lifetime dependence kind based on
+  // type.
+  if (afd->getAttrs().hasAttribute<LifetimeAttr>()) {
+    auto lifetimeKind = getLifetimeDependenceKindFromType(type);
+    bool isCompatible = isLifetimeDependenceCompatibleWithOwnership(
+        lifetimeKind, type, ownership, afd);
+    if (!isCompatible) {
+      assert(lifetimeKind == LifetimeDependenceKind::Scope);
+      diags.diagnose(
+          loc, diag::lifetime_dependence_cannot_use_inferred_scoped_consuming);
+      return std::nullopt;
+    }
+  }
+
+  // For dependsOn type modifier, we check if we had a "scoped" modifier, if not
+  // we determine lifetime dependence kind based on type.
   auto parsedLifetimeKind = specifier.getParsedLifetimeDependenceKind();
   auto lifetimeKind =
       parsedLifetimeKind == ParsedLifetimeDependenceKind::Default
@@ -245,6 +261,10 @@ static bool populateLifetimeDependence(AbstractFunctionDecl *afd,
                                        SmallBitVector &inheritIndices,
                                        SmallBitVector &scopeIndices,
                                        bool &isImmortal) {
+  auto *dc = afd->getDeclContext();
+  auto &ctx = dc->getASTContext();
+  auto &diags = ctx.Diags;
+
   auto updateLifetimeIndices =
       [&](LifetimeEntry entry, unsigned paramIndexToSet,
           std::optional<LifetimeDependenceKind> lifetimeKind) {
@@ -278,7 +298,7 @@ static bool populateLifetimeDependence(AbstractFunctionDecl *afd,
                      diag::lifetime_dependence_immortal_conflict_name);
       return true;
     }
-    if (inheritIndicies.any() || scopeIndices.any()) {
+    if (inheritIndices.any() || scopeIndices.any()) {
       diags.diagnose(entry.getLoc(), diag::lifetime_dependence_immortal_alone);
       return true;
     }
@@ -346,6 +366,45 @@ static bool populateLifetimeDependence(AbstractFunctionDecl *afd,
   }
 }
 
+std::optional<ArrayRef<LifetimeDependenceInfo>>
+LifetimeDependenceInfo::fromLifetimeAttribute(AbstractFunctionDecl *afd) {
+  auto *dc = afd->getDeclContext();
+  auto &ctx = dc->getASTContext();
+
+  auto lifetimeAttrs = afd->getAttrs().getAttributes<LifetimeAttr>();
+
+  auto capacity = afd->hasImplicitSelfDecl()
+                      ? (afd->getParameters()->size() + 1)
+                      : afd->getParameters()->size();
+
+  SmallBitVector inheritIndices(capacity);
+  SmallBitVector scopeIndices(capacity);
+  bool isImmortal = false;
+
+  bool hasError = false;
+  for (auto attr : lifetimeAttrs) {
+    for (auto entry : attr->getLifetimeEntries()) {
+      hasError |= populateLifetimeDependence(afd, entry, inheritIndices,
+                                             scopeIndices, isImmortal);
+    }
+  }
+
+  if (hasError) {
+    return std::nullopt;
+  }
+  // TODO: Handle lifetime dependencies for targets other than function result.
+  SmallVector<LifetimeDependenceInfo> lifetimeDependencies;
+  auto resultIndex = afd->hasImplicitSelfDecl()
+                         ? afd->getParameters()->size() + 1
+                         : afd->getParameters()->size();
+  auto resultDependence = LifetimeDependenceInfo(
+      inheritIndices.any() ? IndexSubset::get(ctx, inheritIndices) : nullptr,
+      scopeIndices.any() ? IndexSubset::get(ctx, scopeIndices) : nullptr,
+      resultIndex, isImmortal);
+  lifetimeDependencies.push_back(resultDependence);
+  return afd->getASTContext().AllocateCopy(lifetimeDependencies);
+}
+
 std::optional<LifetimeDependenceInfo>
 LifetimeDependenceInfo::fromDependsOn(AbstractFunctionDecl *afd,
                                       TypeRepr *targetTypeRepr, Type targetType,
@@ -660,6 +719,10 @@ LifetimeDependenceInfo::get(AbstractFunctionDecl *afd) {
   }
   assert(isa<FuncDecl>(afd) || isa<ConstructorDecl>(afd));
 
+  if (afd->getAttrs().hasAttribute<LifetimeAttr>()) {
+    return LifetimeDependenceInfo::fromLifetimeAttribute(afd);
+  }
+
   SmallVector<LifetimeDependenceInfo> lifetimeDependencies;
 
   for (unsigned targetIndex : indices(*afd->getParameters())) {
@@ -672,20 +735,26 @@ LifetimeDependenceInfo::get(AbstractFunctionDecl *afd) {
     }
   }
 
-  std::optional<LifetimeDependenceInfo> resultDependence =
-      LifetimeDependenceInfo::fromDependsOn(
-          afd, afd->getResultTypeRepr(), getResultOrYield(afd),
-          afd->hasImplicitSelfDecl() ? afd->getParameters()->size() + 1
-                                     : afd->getParameters()->size());
+  std::optional<LifetimeDependenceInfo> resultDependence;
 
-  if (!resultDependence.has_value()) {
+  if (auto *lifetimeTypeRepr = dyn_cast_or_null<LifetimeDependentTypeRepr>(
+          afd->getResultTypeRepr())) {
+    resultDependence = LifetimeDependenceInfo::fromDependsOn(
+        afd, lifetimeTypeRepr, getResultOrYield(afd),
+        afd->hasImplicitSelfDecl() ? afd->getParameters()->size() + 1
+                                   : afd->getParameters()->size());
+  } else {
     resultDependence = LifetimeDependenceInfo::infer(afd);
   }
 
   if (resultDependence.has_value()) {
     lifetimeDependencies.push_back(*resultDependence);
   }
 
+  if (lifetimeDependencies.empty()) {
+    return std::nullopt;
+  }
+
   return afd->getASTContext().AllocateCopy(lifetimeDependencies);
 }
 
diff --git a/test/SIL/lifetime_dependence_span_lifetime_attr.swift b/test/SIL/lifetime_dependence_span_lifetime_attr.swift
@@ -0,0 +1,179 @@
+// RUN: %target-swift-frontend %s -emit-sil \
+// RUN:   -enable-experimental-feature NonescapableTypes
+
+// REQUIRES: asserts
+// REQUIRES: swift_in_compiler
+
+// TODO: Use real Range
+public struct FakeRange<Bound> {
+  public let lowerBound: Bound
+  public let upperBound: Bound
+}
+
+// TODO: Use real Optional
+public enum FakeOptional<T> {
+  case none
+  case some(T)
+}
+
+public struct SpanIndex<Element> {
+  let _rawValue: UnsafeRawPointer
+
+  internal init(rawValue: UnsafeRawPointer) {
+    _rawValue = rawValue
+  }
+
+  var isAligned: Bool {
+    (Int(bitPattern: _rawValue) & (MemoryLayout<Element>.alignment-1)) == 0
+  }
+}
+
+extension SpanIndex: Equatable {}
+
+extension SpanIndex: Hashable {}
+
+extension SpanIndex: Strideable {
+  public typealias Stride = Int
+
+  public func distance(to other: SpanIndex) -> Int {
+    let bytes = _rawValue.distance(to: other._rawValue)
+    let (q, r) = bytes.quotientAndRemainder(dividingBy: MemoryLayout<Element>.stride)
+    precondition(r == 0)
+    return q
+  }
+
+  public func advanced(by n: Int) -> SpanIndex {
+    .init(rawValue: _rawValue.advanced(by: n &* MemoryLayout<Element>.stride))
+  }
+}
+
+extension SpanIndex: Comparable {
+  public static func <(lhs: SpanIndex, rhs: SpanIndex) -> Bool {
+    lhs._rawValue < rhs._rawValue
+  }
+}
+
+public struct Span<Element> : ~Escapable {
+  let start: SpanIndex<Element>
+  public let count: Int
+  private var baseAddress: UnsafeRawPointer { start._rawValue }
+// CHECK-LABEL: sil @$s025lifetime_dependence_span_A5_attr4SpanV11baseAddress5count9dependsOnACyxGSV_Siqd__htcRi_d__Ri0_d__lufC : $@convention(method) <Element><Owner where Owner : ~Copyable, Owner : ~Escapable> (UnsafeRawPointer, Int, @in_guaranteed Owner, @thin Span<Element>.Type) -> _inherit(2) @owned Span<Element> {
+  @lifetime(owner)
+  public init<Owner: ~Copyable & ~Escapable>(
+      baseAddress: UnsafeRawPointer,
+      count: Int,
+      dependsOn owner: borrowing Owner
+    ) {
+      self.init(
+        start: .init(rawValue: baseAddress), count: count, dependsOn: owner
+      )
+  }
+// CHECK-LABEL: sil hidden @$s025lifetime_dependence_span_A5_attr4SpanV5start5count9dependsOnACyxGAA0E5IndexVyxG_Siqd__htcRi_d__Ri0_d__lufC : $@convention(method) <Element><Owner where Owner : ~Copyable, Owner : ~Escapable> (SpanIndex<Element>, Int, @in_guaranteed Owner, @thin Span<Element>.Type) -> _inherit(2) @owned Span<Element> {
+  @lifetime(owner)
+  init<Owner: ~Copyable & ~Escapable>(
+    start index: SpanIndex<Element>,
+    count: Int,
+    dependsOn owner: borrowing Owner
+  ) {
+    precondition(count >= 0, "Count must not be negative")
+    if !_isPOD(Element.self) {
+      precondition(
+        index.isAligned,
+        "baseAddress must be properly aligned for \(Element.self)"
+      )
+    }
+    self.start = index
+    self.count = count
+  }
+  @_unsafeNonescapableResult
+  init(start index: SpanIndex<Element>,
+    count: Int) {
+    precondition(count >= 0, "Count must not be negative")
+    if !_isPOD(Element.self) {
+      precondition(
+        index.isAligned,
+        "baseAddress must be properly aligned for \(Element.self)"
+      )
+    }
+    self.start = index
+    self.count = count
+  }
+}
+
+extension Span {
+  public typealias Index = SpanIndex<Element>
+  public typealias SubSequence = Self
+
+  public var startIndex: Index { start }
+  public var endIndex: Index { start.advanced(by: count) }  
+
+  @inlinable @inline(__always) 
+  public func distance(from start: Index, to end: Index) -> Int {
+    start.distance(to: end)
+  }
+ 
+  public subscript(position: Index) -> Element {
+    get {
+      if _isPOD(Element.self) {
+        return position._rawValue.loadUnaligned(as: Element.self)
+      }
+      else {
+        return position._rawValue.load(as: Element.self)
+      }
+    }
+  }
+ 
+// CHECK-LABEL: sil @$s025lifetime_dependence_span_A5_attr4SpanVyACyxGAA9FakeRangeVyAA0E5IndexVyxGGcig : $@convention(method) <Element> (FakeRange<SpanIndex<Element>>, @guaranteed Span<Element>) -> _inherit(1) @owned Span<Element> {
+  public subscript(bounds: FakeRange<SpanIndex<Element>>) -> Self {
+  @lifetime(self)
+    get {
+      Span(
+        start: bounds.lowerBound,
+        count: bounds.upperBound.distance(to:bounds.lowerBound) / MemoryLayout<Element>.stride
+      )
+    }
+  }
+
+// CHECK-LABEL: sil @$s025lifetime_dependence_span_A5_attr4SpanV6prefix4upToACyxGAA0E5IndexVyxG_tF : $@convention(method) <Element> (SpanIndex<Element>, @guaranteed Span<Element>) -> _inherit(1) @owned Span<Element> {
+  @lifetime(self)
+  borrowing public func prefix(upTo index: SpanIndex<Element>) -> Self {
+    index == startIndex
+    ? Self(start: start, count: 0, dependsOn: copy self)
+    : prefix(through: index.advanced(by: -1))
+  }
+
+// CHECK-LABEL: sil @$s025lifetime_dependence_span_A5_attr4SpanV6prefix7throughACyxGAA0E5IndexVyxG_tF : $@convention(method) <Element> (SpanIndex<Element>, @guaranteed Span<Element>) -> _inherit(1) @owned Span<Element> {
+  @lifetime(self)
+  borrowing public func prefix(through index: Index) -> Self {
+    let nc = distance(from: startIndex, to: index) &+ 1
+    return Self(start: start, count: nc, dependsOn: copy self)
+  }
+
+// CHECK-LABEL: sil @$s025lifetime_dependence_span_A5_attr4SpanV6prefixyACyxGSiF : $@convention(method) <Element> (Int, @owned Span<Element>) -> _inherit(1) @owned Span<Element> {
+  @lifetime(self)
+  consuming public func prefix(_ maxLength: Int) -> Self {
+    precondition(maxLength >= 0, "Can't have a prefix of negative length.")
+    let nc = maxLength < count ? maxLength : count
+    return Self(start: start, count: nc, dependsOn: self)
+  }
+}
+
+extension ContiguousArray {
+  public var view: Span<Element> {
+    @lifetime(self)
+    borrowing _read {
+      yield Span(
+        baseAddress: _baseAddressIfContiguous!, count: count, dependsOn: self
+      )
+    }
+  }
+}
+
+public func array_view_element(a: ContiguousArray<Int> , i: SpanIndex<Int>) -> Int {
+  a.view[i]
+}
+
+public func array_view_slice_element(a: ContiguousArray<Int> , sliceIdx: FakeRange<SpanIndex<Int>>, Idx: SpanIndex<Int>) -> Int {
+  a.view[sliceIdx][Idx]
+}
+
diff --git a/test/Sema/lifetime_attr.swift b/test/Sema/lifetime_attr.swift
@@ -0,0 +1,35 @@
+// RUN: %target-typecheck-verify-swift -disable-availability-checking -enable-experimental-feature NonescapableTypes
+// REQUIRES: asserts
+
+struct NE : ~Escapable {
+  @lifetime(self) // expected-error{{invalid lifetime dependence on self in an initializer}}
+  init() {}
+}
+
+@lifetime(nonexisting) // expected-error{{invalid parameter name specified 'nonexisting'}}
+func invalidAttrOnNonExistingParam(_ ne: NE) -> NE {
+  ne
+}
+
+@lifetime(self) // expected-error{{invalid lifetime dependence specifier on non-existent self}}
+func invalidAttrOnNonExistingSelf(_ ne: NE) -> NE {
+  ne
+}
+
+@lifetime(2) // expected-error{{invalid parameter index specified 2}}
+func invalidAttrOnNonExistingParamIndex(_ ne: NE) -> NE {
+  ne
+}
+
+@lifetime(ne, ne) // expected-error{{duplicate lifetime dependence specifier}}
+func invalidDuplicateLifetimeDependence1(_ ne: borrowing NE) -> NE {
+  ne
+}
+
+class Klass {}
+
+@lifetime(x) // expected-error{{invalid use of lifetime dependence on an Escapable parameter with consuming ownership}}
+func invalidDependence(_ x: consuming Klass) -> NE {
+  NE()
+}
+
