Add an @_addressableForDependencies type attribute.

This attribute makes it so that a parameter of the annotated type, as well as
any type structurally containing that type as a field, becomes passed as
if `@_addressable` if the return value of the function has a dependency on
the parameter. This allows nonescapable values to take interior pointers into
such types.

Author: jckarter

include/swift/AST/DeclAttr.def

@@ -524,18 +524,22 @@ DECL_ATTR(lifetime, Lifetime,
   161)
 
 SIMPLE_DECL_ATTR(_addressableSelf, AddressableSelf,
-  OnAccessor | OnConstructor | OnFunc | OnSubscript | ABIBreakingToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove | UserInaccessible,
+  OnAccessor | OnConstructor | OnFunc | OnSubscript | ABIBreakingToAdd | ABIBreakingToRemove | APIBreakingToAdd | APIBreakingToRemove | UserInaccessible,
   162)
 
+SIMPLE_DECL_ATTR(_addressableForDependencies, AddressableForDependencies,
+  OnNominalType | ABIBreakingToAdd | ABIBreakingToRemove | APIBreakingToAdd | APIBreakingToRemove | UserInaccessible,
+  163)
+
 DECL_ATTR(safe, Safe,
   OnAbstractFunction | OnSubscript | OnVar | OnMacro | OnNominalType |
   OnExtension | OnTypeAlias | OnImport | UserInaccessible |
   ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove,
-  163)
+  164)
 
 DECL_ATTR(abi, ABI,
   OnAbstractFunction | OnVar /* will eventually add types */ | LongAttribute | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
-  164)
+  165)
 DECL_ATTR_FEATURE_REQUIREMENT(ABI, ABIAttribute)
 
 LAST_DECL_ATTR(ABI)

include/swift/AST/DiagnosticsSema.def

@@ -313,6 +313,10 @@ ERROR(addressable_not_enabled,none,
       "@_addressable is an experimental feature", ())
 ERROR(addressableSelf_not_on_method,none,
       "@_addressableSelf cannot be applied to non-member declarations", ())
+ERROR(addressable_types_not_enabled,none,
+      "@_addressableForDependencies is an experimental feature", ())
+ERROR(class_cannot_be_addressable_for_dependencies,none,
+      "a class cannot be @_addressableForDependencies", ())
 
 ERROR(unsupported_closure_attr,none,
       "%select{attribute |}0 '%1' is not supported on a closure",

include/swift/Basic/Features.def

@@ -434,6 +434,7 @@ EXPERIMENTAL_FEATURE(CoroutineAccessorsAllocateInCallee, false)
 EXPERIMENTAL_FEATURE(GenerateForceToMainActorThunks, false)
 
 EXPERIMENTAL_FEATURE(AddressableParameters, true)
+EXPERIMENTAL_FEATURE(AddressableTypes, true)
 
 /// Allow the @abi attribute.
 SUPPRESSIBLE_EXPERIMENTAL_FEATURE(ABIAttribute, true)

include/swift/SIL/AbstractionPattern.h

@@ -1558,6 +1558,16 @@ class AbstractionPattern {
   /// for one of its parameters.
   ParameterTypeFlags getFunctionParamFlags(unsigned index) const;
 
+  /// Given that the value being abstracted is a function type, return whether
+  /// the indicated parameter should be treated as addressable, meaning
+  /// calls should preserve the in-memory address of the argument for as
+  /// long as any dependencies may live.
+  ///
+  /// 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;
+
   /// Given that the value being abstracted is a function type, and that
   /// this is not an opaque abstraction pattern, return the number of
   /// parameters in the pattern.

include/swift/SIL/TypeLowering.h

@@ -176,6 +176,17 @@ enum HasPack_t : bool {
   HasPack = true,
 };
 
+/// Is the type addressable-for-dependencies?
+///
+/// Values of an addressable-for-dependency type are passed indirectly into
+/// functions that specify a return value lifetime dependency on the value.
+/// This allows the dependent value to safely contain pointers to the in-memory
+/// representation of the source of the dependency.
+enum IsAddressableForDependencies_t : bool {
+  IsNotAddressableForDependencies = false,
+  IsAddressableForDependencies = true,
+};
+
 /// Extended type information used by SIL.
 class TypeLowering {
 public:
@@ -184,15 +195,16 @@ class TypeLowering {
     //
     // clang-format off
     enum : unsigned {
-      NonTrivialFlag             = 1 << 0,
-      NonFixedABIFlag            = 1 << 1,
-      AddressOnlyFlag            = 1 << 2,
-      ResilientFlag              = 1 << 3,
-      TypeExpansionSensitiveFlag = 1 << 4,
-      InfiniteFlag               = 1 << 5,
-      HasRawPointerFlag          = 1 << 6,
-      LexicalFlag                = 1 << 7,
-      HasPackFlag                = 1 << 8,
+      NonTrivialFlag                 = 1 << 0,
+      NonFixedABIFlag                = 1 << 1,
+      AddressOnlyFlag                = 1 << 2,
+      ResilientFlag                  = 1 << 3,
+      TypeExpansionSensitiveFlag     = 1 << 4,
+      InfiniteFlag                   = 1 << 5,
+      HasRawPointerFlag              = 1 << 6,
+      LexicalFlag                    = 1 << 7,
+      HasPackFlag                    = 1 << 8,
+      AddressableForDependenciesFlag = 1 << 9,
     };
     // clang-format on
 
@@ -209,15 +221,17 @@ class TypeLowering {
         IsTypeExpansionSensitive_t isTypeExpansionSensitive =
             IsNotTypeExpansionSensitive,
         HasRawPointer_t hasRawPointer = DoesNotHaveRawPointer,
-        IsLexical_t isLexical = IsNotLexical, HasPack_t hasPack = HasNoPack)
+        IsLexical_t isLexical = IsNotLexical, HasPack_t hasPack = HasNoPack,
+        IsAddressableForDependencies_t isAFD = IsAddressableForDependencies)
         : Flags((isTrivial ? 0U : NonTrivialFlag) |
                 (isFixedABI ? 0U : NonFixedABIFlag) |
                 (isAddressOnly ? AddressOnlyFlag : 0U) |
                 (isResilient ? ResilientFlag : 0U) |
                 (isTypeExpansionSensitive ? TypeExpansionSensitiveFlag : 0U) |
                 (hasRawPointer ? HasRawPointerFlag : 0U) |
                 (isLexical ? LexicalFlag : 0U) |
-                (hasPack ? HasPackFlag : 0U)) {}
+                (hasPack ? HasPackFlag : 0U) |
+                (isAFD ? AddressableForDependenciesFlag : 0U)) {}
 
     constexpr bool operator==(RecursiveProperties p) const {
       return Flags == p.Flags;
@@ -227,6 +241,12 @@ class TypeLowering {
       return {IsTrivial, IsFixedABI, IsNotAddressOnly, IsNotResilient};
     }
 
+    static constexpr RecursiveProperties forTrivialOpaque() {
+      return {IsTrivial, IsFixedABI, IsNotAddressOnly, IsNotResilient,
+              IsNotTypeExpansionSensitive, HasRawPointer, IsNotLexical,
+              HasNoPack, IsAddressableForDependencies};
+    }
+
     static constexpr RecursiveProperties forRawPointer() {
       return {IsTrivial, IsFixedABI, IsNotAddressOnly, IsNotResilient,
               IsNotTypeExpansionSensitive, HasRawPointer};
@@ -239,11 +259,14 @@ class TypeLowering {
 
     static constexpr RecursiveProperties forOpaque() {
       return {IsNotTrivial, IsNotFixedABI, IsAddressOnly, IsNotResilient,
-              IsNotTypeExpansionSensitive, DoesNotHaveRawPointer, IsLexical, HasNoPack};
+              IsNotTypeExpansionSensitive, HasRawPointer, IsLexical,
+              HasNoPack, IsAddressableForDependencies};
     }
 
     static constexpr RecursiveProperties forResilient() {
-      return {IsTrivial, IsFixedABI, IsNotAddressOnly, IsResilient};
+      return {IsTrivial, IsFixedABI, IsNotAddressOnly, IsResilient,
+              IsNotTypeExpansionSensitive, HasRawPointer, IsNotLexical,
+              HasNoPack, IsAddressableForDependencies};
     }
 
     void addSubobject(RecursiveProperties other) {
@@ -278,6 +301,10 @@ class TypeLowering {
     HasPack_t isOrContainsPack() const {
       return HasPack_t((Flags & HasPackFlag) != 0);
     }
+    IsAddressableForDependencies_t isAddressableForDependencies() const {
+      return IsAddressableForDependencies_t(
+                                (Flags & AddressableForDependenciesFlag) != 0);
+    }
 
     void setNonTrivial() { Flags |= NonTrivialFlag; }
     void setIsOrContainsRawPointer() { Flags |= HasRawPointerFlag; }
@@ -294,6 +321,9 @@ class TypeLowering {
       Flags = (Flags & ~LexicalFlag) | (isLexical ? LexicalFlag : 0);
     }
     void setHasPack() { Flags |= HasPackFlag; }
+    void setAddressableForDependencies() {
+      Flags |= AddressableForDependenciesFlag;
+    }
   };
 
 private:

lib/AST/ASTDumper.cpp

@@ -3760,6 +3760,7 @@ class PrintAttribute : public AttributeVisitor<PrintAttribute, void, StringRef>,
 
   TRIVIAL_ATTR_PRINTER(Actor, actor)
   TRIVIAL_ATTR_PRINTER(AddressableSelf, _addressableSelf)
+  TRIVIAL_ATTR_PRINTER(AddressableForDependencies, _addressableForDependencies)
   TRIVIAL_ATTR_PRINTER(AlwaysEmitConformanceMetadata,
                        always_emit_conformance_metadata)
   TRIVIAL_ATTR_PRINTER(AlwaysEmitIntoClient, always_emit_into_client)

lib/AST/FeatureSet.cpp

@@ -308,6 +308,14 @@ static bool usesFeatureAddressableParameters(Decl *d) {
   return false;
 }
 
+static bool usesFeatureAddressableTypes(Decl *d) {
+  if (d->getAttrs().hasAttribute<AddressableForDependenciesAttr>()) {
+    return true;
+  }
+  
+  return false;
+}
+
 UNINTERESTING_FEATURE(IsolatedAny2)
 UNINTERESTING_FEATURE(GlobalActorIsolatedTypesUsability)
 UNINTERESTING_FEATURE(ObjCImplementation)

lib/ASTGen/Sources/ASTGen/DeclAttrs.swift

@@ -191,6 +191,7 @@ extension ASTGenVisitor {
 
       // Simple attributes.
       case .addressableSelf,
+        .addressableForDependencies,
         .alwaysEmitConformanceMetadata,
         .alwaysEmitIntoClient,
         .atReasync,

lib/SIL/IR/AbstractionPattern.cpp

@@ -1649,6 +1649,72 @@ AbstractionPattern::getFunctionParamFlags(unsigned index) const {
            .getParameterFlags();
 }
 
+bool
+AbstractionPattern::isFunctionParamAddressable(TypeConverter &TC,
+                                               unsigned index) 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 true;
+
+  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();
+    
+    if (type->isTypeParameter() || type->is<ArchetypeType>()) {
+      // If the function abstraction pattern is completely opaque, assume we
+      // may need to preserve the address for dependencies.
+      return true;    
+    }
+  
+    auto fnTy = cast<AnyFunctionType>(getType());
+  
+    // The parameter might directly be marked addressable.
+    if (fnTy.getParams()[index].getParameterFlags().isAddressable()) {
+      return true;
+    }
+    
+    // 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();
+      }
+    }
+    
+    return false;
+  }
+  }
+  llvm_unreachable("bad kind");
+}
+
 unsigned AbstractionPattern::getNumFunctionParams() const {
   return cast<AnyFunctionType>(getType()).getParams().size();
 }