Add support for @Lifetime(&arg)

Author: meg-gupta

include/swift/AST/ASTBridging.h

@@ -1088,6 +1088,7 @@ enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedParsedLifetimeDependenceKind {
   BridgedParsedLifetimeDependenceKindDefault,
   BridgedParsedLifetimeDependenceKindScope,
   BridgedParsedLifetimeDependenceKindInherit,
+  BridgedParsedLifetimeDependenceKindInout
 };
 
 class BridgedLifetimeDescriptor {

include/swift/AST/DiagnosticsSema.def

@@ -8171,6 +8171,9 @@ ERROR(lifetime_dependence_invalid_inherit_escapable_type, none,
 ERROR(lifetime_dependence_cannot_use_parsed_borrow_consuming, none,
       "invalid use of borrow dependence with consuming ownership",
       ())
+ERROR(lifetime_dependence_cannot_use_default_escapable_consuming, none,
+      "invalid lifetime dependence on an Escapable value with %0 ownership",
+      (StringRef))
 ERROR(lifetime_dependence_cannot_use_parsed_borrow_inout, none,
       "invalid use of borrow dependence on the same inout parameter",
       ())

include/swift/AST/LifetimeDependence.h

@@ -39,7 +39,8 @@ class SILResultInfo;
 enum class ParsedLifetimeDependenceKind : uint8_t {
   Default = 0,
   Scope,
-  Inherit // Only used with deserialized decls
+  Inherit, // Only used with deserialized decls
+  Inout
 };
 
 enum class LifetimeDependenceKind : uint8_t { Inherit = 0, Scope };

lib/AST/Bridging/DeclAttributeBridging.cpp

@@ -474,6 +474,8 @@ unbridged(BridgedParsedLifetimeDependenceKind kind) {
     return swift::ParsedLifetimeDependenceKind::Scope;
   case BridgedParsedLifetimeDependenceKindInherit:
     return swift::ParsedLifetimeDependenceKind::Inherit;
+  case BridgedParsedLifetimeDependenceKindInout:
+    return swift::ParsedLifetimeDependenceKind::Inout;
   }
   llvm_unreachable("unhandled enum value");
 }

lib/AST/LifetimeDependence.cpp

@@ -79,6 +79,20 @@ filterEscapableLifetimeDependencies(GenericSignature sig,
   return didRemoveLifetimeDependencies;
 }
 
+StringRef
+getNameForParsedLifetimeDependenceKind(ParsedLifetimeDependenceKind kind) {
+  switch (kind) {
+  case ParsedLifetimeDependenceKind::Scope:
+    return "borrow";
+  case ParsedLifetimeDependenceKind::Inherit:
+    return "copy";
+  case ParsedLifetimeDependenceKind::Inout:
+    return "inout";
+  default:
+    return "";
+  }
+}
+
 std::string LifetimeDependenceInfo::getString() const {
   std::string lifetimeDependenceString = "@lifetime(";
   auto addressable = getAddressableIndices();
@@ -451,26 +465,43 @@ class LifetimeDependenceChecker {
     }
   }
 
-  bool isCompatibleWithOwnership(LifetimeDependenceKind kind, Type type,
+  bool isCompatibleWithOwnership(ParsedLifetimeDependenceKind kind, Type type,
                                  ValueOwnership ownership) const {
-    if (kind == LifetimeDependenceKind::Inherit) {
+    if (kind == ParsedLifetimeDependenceKind::Inherit) {
       return true;
     }
     // Lifetime dependence always propagates through temporary BitwiseCopyable
     // values, even if the dependence is scoped.
     if (isBitwiseCopyable(type, ctx)) {
       return true;
     }
-    assert(kind == LifetimeDependenceKind::Scope);
     auto loweredOwnership = ownership != ValueOwnership::Default
       ? ownership : getLoweredOwnership(afd);
 
-    if (loweredOwnership == ValueOwnership::InOut ||
-        loweredOwnership == ValueOwnership::Shared) {
+    if (kind == ParsedLifetimeDependenceKind::Scope) {
+      return loweredOwnership == ValueOwnership::Shared;
+    }
+    assert(kind == ParsedLifetimeDependenceKind::Inout);
+    return loweredOwnership == ValueOwnership::InOut;
+  }
+
+  bool isCompatibleWithOwnership(LifetimeDependenceKind kind, Type type,
+                                 ValueOwnership ownership) const {
+    if (kind == LifetimeDependenceKind::Inherit) {
       return true;
     }
-    assert(loweredOwnership == ValueOwnership::Owned);
-    return false;
+    // Lifetime dependence always propagates through temporary BitwiseCopyable
+    // values, even if the dependence is scoped.
+    if (isBitwiseCopyable(type, ctx)) {
+      return true;
+    }
+    auto loweredOwnership = ownership != ValueOwnership::Default
+                                ? ownership
+                                : getLoweredOwnership(afd);
+
+    assert(kind == LifetimeDependenceKind::Scope);
+    return loweredOwnership == ValueOwnership::Shared ||
+           loweredOwnership == ValueOwnership::InOut;
   }
 
   struct TargetDeps {
@@ -546,44 +577,51 @@ class LifetimeDependenceChecker {
   std::optional<LifetimeDependenceKind>
   getDependenceKindFromDescriptor(LifetimeDescriptor descriptor,
                                   ParamDecl *decl) {
-    auto loc = descriptor.getLoc();
-
-    auto ownership = decl->getValueOwnership();
+    auto loc = decl->getLoc();
     auto type = decl->getTypeInContext();
+    auto parsedLifetimeKind = descriptor.getParsedLifetimeDependenceKind();
+    auto ownership = decl->getValueOwnership();
+    auto loweredOwnership = ownership != ValueOwnership::Default
+                                ? ownership
+                                : getLoweredOwnership(afd);
 
-    LifetimeDependenceKind kind;
-    switch (descriptor.getParsedLifetimeDependenceKind()) {
-    case ParsedLifetimeDependenceKind::Default:
+    if (parsedLifetimeKind == ParsedLifetimeDependenceKind::Default) {
       if (type->isEscapable()) {
-        kind = LifetimeDependenceKind::Scope;
+        if (loweredOwnership == ValueOwnership::Shared ||
+            loweredOwnership == ValueOwnership::InOut) {
+          return LifetimeDependenceKind::Scope;
+        } else {
+          diagnose(
+              loc,
+              diag::lifetime_dependence_cannot_use_default_escapable_consuming,
+              getOwnershipSpelling(loweredOwnership));
+          return std::nullopt;
+        }
       } else if (useLazyInference()) {
-        kind = LifetimeDependenceKind::Inherit;
-      } else {
-        diagnose(loc, diag::lifetime_dependence_cannot_infer_kind,
-                 diagnosticQualifier(), descriptor.getString());
-        return std::nullopt;
+        return LifetimeDependenceKind::Inherit;
       }
-      break;
-    case ParsedLifetimeDependenceKind::Scope:
-      kind = LifetimeDependenceKind::Scope;
-      break;
-    case ParsedLifetimeDependenceKind::Inherit:
-      kind = LifetimeDependenceKind::Inherit;
-      break;
+      diagnose(loc, diag::lifetime_dependence_cannot_infer_kind,
+               diagnosticQualifier(), descriptor.getString());
+      return std::nullopt;
     }
-    // @lifetime(borrow x) is invalid for consuming parameters.
-    if (!isCompatibleWithOwnership(kind, type, ownership)) {
+
+    // @lifetime(borrow x) is valid only for borrowing parameters.
+    // @lifetime(inout x) is valid only for inout parameters.
+    if (!isCompatibleWithOwnership(parsedLifetimeKind, type, ownership)) {
       diagnose(loc,
                diag::lifetime_dependence_cannot_use_parsed_borrow_consuming);
       return std::nullopt;
     }
     // @lifetime(copy x) is only invalid for Escapable types.
-    if (kind == LifetimeDependenceKind::Inherit && type->isEscapable()) {
+    if (parsedLifetimeKind == ParsedLifetimeDependenceKind::Inherit &&
+        type->isEscapable()) {
       diagnose(loc, diag::lifetime_dependence_invalid_inherit_escapable_type,
                descriptor.getString());
       return std::nullopt;
     }
-    return kind;
+    return parsedLifetimeKind == ParsedLifetimeDependenceKind::Inherit
+               ? LifetimeDependenceKind::Inherit
+               : LifetimeDependenceKind::Scope;
   }
 
   // Finds the ParamDecl* and its index from a LifetimeDescriptor

lib/ASTGen/Sources/ASTGen/DeclAttrs.swift

@@ -1106,6 +1106,9 @@ extension ASTGenVisitor {
     } else if let borrowExpr = node.as(BorrowExprSyntax.self) {
       lifetimeDependenceKind = .scope
       descriptorExpr = borrowExpr.expression
+    } else if let inoutExpr = node.as(InOutExprSyntax.self) {
+      lifetimeDependenceKind = .inout
+      descriptorExpr = inoutExpr.expression
     } else {
       lifetimeDependenceKind = .default
       descriptorExpr = node

lib/Parse/ParseDecl.cpp

@@ -5021,6 +5021,11 @@ ParserResult<LifetimeEntry> Parser::parseLifetimeEntry(SourceLoc loc) {
                           tok::kw_self)) {
       return ParsedLifetimeDependenceKind::Scope;
     }
+    if (Tok.is(tok::amp_prefix) &&
+        peekToken().isAny(tok::identifier, tok::integer_literal,
+                          tok::kw_self)) {
+      return ParsedLifetimeDependenceKind::Inout;
+    }
     return std::nullopt;
   };