Add support for parsing @lifetime attribute to specify lifetime dependencies on 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,27 @@ public:
   }
 };
 
+class LifetimeAttr final
+    : public DeclAttribute,
+      private llvm::TrailingObjects<LifetimeAttr, LifetimeDependenceSpecifier> {
+
+  friend TrailingObjects;
+
+  unsigned NumEntries = 0;
+
+  explicit LifetimeAttr(SourceLoc atLoc, SourceRange baseRange, bool implicit,
+                        ArrayRef<LifetimeDependenceSpecifier> entries);
+
+public:
+  static LifetimeAttr *create(ASTContext &context, SourceLoc atLoc,
+                              SourceRange baseRange, bool implicit,
+                              ArrayRef<LifetimeDependenceSpecifier> entries);
+
+  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/LifetimeDependence.h

@@ -74,7 +74,9 @@ private:
 
 public:
   static LifetimeDependenceSpecifier getNamedLifetimeDependenceSpecifier(
-      SourceLoc loc, ParsedLifetimeDependenceKind kind, Identifier name) {
+      SourceLoc loc, Identifier name,
+      ParsedLifetimeDependenceKind kind =
+          ParsedLifetimeDependenceKind::Default) {
     return {loc, SpecifierKind::Named, kind, name};
   }
 
@@ -84,13 +86,15 @@ public:
   }
 
   static LifetimeDependenceSpecifier getOrderedLifetimeDependenceSpecifier(
-      SourceLoc loc, ParsedLifetimeDependenceKind kind, unsigned index) {
+      SourceLoc loc, unsigned index,
+      ParsedLifetimeDependenceKind kind =
+          ParsedLifetimeDependenceKind::Default) {
     return {loc, SpecifierKind::Ordered, kind, index};
   }
 
-  static LifetimeDependenceSpecifier
+  static LifetimeDependenceSpecifier getSelfLifetimeDependenceSpecifier(
-  getSelfLifetimeDependenceSpecifier(SourceLoc loc,
+      SourceLoc loc, ParsedLifetimeDependenceKind kind =
-                                     ParsedLifetimeDependenceKind kind) {
+                         ParsedLifetimeDependenceKind::Default) {
     return {loc, SpecifierKind::Self, kind, {}};
   }
 

include/swift/Parse/Parser.h

@@ -1178,6 +1178,10 @@ public:
       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,

lib/AST/Attr.cpp

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

lib/ASTGen/Sources/ASTGen/DeclAttrs.swift

@@ -189,6 +189,8 @@ extension ASTGenVisitor {
         fatalError("unimplemented")
       case .allowFeatureSuppression:
         return self.generateAllowFeatureSuppressionAttr(attribute: node)?.asDeclAttribute
+      case .lifetime:
+        fatalError("unimplemented")
 
       // Simple attributes.
       case .alwaysEmitConformanceMetadata,

lib/Parse/ParseDecl.cpp

@@ -2801,6 +2801,96 @@ static std::optional<Identifier> parseSingleAttrOptionImpl(
   return P.Context.getIdentifier(parsedName);
 }
 
+ParserResult<LifetimeAttr> Parser::parseLifetimeAttribute(SourceLoc atLoc,
+                                                          SourceLoc loc) {
+  ParserStatus status;
+  SmallVector<LifetimeDependenceSpecifier> lifetimeEntries;
+
+  if (!Context.LangOpts.hasFeature(Feature::NonescapableTypes)) {
+    diagnose(loc, diag::requires_experimental_feature, "lifetime attribute",
+             false, getFeatureName(Feature::NonescapableTypes));
+    status.setIsParseError();
+    return status;
+  }
+
+  if (!Tok.isFollowingLParen()) {
+    diagnose(loc, diag::expected_lparen_after_lifetime_dependence);
+    status.setIsParseError();
+    return status;
+  }
+  // consume the l_paren
+  auto lParenLoc = consumeToken();
+
+  SourceLoc rParenLoc;
+  bool foundParamId = false;
+  status = parseList(
+      tok::r_paren, lParenLoc, rParenLoc, /*AllowSepAfterLast*/ false,
+      diag::expected_rparen_after_lifetime_dependence, [&]() -> ParserStatus {
+        ParserStatus listStatus;
+        foundParamId = true;
+        switch (Tok.getKind()) {
+        case tok::identifier: {
+          Identifier paramName;
+          auto paramLoc =
+              consumeIdentifier(paramName, /*diagnoseDollarPrefix=*/false);
+          if (paramName.is("immortal")) {
+            lifetimeEntries.push_back(
+                LifetimeDependenceSpecifier::
+                    getImmortalLifetimeDependenceSpecifier(paramLoc));
+          } else {
+            lifetimeEntries.push_back(
+                LifetimeDependenceSpecifier::
+                    getNamedLifetimeDependenceSpecifier(paramLoc, paramName));
+          }
+          break;
+        }
+        case tok::integer_literal: {
+          SourceLoc paramLoc;
+          unsigned paramNum;
+          if (parseUnsignedInteger(
+                  paramNum, paramLoc,
+                  diag::expected_param_index_lifetime_dependence)) {
+            listStatus.setIsParseError();
+            return listStatus;
+          }
+          lifetimeEntries.push_back(
+              LifetimeDependenceSpecifier::
+                  getOrderedLifetimeDependenceSpecifier(paramLoc, paramNum));
+          break;
+        }
+        case tok::kw_self: {
+          auto paramLoc = consumeToken(tok::kw_self);
+          lifetimeEntries.push_back(
+              LifetimeDependenceSpecifier::getSelfLifetimeDependenceSpecifier(
+                  paramLoc));
+          break;
+        }
+        default:
+          diagnose(
+              Tok,
+              diag::
+                  expected_identifier_or_index_or_self_after_lifetime_dependence);
+          listStatus.setIsParseError();
+          return listStatus;
+        }
+        return listStatus;
+      });
+
+  if (!foundParamId) {
+    diagnose(
+        Tok,
+        diag::expected_identifier_or_index_or_self_after_lifetime_dependence);
+    status.setIsParseError();
+    return status;
+  }
+
+  assert(!lifetimeEntries.empty());
+  SourceRange range(loc, rParenLoc);
+  return ParserResult<LifetimeAttr>(
+      LifetimeAttr::create(Context, atLoc, SourceRange(loc, rParenLoc),
+                           /* implicit */ false, lifetimeEntries));
+}
+
 /// Parses a (possibly optional) argument for an attribute containing a single, arbitrary identifier.
 ///
 /// \param P The parser object.
@@ -4062,6 +4152,13 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
     Attributes.add(attr);
     break;
   }
+  case DeclAttrKind::Lifetime: {
+    auto Attr = parseLifetimeAttribute(AtLoc, Loc);
+    Status |= Attr;
+    if (Attr.isNonNull())
+      Attributes.add(Attr.get());
+    break;
+  }
   }
 
   if (DuplicateAttribute) {
@@ -5147,7 +5244,7 @@ ParserStatus Parser::parseLifetimeDependenceSpecifiers(
               specifierList.push_back(
                   LifetimeDependenceSpecifier::
                       getNamedLifetimeDependenceSpecifier(
-                          paramLoc, lifetimeDependenceKind, paramName));
+                          paramLoc, paramName, lifetimeDependenceKind));
             }
             break;
           }
@@ -5163,7 +5260,7 @@ ParserStatus Parser::parseLifetimeDependenceSpecifiers(
             specifierList.push_back(
                 LifetimeDependenceSpecifier::
                     getOrderedLifetimeDependenceSpecifier(
-                        paramLoc, lifetimeDependenceKind, paramNum));
+                        paramLoc, paramNum, lifetimeDependenceKind));
             break;
           }
           case tok::kw_self: {

lib/Sema/TypeCheckAttr.cpp

@@ -366,6 +366,7 @@ public:
   void visitWeakLinkedAttr(WeakLinkedAttr *attr);
   void visitSILGenNameAttr(SILGenNameAttr *attr);
   void visitUnsafeAttr(UnsafeAttr *attr);
+  void visitLifetimeAttr(LifetimeAttr *attr);
 };
 
 } // end anonymous namespace
@@ -7674,6 +7675,8 @@ void AttributeChecker::visitUnsafeAttr(UnsafeAttr *attr) {
   diagnoseAndRemoveAttr(attr, diag::unsafe_attr_disabled);
 }
 
+void AttributeChecker::visitLifetimeAttr(LifetimeAttr *attr) {}
+
 namespace {
 
 class ClosureAttributeChecker

lib/Sema/TypeCheckDeclOverride.cpp

@@ -1701,6 +1701,7 @@ namespace  {
     UNINTERESTING_ATTR(UnsafeNonEscapableResult)
     UNINTERESTING_ATTR(StaticExclusiveOnly)
     UNINTERESTING_ATTR(PreInverseGenerics)
+    UNINTERESTING_ATTR(Lifetime)
 #undef UNINTERESTING_ATTR
 
     void visitAvailableAttr(AvailableAttr *attr) {

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 = 889; // public-module-name
+const uint16_t SWIFTMODULE_VERSION_MINOR = 890; // @lifetime attribute
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -2490,6 +2490,15 @@ namespace decls_block {
                                // trialed by introduced conformances
   >;
 
+  using LifetimeDeclAttrLayout =
+      BCRecordLayout<Lifetime_DECL_ATTR,
+                     BCVBR<4>,           // targetIndex
+                     BCFixed<1>,         // isImmortal
+                     BCFixed<1>,         // hasInheritLifetimeParamIndices
+                     BCFixed<1>,         // hasScopeLifetimeParamIndices
+                     BCArray<BCFixed<1>> // concatenated param indices
+                     >;
+
 #undef SYNTAX_SUGAR_TYPE_LAYOUT
 #undef TYPE_LAYOUT
 #undef TYPE_LAYOUT_IMPL

lib/Serialization/Serialization.cpp

@@ -3382,6 +3382,10 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       RawLayoutDeclAttrLayout::emitRecord(
         S.Out, S.ScratchRecord, abbrCode, attr->isImplicit(),
         typeID, countID, rawSize, rawAlign, attr->shouldMoveAsLikeType());
+      return;
+    }
+    case DeclAttrKind::Lifetime: {
+      return;
     }
     }
   }

test/Parse/lifetime_attr.swift

@@ -0,0 +1,22 @@
+// RUN: %target-typecheck-verify-swift -disable-availability-checking -enable-experimental-feature NonescapableTypes 
+// REQUIRES: asserts
+
+struct NE : ~Escapable {
+
+}
+
+@lifetime(ne)
+func derive(_ ne: NE) -> NE {
+  ne
+}
+
+@lifetime // expected-error{{expected '(' after lifetime dependence specifier}} 
+func testMissingLParenError(_ ne: NE) -> NE {
+  ne
+}
+
+@lifetime() // expected-error{{expected identifier, index or self in lifetime dependence specifier}}
+func testMissingDependence(_ ne: NE) -> NE {
+  ne
+}
+