Merge pull request swiftlang#80438 from jckarter/substitute-away-escapable-lifetime-deps

Type substitution eliminates dependencies with Escapable targets.

Author: jckarter

include/swift/AST/LifetimeDependence.h

@@ -354,6 +354,18 @@ class LifetimeDependenceInfo {
 std::optional<LifetimeDependenceInfo>
 getLifetimeDependenceFor(ArrayRef<LifetimeDependenceInfo> lifetimeDependencies,
                          unsigned index);
+
+/// Helper to remove lifetime dependencies whose target references an
+/// Escapable type.
+///
+/// Returns true if the set of output lifetime dependencies is different from
+/// the set of input lifetime dependencies, false if there was no change.
+bool
+filterEscapableLifetimeDependencies(GenericSignature sig,
+        ArrayRef<LifetimeDependenceInfo> inputs,
+        SmallVectorImpl<LifetimeDependenceInfo> &outputs,
+        llvm::function_ref<Type (unsigned targetIndex)> getSubstTargetType);
+
 } // namespace swift
 
 #endif

include/swift/AST/TypeTransform.h

@@ -20,6 +20,7 @@
 
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/SILLayout.h"
+#include "swift/AST/LifetimeDependence.h"
 
 namespace swift {
 
@@ -333,11 +334,37 @@ case TypeKind::Id:
         transErrorResult = result;
       }
 
-      if (!changed) return t;
+      if (!changed) {
+        return t;
+      }
+      
+      // Lifetime dependencies get eliminated if their target type was
+      // substituted with an escapable type.
+      auto extInfo = fnTy->getExtInfo();
+      if (!extInfo.getLifetimeDependencies().empty()) {
+        SmallVector<LifetimeDependenceInfo, 2> substDependenceInfos;
+        bool didRemoveLifetimeDependencies
+          = filterEscapableLifetimeDependencies(GenericSignature(),
+                                              extInfo.getLifetimeDependencies(),
+                                              substDependenceInfos,
+                                              [&](unsigned targetIndex) {
+            if (targetIndex >= transInterfaceParams.size()) {
+              // Target is a return type.
+              return transInterfaceResults[targetIndex - transInterfaceParams.size()]
+                .getInterfaceType();
+            } else {
+              // Target is a parameter.
+              return transInterfaceParams[targetIndex].getInterfaceType();
+            }
+          });
+        if (didRemoveLifetimeDependencies) {
+          extInfo = extInfo.withLifetimeDependencies(substDependenceInfos);
+        }
+      }
 
       return SILFunctionType::get(
           fnTy->getInvocationGenericSignature(),
-          fnTy->getExtInfo(),
+          extInfo,
           fnTy->getCoroutineKind(),
           fnTy->getCalleeConvention(),
           transInterfaceParams,
@@ -809,7 +836,7 @@ case TypeKind::Id:
       }
 
       // Transform result type.
-      auto resultTy = doIt(function->getResult(), pos);
+      Type resultTy = doIt(function->getResult(), pos);
       if (!resultTy)
         return Type();
 
@@ -876,8 +903,45 @@ case TypeKind::Id:
         return GenericFunctionType::get(
             genericSig, substParams, resultTy, extInfo);
       }
+      
+      if (isUnchanged) {
+        return t;
+      }
 
-      if (isUnchanged) return t;
+      // Substitution may have replaced parameter or return types that had
+      // lifetime dependencies with Escapable types, which render those
+      // dependencies inactive.
+      if (extInfo && !extInfo->getLifetimeDependencies().empty()) {
+        SmallVector<LifetimeDependenceInfo, 2> substDependenceInfos;
+        bool didRemoveLifetimeDependencies
+          = filterEscapableLifetimeDependencies(GenericSignature(),
+                                                extInfo->getLifetimeDependencies(),
+                                                substDependenceInfos,
+                                                [&](unsigned targetIndex) {
+            // Traverse potentially-curried function types.
+            ArrayRef<AnyFunctionType::Param> params = substParams;
+            auto result = resultTy;
+            while (targetIndex >= params.size()) {
+              if (auto curriedTy = result->getAs<AnyFunctionType>()) {
+                targetIndex -= params.size();
+                params = curriedTy->getParams();
+                result = curriedTy->getResult();
+                continue;
+              } else {
+                // The last lifetime dependency targets the result at the end
+                // of the curried chain.
+                ASSERT(targetIndex == params.size()
+                       && "invalid lifetime dependence target");
+                return result;
+              }
+            }
+            return params[targetIndex].getParameterType();
+          });
+
+        if (didRemoveLifetimeDependencies) {
+          extInfo = extInfo->withLifetimeDependencies(substDependenceInfos);
+        }
+      }
 
       return FunctionType::get(substParams, resultTy, extInfo);
     }

include/swift/AST/Types.h

@@ -688,6 +688,10 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Returns true if this contextual type satisfies a conformance to Escapable.
   bool isEscapable();
 
+  /// Returns true if this type satisfies a conformance to Escapable in the
+  /// given generic signature.
+  bool isEscapable(GenericSignature sig);
+
   /// Returns true if this contextual type is (Escapable && !isNoEscape).
   bool mayEscape() { return !isNoEscape() && isEscapable(); }
 
@@ -3566,6 +3570,16 @@ class AnyFunctionType : public TypeBase {
     }
     Bits.AnyFunctionType.NumParams = NumParams;
     assert(Bits.AnyFunctionType.NumParams == NumParams && "Params dropped!");
+    
+    if (Info) {
+      unsigned maxLifetimeTarget = NumParams + 1;
+      if (auto outputFn = Output->getAs<AnyFunctionType>()) {
+        maxLifetimeTarget += outputFn->getNumParams();
+      }
+      for (auto &dep : Info->getLifetimeDependencies()) {
+        assert(dep.getTargetIndex() < maxLifetimeTarget);
+      }
+    }
   }
 
 public:
@@ -5170,6 +5184,16 @@ class SILFunctionType final
       const ASTContext &ctx,
       ProtocolConformanceRef witnessMethodConformance =
           ProtocolConformanceRef());
+          
+  /// Given an existing ExtInfo, and a set of interface parameters and results
+  /// destined for a new SILFunctionType, return a new ExtInfo with only the
+  /// lifetime dependencies relevant after substitution.
+  static ExtInfo
+  getSubstLifetimeDependencies(GenericSignature genericSig,
+                               ExtInfo origExtInfo,
+                               ArrayRef<SILParameterInfo> params,
+                               ArrayRef<SILYieldInfo> yields,
+                               ArrayRef<SILResultInfo> results);
 
   /// Return a structurally-identical function type with a slightly tweaked
   /// ExtInfo.

include/swift/SIL/AbstractionPattern.h

@@ -1566,7 +1566,9 @@ class AbstractionPattern {
   /// This may be true either because the type is structurally addressable for
   /// dependencies, or because it was explicitly marked as `@_addressable`
   /// in its declaration.
-  bool isFunctionParamAddressable(TypeConverter &TC, unsigned index) const;
+  bool isFunctionParamAddressable(unsigned index) const;
+  
+  ArrayRef<LifetimeDependenceInfo> getLifetimeDependencies() const;
 
   /// Given that the value being abstracted is a function type, and that
   /// this is not an opaque abstraction pattern, return the number of

lib/AST/ConformanceLookup.cpp

@@ -916,3 +916,11 @@ bool TypeBase::isEscapable() {
   auto canType = preprocessTypeForInvertibleQuery(this);
   return conformsToInvertible(canType, InvertibleProtocolKind::Escapable);
 }
+
+bool TypeBase::isEscapable(GenericSignature sig) {
+  Type contextTy = this;
+  if (sig) {
+    contextTy = sig.getGenericEnvironment()->mapTypeIntoContext(contextTy);
+  }
+  return contextTy->isEscapable();
+}

lib/AST/LifetimeDependence.cpp

@@ -48,6 +48,32 @@ getLifetimeDependenceFor(ArrayRef<LifetimeDependenceInfo> lifetimeDependencies,
   return std::nullopt;
 }
 
+bool
+filterEscapableLifetimeDependencies(GenericSignature sig,
+        ArrayRef<LifetimeDependenceInfo> inputs,
+        SmallVectorImpl<LifetimeDependenceInfo> &outputs,
+        llvm::function_ref<Type (unsigned targetIndex)> getSubstTargetType) {
+  bool didRemoveLifetimeDependencies = false;
+
+  for (auto &depInfo : inputs) {
+    auto targetIndex = depInfo.getTargetIndex();
+    Type substTy = getSubstTargetType(targetIndex);
+    
+    // Drop the dependency if the target type is Escapable.
+    if (sig || !substTy->hasTypeParameter()) {
+      if (substTy->isEscapable(sig)) {
+        didRemoveLifetimeDependencies = true;
+        continue;
+      }
+    }
+    
+    // Otherwise, keep the dependency.
+    outputs.push_back(depInfo);
+  }
+
+  return didRemoveLifetimeDependencies;
+}
+
 std::string LifetimeDependenceInfo::getString() const {
   std::string lifetimeDependenceString = "@lifetime(";
   auto addressable = getAddressableIndices();

lib/SIL/IR/AbstractionPattern.cpp

@@ -1650,8 +1650,7 @@ AbstractionPattern::getFunctionParamFlags(unsigned index) const {
 }
 
 bool
-AbstractionPattern::isFunctionParamAddressable(TypeConverter &TC,
-                                               unsigned index) const {
+AbstractionPattern::isFunctionParamAddressable(unsigned index) const {
   switch (getKind()) {
   case Kind::Invalid:
   case Kind::Tuple:
@@ -1661,7 +1660,7 @@ AbstractionPattern::isFunctionParamAddressable(TypeConverter &TC,
   case Kind::OpaqueDerivativeFunction:
     // If the function abstraction pattern is completely opaque, assume we
     // may need to preserve the address for dependencies.
-    return true;
+    return false;
 
   case Kind::ClangType:
   case Kind::ObjCCompletionHandlerArgumentsType:
@@ -1685,29 +1684,47 @@ AbstractionPattern::isFunctionParamAddressable(TypeConverter &TC,
     auto fnTy = cast<AnyFunctionType>(getType());
 
     // The parameter might directly be marked addressable.
-    if (fnTy.getParams()[index].getParameterFlags().isAddressable()) {
-      return true;
-    }
+    return fnTy.getParams()[index].getParameterFlags().isAddressable();
+  }
+  }
+  llvm_unreachable("bad kind");
+}
+
+ArrayRef<LifetimeDependenceInfo>
+AbstractionPattern::getLifetimeDependencies() const {
+  switch (getKind()) {
+  case Kind::Invalid:
+  case Kind::Tuple:
+    llvm_unreachable("not any kind of function!");
+  case Kind::Opaque:
+  case Kind::OpaqueFunction:
+  case Kind::OpaqueDerivativeFunction:
+    // If the function abstraction pattern is completely opaque, assume we
+    // may need to preserve the address for dependencies.
+    return {};
+
+  case Kind::ClangType:
+  case Kind::ObjCCompletionHandlerArgumentsType:
+  case Kind::CurriedObjCMethodType:
+  case Kind::PartialCurriedObjCMethodType:
+  case Kind::ObjCMethodType:
+  case Kind::CFunctionAsMethodType:
+  case Kind::CurriedCFunctionAsMethodType:
+  case Kind::PartialCurriedCFunctionAsMethodType:
+  case Kind::CXXMethodType:
+  case Kind::CurriedCXXMethodType:
+  case Kind::PartialCurriedCXXMethodType:
+  case Kind::Type:
+  case Kind::Discard: {
+    auto type = getType();
 
-    // The parameter could be of a type that is addressable for dependencies,
-    // in which case it becomes addressable when a return has a scoped
-    // dependency on it.
-    for (auto &dep : fnTy->getLifetimeDependencies()) {
-      auto scoped = dep.getScopeIndices();
-      if (!scoped) {
-        continue;
-      }
-      
-      if (scoped->contains(index)) {
-        auto paramTy = getFunctionParamType(index);
-        
-        return TC.getTypeLowering(paramTy, paramTy.getType(),
-                                  TypeExpansionContext::minimal())
-          .getRecursiveProperties().isAddressableForDependencies();
-      }
+    if (type->isTypeParameter() || type->is<ArchetypeType>()) {
+      return {};
     }
-    
-    return false;
+  
+    auto fnTy = cast<AnyFunctionType>(getType());
+  
+    return fnTy->getExtInfo().getLifetimeDependencies();
   }
   }
   llvm_unreachable("bad kind");