Introduce @lifetime attribute to specify lifetime dependence on function declarations

include/swift/AST/Attr.h

@@ -24,6 +24,7 @@
 #include "swift/AST/DeclNameLoc.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/KnownProtocols.h"
+#include "swift/AST/LifetimeDependence.h"
 #include "swift/AST/MacroDeclaration.h"
 #include "swift/AST/Ownership.h"
 #include "swift/AST/PlatformKind.h"
@@ -2627,6 +2628,31 @@ class RawLayoutAttr final : public DeclAttribute {
   }
 };
 
+class LifetimeAttr final
+    : public DeclAttribute,
+      private llvm::TrailingObjects<LifetimeAttr, LifetimeEntry> {
+
+  friend TrailingObjects;
+
+  unsigned NumEntries = 0;
+
+  explicit LifetimeAttr(SourceLoc atLoc, SourceRange baseRange, bool implicit,
+                        ArrayRef<LifetimeEntry> entries);
+
+public:
+  static LifetimeAttr *create(ASTContext &context, SourceLoc atLoc,
+                              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;
+  }
+};
+
 /// Predicate used to filter MatchingAttributeRange.
 template <typename ATTR, bool AllowInvalid> struct ToAttributeKind {
   ToAttributeKind() {}

include/swift/AST/DeclAttr.def

@@ -506,6 +506,10 @@ SIMPLE_DECL_ATTR(unsafe, Unsafe,
   ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove,
   160)
 
+DECL_ATTR(lifetime, Lifetime,
+  OnAccessor | OnConstructor | OnFunc | OnSubscript | LongAttribute | ABIBreakingToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove,
+  161)
+
 LAST_DECL_ATTR(Unsafe)
 
 #undef DECL_ATTR_ALIAS

include/swift/AST/DiagnosticsSema.def

@@ -7952,6 +7952,9 @@ ERROR(lifetime_dependence_function_type, none,
       "lifetime dependencies on function types are not supported",
       ())
 
+ERROR(lifetime_dependence_immortal_alone, none,
+      "cannot specify any other dependence source along with immortal", ())
+
 //===----------------------------------------------------------------------===//
 //                             MARK: Sending
 //===----------------------------------------------------------------------===//

include/swift/AST/LifetimeDependence.h

@@ -1,4 +1,4 @@
-//===--- LifetimeDependenceSpecifiers.h ------------------------*- C++ -*-===//
+//===--- LifetimeDependence.h ---------------------------------*- C++ -*-===//
 //
 // This source file is part of the Swift.org open source project
 //
@@ -43,13 +43,12 @@ enum class ParsedLifetimeDependenceKind : uint8_t {
 
 enum class LifetimeDependenceKind : uint8_t { Inherit = 0, Scope };
 
-class LifetimeDependenceSpecifier {
-public:
-  enum class SpecifierKind { Named, Ordered, Self, Immortal };
+enum class LifetimeEntryKind { Named, Ordered, Self, Immortal };
 
+class LifetimeEntry {
 private:
   SourceLoc loc;
-  SpecifierKind specifierKind;
+  LifetimeEntryKind lifetimeEntryKind;
   ParsedLifetimeDependenceKind parsedLifetimeDependenceKind;
   union Value {
     struct {
@@ -65,68 +64,72 @@ class LifetimeDependenceSpecifier {
     Value() {}
   } value;
 
-  LifetimeDependenceSpecifier(
-      SourceLoc loc, SpecifierKind specifierKind,
-      ParsedLifetimeDependenceKind parsedLifetimeDependenceKind, Value value)
-      : loc(loc), specifierKind(specifierKind),
+  LifetimeEntry(SourceLoc loc, LifetimeEntryKind lifetimeEntryKind,
+                ParsedLifetimeDependenceKind parsedLifetimeDependenceKind,
+                Value value)
+      : loc(loc), lifetimeEntryKind(lifetimeEntryKind),
         parsedLifetimeDependenceKind(parsedLifetimeDependenceKind),
         value(value) {}
 
 public:
-  static LifetimeDependenceSpecifier getNamedLifetimeDependenceSpecifier(
-      SourceLoc loc, ParsedLifetimeDependenceKind kind, Identifier name) {
-    return {loc, SpecifierKind::Named, kind, name};
+  static LifetimeEntry
+  getNamedLifetimeEntry(SourceLoc loc, Identifier name,
+                        ParsedLifetimeDependenceKind kind =
+                            ParsedLifetimeDependenceKind::Default) {
+    return {loc, LifetimeEntryKind::Named, kind, name};
   }
 
-  static LifetimeDependenceSpecifier
-  getImmortalLifetimeDependenceSpecifier(SourceLoc loc) {
-    return {loc, SpecifierKind::Immortal, {}, {}};
+  static LifetimeEntry getImmortalLifetimeEntry(SourceLoc loc) {
+    return {loc, LifetimeEntryKind::Immortal, {}, {}};
   }
 
-  static LifetimeDependenceSpecifier getOrderedLifetimeDependenceSpecifier(
-      SourceLoc loc, ParsedLifetimeDependenceKind kind, unsigned index) {
-    return {loc, SpecifierKind::Ordered, kind, index};
+  static LifetimeEntry
+  getOrderedLifetimeEntry(SourceLoc loc, unsigned index,
+                          ParsedLifetimeDependenceKind kind =
+                              ParsedLifetimeDependenceKind::Default) {
+    return {loc, LifetimeEntryKind::Ordered, kind, index};
   }
 
-  static LifetimeDependenceSpecifier
-  getSelfLifetimeDependenceSpecifier(SourceLoc loc,
-                                     ParsedLifetimeDependenceKind kind) {
-    return {loc, SpecifierKind::Self, kind, {}};
+  static LifetimeEntry
+  getSelfLifetimeEntry(SourceLoc loc,
+                       ParsedLifetimeDependenceKind kind =
+                           ParsedLifetimeDependenceKind::Default) {
+    return {loc, LifetimeEntryKind::Self, kind, {}};
   }
 
   SourceLoc getLoc() const { return loc; }
 
-  SpecifierKind getSpecifierKind() const { return specifierKind; }
+  LifetimeEntryKind getLifetimeEntryKind() const { return lifetimeEntryKind; }
 
   ParsedLifetimeDependenceKind getParsedLifetimeDependenceKind() const {
     return parsedLifetimeDependenceKind;
   }
 
   Identifier getName() const {
-    assert(specifierKind == SpecifierKind::Named);
+    assert(lifetimeEntryKind == LifetimeEntryKind::Named);
     return value.Named.name;
   }
 
   unsigned getIndex() const {
-    assert(specifierKind == SpecifierKind::Ordered);
+    assert(lifetimeEntryKind == LifetimeEntryKind::Ordered);
     return value.Ordered.index;
   }
 
   std::string getParamString() const {
-    switch (specifierKind) {
-    case SpecifierKind::Named:
+    switch (lifetimeEntryKind) {
+    case LifetimeEntryKind::Named:
       return value.Named.name.str().str();
-    case SpecifierKind::Self:
+    case LifetimeEntryKind::Self:
       return "self";
-    case SpecifierKind::Ordered:
+    case LifetimeEntryKind::Ordered:
       return std::to_string(value.Ordered.index);
-    case SpecifierKind::Immortal:
+    case LifetimeEntryKind::Immortal:
       return "immortal";
     }
-    llvm_unreachable("Invalid LifetimeDependenceSpecifier::SpecifierKind");
+    llvm_unreachable("Invalid LifetimeEntryKind");
   }
 
-  std::string getLifetimeDependenceSpecifierString() const {
+  std::string getDependsOnString() const {
     switch (parsedLifetimeDependenceKind) {
     case ParsedLifetimeDependenceKind::Default:
       return "dependsOn(" + getParamString() + ")";
@@ -135,8 +138,7 @@ class LifetimeDependenceSpecifier {
     case ParsedLifetimeDependenceKind::Inherit:
       return "dependsOn(inherited " + getParamString() + ")";
     }
-    llvm_unreachable(
-        "Invalid LifetimeDependenceSpecifier::ParsedLifetimeDependenceKind");
+    llvm_unreachable("Invalid LifetimeEntry::ParsedLifetimeDependenceKind");
   }
 };
 
@@ -151,10 +153,14 @@ class LifetimeDependenceInfo {
                                             unsigned sourceIndex,
                                             LifetimeDependenceKind kind);
 
-  /// Builds LifetimeDependenceInfo on a result or parameter from a swift decl
+  /// Builds LifetimeDependenceInfo from @lifetime attribute
+  static std::optional<ArrayRef<LifetimeDependenceInfo>>
+  fromLifetimeAttribute(AbstractFunctionDecl *afd);
+
+  /// Builds LifetimeDependenceInfo from dependsOn type modifier
   static std::optional<LifetimeDependenceInfo>
-  fromTypeRepr(AbstractFunctionDecl *afd, LifetimeDependentTypeRepr *typeRepr,
-               unsigned targetIndex);
+  fromDependsOn(AbstractFunctionDecl *afd, TypeRepr *targetRepr,
+                Type targetType, unsigned targetIndex);
 
   /// Infer LifetimeDependenceInfo on result
   static std::optional<LifetimeDependenceInfo> infer(AbstractFunctionDecl *afd);
@@ -169,9 +175,9 @@ class LifetimeDependenceInfo {
 
   /// Builds LifetimeDependenceInfo from SIL function type
   static std::optional<LifetimeDependenceInfo>
-  fromTypeRepr(LifetimeDependentTypeRepr *lifetimeDependentRepr,
-               unsigned targetIndex, ArrayRef<SILParameterInfo> params,
-               DeclContext *dc);
+  fromDependsOn(LifetimeDependentTypeRepr *lifetimeDependentRepr,
+                unsigned targetIndex, ArrayRef<SILParameterInfo> params,
+                DeclContext *dc);
 
 public:
   LifetimeDependenceInfo(IndexSubset *inheritLifetimeParamIndices,

include/swift/AST/TypeRepr.h

@@ -1533,31 +1533,28 @@ class SILBoxTypeRepr final : public TypeRepr,
 
 class LifetimeDependentTypeRepr final
     : public SpecifierTypeRepr,
-      private llvm::TrailingObjects<LifetimeDependentTypeRepr,
-                                    LifetimeDependenceSpecifier> {
+      private llvm::TrailingObjects<LifetimeDependentTypeRepr, LifetimeEntry> {
   friend TrailingObjects;
 
   size_t numTrailingObjects(OverloadToken<LifetimeDependentTypeRepr>) const {
     return Bits.LifetimeDependentTypeRepr.NumDependencies;
   }
 
 public:
-  LifetimeDependentTypeRepr(TypeRepr *base,
-                            ArrayRef<LifetimeDependenceSpecifier> specifiers)
+  LifetimeDependentTypeRepr(TypeRepr *base, ArrayRef<LifetimeEntry> specifiers)
       : SpecifierTypeRepr(TypeReprKind::LifetimeDependent, base,
                           specifiers.front().getLoc()) {
     assert(base);
     Bits.LifetimeDependentTypeRepr.NumDependencies = specifiers.size();
     std::uninitialized_copy(specifiers.begin(), specifiers.end(),
-                            getTrailingObjects<LifetimeDependenceSpecifier>());
+                            getTrailingObjects<LifetimeEntry>());
   }
 
-  static LifetimeDependentTypeRepr *
-  create(ASTContext &C, TypeRepr *base,
-         ArrayRef<LifetimeDependenceSpecifier> specifiers);
+  static LifetimeDependentTypeRepr *create(ASTContext &C, TypeRepr *base,
+                                           ArrayRef<LifetimeEntry> specifiers);
 
-  ArrayRef<LifetimeDependenceSpecifier> getLifetimeDependencies() const {
-    return {getTrailingObjects<LifetimeDependenceSpecifier>(),
+  ArrayRef<LifetimeEntry> getLifetimeDependencies() const {
+    return {getTrailingObjects<LifetimeEntry>(),
             Bits.LifetimeDependentTypeRepr.NumDependencies};
   }
 

include/swift/Parse/Parser.h

@@ -1178,6 +1178,10 @@ class Parser {
       MacroSyntax syntax, SourceLoc AtLoc, SourceLoc Loc
   );
 
+  /// Parse the @lifetime attribute.
+  ParserResult<LifetimeAttr> parseLifetimeAttribute(SourceLoc AtLoc,
+                                                    SourceLoc Loc);
+
   /// Parse a specific attribute.
   ParserStatus parseDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc,
                                   SourceLoc AtEndLoc,
@@ -1266,8 +1270,8 @@ class Parser {
                                         ConventionTypeAttr *&result,
                                         bool justChecking);
 
-  ParserStatus parseLifetimeDependenceSpecifiers(
-      SmallVectorImpl<LifetimeDependenceSpecifier> &specifierList);
+  ParserStatus
+  parseLifetimeEntries(SmallVectorImpl<LifetimeEntry> &specifierList);
 
   ParserResult<ImportDecl> parseDeclImport(ParseDeclOptions Flags,
                                            DeclAttributes &Attributes);
@@ -1470,7 +1474,7 @@ class Parser {
     SourceLoc ConstLoc;
     SourceLoc SendingLoc;
     SmallVector<TypeOrCustomAttr> Attributes;
-    SmallVector<LifetimeDependenceSpecifier> lifetimeDependenceSpecifiers;
+    SmallVector<LifetimeEntry> lifetimeEntries;
 
     ParsedTypeAttributeList(ParseTypeReason reason) : ParseReason(reason) {}
 

lib/AST/ASTDumper.cpp

@@ -3579,7 +3579,7 @@ class PrintTypeRepr : public TypeReprVisitor<PrintTypeRepr, void, StringRef>,
     for (auto &dep : T->getLifetimeDependencies()) {
       printFieldRaw(
           [&](raw_ostream &out) {
-            out << " " << dep.getLifetimeDependenceSpecifierString() << " ";
+            out << " " << dep.getDependsOnString() << " ";
           },
           "");
     }

lib/AST/ASTPrinter.cpp

@@ -4110,7 +4110,7 @@ void PrintAST::visitFuncDecl(FuncDecl *decl) {
         if (auto *typeRepr = dyn_cast_or_null<LifetimeDependentTypeRepr>(
                 decl->getResultTypeRepr())) {
           for (auto &dep : typeRepr->getLifetimeDependencies()) {
-            Printer << " " << dep.getLifetimeDependenceSpecifierString() << " ";
+            Printer << " " << dep.getDependsOnString() << " ";
           }
         }
       }
@@ -4335,7 +4335,7 @@ void PrintAST::visitConstructorDecl(ConstructorDecl *decl) {
         auto *typeRepr =
             cast<LifetimeDependentTypeRepr>(decl->getResultTypeRepr());
         for (auto &dep : typeRepr->getLifetimeDependencies()) {
-          Printer << dep.getLifetimeDependenceSpecifierString() << " ";
+          Printer << dep.getDependsOnString() << " ";
         }
         // TODO: Handle failable initializers with lifetime dependent returns
         Printer << "Self";

lib/AST/Attr.cpp

@@ -1813,6 +1813,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");
@@ -2011,6 +2025,8 @@ StringRef DeclAttribute::getAttrName() const {
     } else {
       return "_allowFeatureSuppression";
     }
+  case DeclAttrKind::Lifetime:
+    return "lifetime";
   }
   llvm_unreachable("bad DeclAttrKind");
 }
@@ -3040,6 +3056,22 @@ AllowFeatureSuppressionAttr *AllowFeatureSuppressionAttr::create(
       AllowFeatureSuppressionAttr(atLoc, range, implicit, inverted, features);
 }
 
+LifetimeAttr::LifetimeAttr(SourceLoc atLoc, SourceRange baseRange,
+                           bool implicit, ArrayRef<LifetimeEntry> entries)
+    : DeclAttribute(DeclAttrKind::Lifetime, atLoc, baseRange, implicit),
+      NumEntries(entries.size()) {
+  std::copy(entries.begin(), entries.end(),
+            getTrailingObjects<LifetimeEntry>());
+}
+
+LifetimeAttr *LifetimeAttr::create(ASTContext &context, SourceLoc atLoc,
+                                   SourceRange baseRange, bool implicit,
+                                   ArrayRef<LifetimeEntry> entries) {
+  unsigned size = totalSizeToAlloc<LifetimeEntry>(entries.size());
+  void *mem = context.Allocate(size, alignof(LifetimeEntry));
+  return new (mem) LifetimeAttr(atLoc, baseRange, implicit, entries);
+}
+
 void swift::simple_display(llvm::raw_ostream &out, const DeclAttribute *attr) {
   if (attr)
     attr->print(out);