[AST] Add @_lifetime type attribute

Author: aidan-hall

include/swift/AST/ASTBridging.h

@@ -2496,6 +2496,13 @@ BridgedDifferentiableTypeAttr BridgedDifferentiableTypeAttr_createParsed(
     swift::SourceLoc nameLoc, swift::SourceRange parensRange,
     BridgedDifferentiabilityKind cKind, swift::SourceLoc kindLoc);
 
+SWIFT_NAME("BridgedLifetimeTypeAttr.createParsed(_:atLoc:nameLoc:"
+           "parensRange:entry:)")
+BridgedLifetimeTypeAttr BridgedLifetimeTypeAttr_createParsed(
+    BridgedASTContext cContext, swift::SourceLoc atLoc,
+    swift::SourceLoc nameLoc, swift::SourceRange parensRange,
+    BridgedLifetimeEntry entry);
+
 SWIFT_NAME("BridgedIsolatedTypeAttr.createParsed(_:atLoc:nameLoc:parensRange:"
            "isolationKind:isolationKindLoc:)")
 BridgedIsolatedTypeAttr BridgedIsolatedTypeAttr_createParsed(

include/swift/AST/Attr.h

@@ -4349,6 +4349,19 @@ class DifferentiableTypeAttr
   void printImpl(ASTPrinter &printer, const PrintOptions &options) const;
 };
 
+class LifetimeTypeAttr : public SimpleTypeAttrWithArgs<TypeAttrKind::Lifetime> {
+  LifetimeEntry *entry;
+
+public:
+  LifetimeTypeAttr(SourceLoc atLoc, SourceLoc kwLoc, SourceRange parens,
+                   LifetimeEntry *entry)
+      : SimpleTypeAttr(atLoc, kwLoc, parens), entry(entry) {}
+
+  LifetimeEntry *getLifetimeEntry() const { return entry; }
+
+  void printImpl(ASTPrinter &printer, const PrintOptions &options) const;
+};
+
 class OpaqueReturnTypeOfTypeAttr
     : public SimpleTypeAttrWithArgs<TypeAttrKind::OpaqueReturnTypeOf> {
   Located<StringRef> MangledName;

include/swift/AST/TypeAttr.def

@@ -61,6 +61,7 @@ SIMPLE_TYPE_ATTR(retroactive, Retroactive)
 SIMPLE_TYPE_ATTR(unchecked, Unchecked)
 SIMPLE_TYPE_ATTR(preconcurrency, Preconcurrency)
 SIMPLE_TYPE_ATTR(unsafe, Unsafe)
+TYPE_ATTR(_lifetime, Lifetime)
 SIMPLE_TYPE_ATTR(_local, Local)
 SIMPLE_TYPE_ATTR(_noMetadata, NoMetadata)
 TYPE_ATTR(_opaqueReturnTypeOf, OpaqueReturnTypeOf)

lib/AST/Attr.cpp

@@ -238,6 +238,16 @@ void DifferentiableTypeAttr::printImpl(ASTPrinter &printer,
   printer.printStructurePost(PrintStructureKind::BuiltinAttribute);
 }
 
+void LifetimeTypeAttr::printImpl(ASTPrinter &printer,
+                                 const PrintOptions &options) const {
+  printer.callPrintStructurePre(PrintStructureKind::BuiltinAttribute);
+  printer.printAttrName("@_lifetime");
+  printer << "(";
+  printer << entry->getString();
+  printer << ")";
+  printer.printStructurePost(PrintStructureKind::BuiltinAttribute);
+}
+
 void ConventionTypeAttr::printImpl(ASTPrinter &printer,
                                    const PrintOptions &options) const {
   printer.callPrintStructurePre(PrintStructureKind::BuiltinAttribute);

lib/AST/Bridging/TypeAttributeBridging.cpp

@@ -69,6 +69,14 @@ BridgedDifferentiableTypeAttr BridgedDifferentiableTypeAttr_createParsed(
       atLoc, nameLoc, parensRange, {unbridged(cKind), kindLoc});
 }
 
+BridgedLifetimeTypeAttr BridgedLifetimeTypeAttr_createParsed(
+    BridgedASTContext cContext, swift::SourceLoc atLoc,
+    swift::SourceLoc nameLoc, swift::SourceRange parensRange,
+    BridgedLifetimeEntry entry) {
+  return new (cContext.unbridged())
+      LifetimeTypeAttr(atLoc, nameLoc, parensRange, entry.unbridged());
+}
+
 BridgedIsolatedTypeAttr BridgedIsolatedTypeAttr_createParsed(
     BridgedASTContext cContext, SourceLoc atLoc, SourceLoc nameLoc,
     SourceRange parensRange, BridgedIsolatedTypeAttrIsolationKind cIsolation,

lib/ASTGen/Sources/ASTGen/TypeAttrs.swift

@@ -84,6 +84,9 @@ extension ASTGenVisitor {
       case .Differentiable:
         return (self.generateDifferentiableTypeAttr(attribute: node)?.asTypeAttribute)
           .map(BridgedTypeOrCustomAttr.typeAttr(_:))
+      case .Lifetime:
+        return (self.generateLifetimeTypeAttr(attribute: node)?.asTypeAttribute)
+          .map(BridgedTypeOrCustomAttr.typeAttr(_:))
       case .OpaqueReturnTypeOf:
         return (self.generateOpaqueReturnTypeOfTypeAttr(attribute: node)?.asTypeAttribute)
           .map(BridgedTypeOrCustomAttr.typeAttr(_:))
@@ -246,6 +249,20 @@ extension ASTGenVisitor {
       kindLoc: differentiabilityLoc
     )
   }
+
+  func generateLifetimeTypeAttr(attribute node: AttributeSyntax) -> BridgedLifetimeTypeAttr? {
+    guard let entry = self.generateLifetimeEntry(attribute: node) else {
+      return nil
+    }
+
+    return .createParsed(
+      self.ctx,
+      atLoc: self.generateSourceLoc(node.atSign),
+      nameLoc: self.generateSourceLoc(node.attributeName),
+      parensRange: self.generateAttrParensRange(attribute: node),
+      entry: entry
+    )
+  }
 
   func generateIsolatedTypeAttr(attribute node: AttributeSyntax) -> BridgedIsolatedTypeAttr? {
     let isolationKindLoc = self.generateSourceLoc(node.arguments)

lib/Parse/ParseDecl.cpp

@@ -39,6 +39,7 @@
 #include "swift/Parse/ParseDeclName.h"
 #include "swift/Parse/ParseSILSupport.h"
 #include "swift/Parse/Parser.h"
+#include "swift/Parse/ParserResult.h"
 #include "swift/Strings.h"
 #include "swift/Subsystems.h"
 #include "llvm/ADT/PointerUnion.h"
@@ -4978,6 +4979,22 @@ ParserStatus Parser::parseTypeAttribute(TypeOrCustomAttr &result,
     return makeParserSuccess();
   }
 
+  case TypeAttrKind::Lifetime: {
+    const auto entryResult = parseLifetimeEntry(AtLoc);
+    if (entryResult.isNull()) {
+      return makeParserError();
+    }
+
+    auto *entry = entryResult.get();
+
+    if (!justChecking) {
+      result = new (Context) LifetimeTypeAttr(
+          AtLoc, attrLoc, SourceRange(entry->getStartLoc(), entry->getEndLoc()),
+          entryResult.get());
+    }
+    return makeParserSuccess();
+  }
+
   case TypeAttrKind::Convention: {
     ConventionTypeAttr *convention = nullptr;
     if (parseConventionAttributeInternal(AtLoc, attrLoc, convention,

test/expr/closure/lifetimes.swift

@@ -0,0 +1,9 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-feature Lifetimes
+
+struct NE: ~Escapable {
+
+}
+
+func bar1(_ closure: @_lifetime(copy ne) (_ ne: NE) -> NE, _ ne: NE) {
+  _ = closure(ne)
+}