Fix getLoweredOwnership() for setters of ~Escapable types

lib/AST/LifetimeDependence.cpp

@@ -186,15 +186,18 @@ void LifetimeDependenceInfo::Profile(llvm::FoldingSetNodeID &ID) const {
   }
 }
 
-// Warning: this is incorrect for Setter 'newValue' parameters. It should only
-// be called for a Setter's 'self'.
-static ValueOwnership getLoweredOwnership(AbstractFunctionDecl *afd) {
+static ValueOwnership getLoweredOwnership(ParamDecl *param,
+                                          AbstractFunctionDecl *afd) {
   if (isa<ConstructorDecl>(afd)) {
     return ValueOwnership::Owned;
   }
   if (auto *ad = dyn_cast<AccessorDecl>(afd)) {
-    if (ad->getAccessorKind() == AccessorKind::Set ||
-        isYieldingMutableAccessor(ad->getAccessorKind())) {
+    if (ad->getAccessorKind() == AccessorKind::Set) {
+      return param->isSelfParameter() ? ValueOwnership::InOut
+                                      : ValueOwnership::Owned;
+    }
+    if (isYieldingMutableAccessor(ad->getAccessorKind())) {
+      assert(param->isSelfParameter());
       return ValueOwnership::InOut;
     }
   }
@@ -565,17 +568,25 @@ class LifetimeDependenceChecker {
     }
   }
 
-  bool isCompatibleWithOwnership(ParsedLifetimeDependenceKind kind, Type type,
-                                 ValueOwnership loweredOwnership,
+  bool isCompatibleWithOwnership(ParsedLifetimeDependenceKind kind,
+                                 ParamDecl *param,
                                  bool isInterfaceFile = false) const {
     if (kind == ParsedLifetimeDependenceKind::Inherit) {
       return true;
     }
+
+    auto *afd = cast<AbstractFunctionDecl>(decl);
+    auto paramType = param->getTypeInContext();
+    auto ownership = param->getValueOwnership();
+    auto loweredOwnership = ownership != ValueOwnership::Default
+                                ? ownership
+                                : getLoweredOwnership(param, afd);
+
     if (kind == ParsedLifetimeDependenceKind::Borrow) {
       // An owned/consumed BitwiseCopyable value can be effectively borrowed
       // because its lifetime can be indefinitely extended.
-      if (loweredOwnership == ValueOwnership::Owned
-          && isBitwiseCopyable(type, ctx)) {
+      if (loweredOwnership == ValueOwnership::Owned &&
+          isBitwiseCopyable(paramType, ctx)) {
         return true;
       }
       if (isInterfaceFile) {
@@ -588,21 +599,23 @@ class LifetimeDependenceChecker {
     return loweredOwnership == ValueOwnership::InOut;
   }
 
-  bool isCompatibleWithOwnership(LifetimeDependenceKind kind, Type type,
-                                 ValueOwnership ownership) const {
-    auto *afd = cast<AbstractFunctionDecl>(decl);
+  bool isCompatibleWithOwnership(LifetimeDependenceKind kind,
+                                 ParamDecl *param) const {
     if (kind == LifetimeDependenceKind::Inherit) {
       return true;
     }
+
+    auto *afd = cast<AbstractFunctionDecl>(decl);
+    auto paramType = param->getTypeInContext();
+    auto ownership = param->getValueOwnership();
+    auto loweredOwnership = ownership != ValueOwnership::Default
+                                ? ownership
+                                : getLoweredOwnership(param, afd);
     // Lifetime dependence always propagates through temporary BitwiseCopyable
     // values, even if the dependence is scoped.
-    if (isBitwiseCopyable(type, ctx)) {
+    if (isBitwiseCopyable(paramType, ctx)) {
       return true;
     }
-    auto loweredOwnership = ownership != ValueOwnership::Default
-                                ? ownership
-                                : getLoweredOwnership(afd);
-
     assert(kind == LifetimeDependenceKind::Scope);
     return loweredOwnership == ValueOwnership::Shared ||
            loweredOwnership == ValueOwnership::InOut;
@@ -690,7 +703,7 @@ class LifetimeDependenceChecker {
     auto ownership = paramDecl->getValueOwnership();
     auto loweredOwnership = ownership != ValueOwnership::Default
                                 ? ownership
-                                : getLoweredOwnership(afd);
+                                : getLoweredOwnership(paramDecl, afd);
 
     switch (parsedLifetimeKind) {
     case ParsedLifetimeDependenceKind::Default: {
@@ -717,9 +730,7 @@ class LifetimeDependenceChecker {
     case ParsedLifetimeDependenceKind::Inout: {
       // @lifetime(borrow x) is valid only for borrowing parameters.
       // @lifetime(&x) is valid only for inout parameters.
-      auto loweredOwnership = ownership != ValueOwnership::Default
-        ? ownership : getLoweredOwnership(afd);
-      if (isCompatibleWithOwnership(parsedLifetimeKind, type, loweredOwnership,
+      if (isCompatibleWithOwnership(parsedLifetimeKind, paramDecl,
                                     isInterfaceFile())) {
         return LifetimeDependenceKind::Scope;
       }
@@ -1039,8 +1050,7 @@ class LifetimeDependenceChecker {
     }
     // Infer based on ownership if possible for either explicit accessors or
     // methods as long as they pass preceding ambiguity checks.
-    auto kind = inferLifetimeDependenceKind(
-        selfTypeInContext, afd->getImplicitSelfDecl()->getValueOwnership());
+    auto kind = inferLifetimeDependenceKind(afd->getImplicitSelfDecl());
     if (!kind) {
       // Special diagnostic for an attempt to depend on a consuming parameter.
       diagnose(returnLoc,
@@ -1054,19 +1064,21 @@ class LifetimeDependenceChecker {
   // Infer the kind of dependence that makes sense for reading or writing a
   // stored property (for getters or initializers).
   std::optional<LifetimeDependenceKind>
-  inferLifetimeDependenceKind(Type sourceType, ValueOwnership ownership) {
+  inferLifetimeDependenceKind(ParamDecl *param) {
     auto *afd = cast<AbstractFunctionDecl>(decl);
-    if (!sourceType->isEscapable()) {
+    Type paramType = param->getTypeInContext();
+    ValueOwnership ownership = param->getValueOwnership();
+    if (!paramType->isEscapable()) {
       return LifetimeDependenceKind::Inherit;
     }
     // Lifetime dependence always propagates through temporary BitwiseCopyable
     // values, even if the dependence is scoped.
-    if (isBitwiseCopyable(sourceType, ctx)) {
+    if (isBitwiseCopyable(paramType, ctx)) {
       return LifetimeDependenceKind::Scope;
     }
     auto loweredOwnership = ownership != ValueOwnership::Default
                                 ? ownership
-                                : getLoweredOwnership(afd);
+                                : getLoweredOwnership(param, afd);
     // It is impossible to depend on a consumed Escapable value (unless it is
     // BitwiseCopyable as checked above).
     if (loweredOwnership == ValueOwnership::Owned) {
@@ -1114,8 +1126,7 @@ class LifetimeDependenceChecker {
         return;
       }
       // A single Escapable parameter must be borrowed.
-      auto kind = inferLifetimeDependenceKind(paramTypeInContext,
-                                              param->getValueOwnership());
+      auto kind = inferLifetimeDependenceKind(param);
       if (!kind) {
         diagnose(returnLoc,
                  diag::lifetime_dependence_cannot_infer_scope_ownership,
@@ -1172,9 +1183,7 @@ class LifetimeDependenceChecker {
       return;
     }
     auto kind = LifetimeDependenceKind::Scope;
-    auto paramOwnership = param->getValueOwnership();
-    if (!isCompatibleWithOwnership(kind, paramTypeInContext, paramOwnership))
-    {
+    if (!isCompatibleWithOwnership(kind, param)) {
       diagnose(returnLoc,
                diag::lifetime_dependence_cannot_infer_scope_ownership,
                param->getParameterName().str(), diagnosticQualifier());
@@ -1207,8 +1216,7 @@ class LifetimeDependenceChecker {
         }
       }
 
-      candidateLifetimeKind =
-          inferLifetimeDependenceKind(paramTypeInContext, paramOwnership);
+      candidateLifetimeKind = inferLifetimeDependenceKind(param);
       if (!candidateLifetimeKind) {
         continue;
       }
@@ -1383,11 +1391,9 @@ class LifetimeDependenceChecker {
         }
       }
     }
-    auto *afd = cast<AbstractFunctionDecl>(decl);
     // Either a Get or Modify without any wrapped accessor. Handle these like a
     // read of the stored property.
-    return inferLifetimeDependenceKind(
-      selfTypeInContext, afd->getImplicitSelfDecl()->getValueOwnership());
+    return inferLifetimeDependenceKind(accessor->getImplicitSelfDecl());
   }
 
   // Infer 'inout' parameter dependency when the only parameter is

test/Sema/lifetime_attr.swift

@@ -125,4 +125,20 @@ struct Wrapper : ~Escapable {
     nonmutating _modify {// expected-error{{lifetime-dependent parameter 'self' must be 'inout'}}
     }
   }
+
+  var otherNE: NE {
+    @_lifetime(copy self)
+    get {
+      _ne
+    }
+    @_lifetime(self: borrow newValue)
+    set {
+      self._ne = newValue
+    }
+    @_lifetime(&self)
+    _modify {
+      yield &self._ne
+    }
+  }
 }
+