Prototype an @_addressable attribute that puts an argument at a stable address.

jckarter · jckarter · commit 3c0b08dbcb6e · 2024-12-03T20:39:23.000-08:00
Many APIs using nonescapable types would like to vend interior pointers to their
parameter bindings, but this isn't normally always possible because of representation
changes the caller may do around the call, such as moving the value in or out of memory,
bridging or reabstracting it, etc. `@_addressable` forces the corresponding parameter
to be passed indirectly in memory, in its maximally-abstracted representation.
[TODO] If return values have a lifetime dependency on this parameter, the caller must
keep this in-memory representation alive for the duration of the dependent value's
lifetime.
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -6817,6 +6817,9 @@ class ParamDecl : public VarDecl {
 
     /// Whether or not this parameter is 'isolated'.
     IsIsolated = 1 << 2,
+    
+    /// Whether this parameter is `@_addressable`.
+    IsAddressable = 1 << 3,
 
     /// Whether or not this parameter is 'sending'.
     IsSending = 1 << 4,
@@ -7107,6 +7110,18 @@ class ParamDecl : public VarDecl {
       removeFlag(Flag::IsSending);
   }
 
+  /// Whether or not this parameter is marked with '@_addressable'.
+  bool isAddressable() const {
+    return getOptions().contains(Flag::IsAddressable);
+  }
+  
+  void setAddressable(bool value = true) {
+    if (value)
+      addFlag(Flag::IsAddressable);
+    else
+      removeFlag(Flag::IsAddressable);
+  }
+
   /// Whether or not this parameter is marked with '_const'.
   bool isCompileTimeConst() const {
     return ArgumentNameAndFlags.getInt().contains(
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -1834,6 +1834,10 @@ ERROR(attr_implementation_category_goes_on_objc_attr,none,
       "'@implementation'",
       ())
 
+ERROR(attr_not_allowed_in_c_functions,none,
+      "attribute %0 is not allowed in C function conventions",
+      (StringRef))
+
 WARNING(objc_implementation_will_become_objc,none,
         "%kind0 will become implicitly '@objc' after adopting '@objc "
         "@implementation'; add '%select{@nonobjc|final}1' to keep the current "
diff --git a/include/swift/AST/TypeAttr.def b/include/swift/AST/TypeAttr.def
@@ -64,6 +64,7 @@ SIMPLE_TYPE_ATTR(_local, Local)
 SIMPLE_TYPE_ATTR(_noMetadata, NoMetadata)
 TYPE_ATTR(_opaqueReturnTypeOf, OpaqueReturnTypeOf)
 TYPE_ATTR(isolated, Isolated)
+SIMPLE_TYPE_ATTR(_addressable, Addressable)
 
 // SIL-specific attributes
 SIMPLE_SIL_TYPE_ATTR(async, Async)
diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h
@@ -2388,7 +2388,8 @@ class ParameterTypeFlags {
     Isolated = 1 << 7,
     CompileTimeConst = 1 << 8,
     Sending = 1 << 9,
-    NumBits = 10
+    Addressable = 1 << 10,
+    NumBits = 11
   };
   OptionSet<ParameterFlags> value;
   static_assert(NumBits <= 8*sizeof(OptionSet<ParameterFlags>), "overflowed");
@@ -2403,20 +2404,21 @@ class ParameterTypeFlags {
 
   ParameterTypeFlags(bool variadic, bool autoclosure, bool nonEphemeral,
                      ParamSpecifier specifier, bool isolated, bool noDerivative,
-                     bool compileTimeConst, bool isSending)
+                     bool compileTimeConst, bool isSending, bool isAddressable)
       : value((variadic ? Variadic : 0) | (autoclosure ? AutoClosure : 0) |
               (nonEphemeral ? NonEphemeral : 0) |
               uint8_t(specifier) << SpecifierShift | (isolated ? Isolated : 0) |
               (noDerivative ? NoDerivative : 0) |
               (compileTimeConst ? CompileTimeConst : 0) |
-              (isSending ? Sending : 0)) {}
+              (isSending ? Sending : 0) |
+              (isAddressable ? Addressable : 0)) {}
 
   /// Create one from what's present in the parameter type
   inline static ParameterTypeFlags
   fromParameterType(Type paramTy, bool isVariadic, bool isAutoClosure,
                     bool isNonEphemeral, ParamSpecifier ownership,
                     bool isolated, bool isNoDerivative, bool compileTimeConst,
-                    bool isSending);
+                    bool isSending, bool isAddressable);
 
   bool isNone() const { return !value; }
   bool isVariadic() const { return value.contains(Variadic); }
@@ -2429,6 +2431,7 @@ class ParameterTypeFlags {
   bool isCompileTimeConst() const { return value.contains(CompileTimeConst); }
   bool isNoDerivative() const { return value.contains(NoDerivative); }
   bool isSending() const { return value.contains(Sending); }
+  bool isAddressable() const { return value.contains(Addressable); }
 
   /// Get the spelling of the parameter specifier used on the parameter.
   ParamSpecifier getOwnershipSpecifier() const {
@@ -2497,6 +2500,12 @@ class ParameterTypeFlags {
                                   : value - ParameterTypeFlags::Sending);
   }
 
+  ParameterTypeFlags withAddressable(bool withAddressable) const {
+    return ParameterTypeFlags(withAddressable
+                                  ? value | ParameterTypeFlags::Addressable
+                                  : value - ParameterTypeFlags::Addressable);
+  }
+
   bool operator ==(const ParameterTypeFlags &other) const {
     return value.toRaw() == other.value.toRaw();
   }
@@ -2590,7 +2599,8 @@ class YieldTypeFlags {
                               /*nonEphemeral*/ false, getOwnershipSpecifier(),
                               /*isolated*/ false, /*noDerivative*/ false,
                               /*compileTimeConst*/ false,
-                              /*is sending*/ false);
+                              /*is sending*/ false,
+                              /*is addressable*/ false);
   }
 
   bool operator ==(const YieldTypeFlags &other) const {
@@ -3386,6 +3396,8 @@ class AnyFunctionType : public TypeBase {
 
     /// Whether the parameter is marked '@noDerivative'.
     bool isNoDerivative() const { return Flags.isNoDerivative(); }
+    
+    bool isAddressable() const { return Flags.isAddressable(); }
 
     /// Whether the parameter might be a semantic result for autodiff purposes.
     /// This includes inout parameters.
@@ -8092,7 +8104,7 @@ inline TupleTypeElt TupleTypeElt::getWithType(Type T) const {
 inline ParameterTypeFlags ParameterTypeFlags::fromParameterType(
     Type paramTy, bool isVariadic, bool isAutoClosure, bool isNonEphemeral,
     ParamSpecifier ownership, bool isolated, bool isNoDerivative,
-    bool compileTimeConst, bool isSending) {
+    bool compileTimeConst, bool isSending, bool isAddressable) {
   // FIXME(Remove InOut): The last caller that needs this is argument
   // decomposition.  Start by enabling the assertion there and fixing up those
   // callers, then remove this, then remove
@@ -8103,7 +8115,8 @@ inline ParameterTypeFlags ParameterTypeFlags::fromParameterType(
     ownership = ParamSpecifier::InOut;
   }
   return {isVariadic, isAutoClosure,  isNonEphemeral,   ownership,
-          isolated,   isNoDerivative, compileTimeConst, isSending};
+          isolated,   isNoDerivative, compileTimeConst, isSending,
+          isAddressable};
 }
 
 inline const Type *BoundGenericType::getTrailingObjectsPointer() const {
diff --git a/include/swift/Basic/Features.def b/include/swift/Basic/Features.def
@@ -430,6 +430,8 @@ EXPERIMENTAL_FEATURE(CoroutineAccessorsAllocateInCallee, false)
 // When a parameter has unspecified isolation, infer it as main actor isolated.
 EXPERIMENTAL_FEATURE(GenerateForceToMainActorThunks, false)
 
+EXPERIMENTAL_FEATURE(AddressableParameters, true)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -8576,7 +8576,7 @@ AnyFunctionType::Param ParamDecl::toFunctionParam(Type type) const {
   auto flags = ParameterTypeFlags::fromParameterType(
       type, isVariadic(), isAutoClosure(), isNonEphemeral(), getSpecifier(),
       isIsolated(), /*isNoDerivative*/ false, isCompileTimeConst(),
-      isSending());
+      isSending(), isAddressable());
   return AnyFunctionType::Param(type, label, flags, internalLabel);
 }
 
diff --git a/lib/AST/FeatureSet.cpp b/lib/AST/FeatureSet.cpp
@@ -290,6 +290,20 @@ static bool usesFeatureIsolatedAny(Decl *decl) {
   });
 }
 
+static bool usesFeatureAddressableParameters(Decl *d) {
+  auto fd = dyn_cast<AbstractFunctionDecl>(d);
+  if (!fd) {
+    return false;
+  }
+  
+  for (auto pd : *fd->getParameters()) {
+    if (pd->isAddressable()) {
+      return true;
+    }
+  }
+  return false;
+}
+
 UNINTERESTING_FEATURE(IsolatedAny2)
 UNINTERESTING_FEATURE(GlobalActorIsolatedTypesUsability)
 UNINTERESTING_FEATURE(ObjCImplementation)
diff --git a/lib/ASTGen/Sources/ASTGen/TypeAttrs.swift b/lib/ASTGen/Sources/ASTGen/TypeAttrs.swift
@@ -51,6 +51,7 @@ extension ASTGenVisitor {
       switch attrKind {
       // Simple type attributes.
       case .autoclosure,
+        .addressable,
         .escaping,
         .noEscape,
         .noDerivative,
diff --git a/lib/Parse/ParsePattern.cpp b/lib/Parse/ParsePattern.cpp
@@ -631,6 +631,8 @@ mapParsedParameters(Parser &parser,
             // or typealias with underlying function type.
             if (ATR->has(TypeAttrKind::Autoclosure))
               param->setAutoClosure(true);
+            if (ATR->has(TypeAttrKind::Addressable))
+              param->setAddressable(true);
 
             unwrappedType = ATR->getTypeRepr();
             continue;
diff --git a/lib/SIL/IR/SILFunctionType.cpp b/lib/SIL/IR/SILFunctionType.cpp
@@ -1720,6 +1720,12 @@ class DestructureInputs {
              ParameterTypeFlags origFlags) {
     assert(!isa<InOutType>(substType));
 
+    // If the parameter is marked addressable, lower it with maximal
+    // abstraction.
+    if (origFlags.isAddressable()) {
+      origType = AbstractionPattern::getOpaque();
+    }
+
     // Tuples get expanded unless they're inout.
     if (origType.isTuple() && ownership != ValueOwnership::InOut) {
       expandTuple(ownership, forSelf, origType, substType, origFlags);
diff --git a/lib/SILGen/SILGenApply.cpp b/lib/SILGen/SILGenApply.cpp
@@ -3506,6 +3506,31 @@ class ArgEmitter {
 private:
   void emit(ArgumentSource &&arg, AbstractionPattern origParamType,
             std::optional<AnyFunctionType::Param> origParam = std::nullopt) {
+    if (origParam && origParam->isAddressable()) {
+      // If the function takes an addressable parameter, and its argument is
+      // a reference to an addressable declaration with compatible ownership,
+      // forward the address along in-place.
+      if (arg.isExpr()) {
+        auto expr = std::move(arg).asKnownExpr();
+        
+        if (auto le = dyn_cast<LoadExpr>(expr)) {
+          expr = le->getSubExpr();
+        }
+        if (auto dre = dyn_cast<DeclRefExpr>(expr)) {
+          if (auto param = dyn_cast<ParamDecl>(dre->getDecl())) {
+            if (param->isAddressable()
+              && param->getValueOwnership() == origParam->getValueOwnership()) {
+              auto addr = SGF.VarLocs[param].value;
+              claimNextParameters(1);
+              Args.push_back(ManagedValue::forBorrowedAddressRValue(addr));
+              return;
+            }
+          }
+        }
+        arg = ArgumentSource(expr);
+      }
+    }
+            
     if (!arg.hasLValueType()) {
       // If the unsubstituted function type has a parameter of tuple type,
       // explode the tuple value.
diff --git a/lib/SILGen/SILGenProlog.cpp b/lib/SILGen/SILGenProlog.cpp
@@ -194,8 +194,15 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
   ManagedValue handleParam(AbstractionPattern origType, CanType substType,
                            ParamDecl *pd) {
     // Note: inouts of tuples are not exploded, so we bypass visit().
-    if (pd->isInOut())
-      return handleInOut(origType, substType);
+    if (pd->isInOut()) {
+      return handleInOut(origType, substType, pd->isAddressable());
+    }
+    // `@_addressable` also suppresses exploding the parameter.
+    if (pd->isAddressable()) {
+      return handleScalar(claimNextParameter(), origType, substType,
+                          /*emitInto*/ nullptr,
+                          /*inout*/ false, /*addressable*/ true);
+    }
     return visit(substType, origType, /*emitInto*/ nullptr);
   }
 
@@ -213,7 +220,9 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
     // Handle scalar components.
     if (!isa<PackExpansionType>(substType)) {
       return handleScalar(componentValue, origPatternType, substType,
-                          /*emit into*/ nullptr, substParam.isInOut());
+                          /*emit into*/ nullptr,
+                          substParam.isInOut(),
+                          /*is addressable*/ false);
     }
 
     auto componentPackTy = componentValue.getType().castTo<SILPackType>();
@@ -259,24 +268,33 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
       }
 
       return handleScalar(input, origPatternType, substEltType,
-                          context.getEmitInto(), /*inout*/ false);
+                          context.getEmitInto(),
+                          /*inout*/ false,
+                          /*addressable*/ false);
     });
   }
 
   ManagedValue visitType(CanType t, AbstractionPattern orig,
                          Initialization *emitInto) {
     auto mv = claimNextParameter();
-    return handleScalar(mv, orig, t, emitInto, /*inout*/ false);
+    return handleScalar(mv, orig, t, emitInto,
+                        /*inout*/ false,
+                        /*addressable*/ false);
   }
 
-  ManagedValue handleInOut(AbstractionPattern orig, CanType t) {
+  ManagedValue handleInOut(AbstractionPattern orig, CanType t,
+                           bool isAddressable) {
     auto mv = claimNextParameter();
-    return handleScalar(mv, orig, t, /*emitInto*/ nullptr, /*inout*/ true);
+    return handleScalar(mv, orig, t, /*emitInto*/ nullptr,
+                        /*inout*/ true,
+                        isAddressable);
   }
 
   ManagedValue handleScalar(ManagedValue mv,
                             AbstractionPattern orig, CanType t,
-                            Initialization *emitInto, bool isInOut) {
+                            Initialization *emitInto,
+                            bool isInOut,
+                            bool isAddressable) {
     assert(!(isInOut && emitInto != nullptr));
 
     auto argType = SGF.getLoweredType(t, mv.getType().getCategory());
@@ -321,20 +339,26 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
       return mv;
     }
 
-    // This can happen if the value is resilient in the calling convention
-    // but not resilient locally.
-    bool argIsLoadable = argType.isLoadable(SGF.F);
-    if (argIsLoadable) {
-      if (argType.isAddress()) {
-        mv = SGF.B.createLoadWithSameOwnership(loc, mv);
-        argType = argType.getObjectType();
+    // If the parameter is marked `@_addressable`, then we want to defer any
+    // reabstraction of the parameter as received, so that we can use the
+    // original value at its stable address when possible.
+    bool argIsLoadable = false;
+    if (!isAddressable) {
+      argIsLoadable = argType.isLoadable(SGF.F);
+      // This can happen if the value is resilient in the calling convention
+      // but not resilient locally.
+      if (argIsLoadable) {
+        if (argType.isAddress()) {
+          mv = SGF.B.createLoadWithSameOwnership(loc, mv);
+          argType = argType.getObjectType();
+        }
       }
-    }
 
-    assert(argType.getCategory() == mv.getType().getCategory());
-    if (argType.getASTType() != mv.getType().getASTType()) {
-      // Reabstract the value if necessary.
-      mv = SGF.emitOrigToSubstValue(loc, mv.ensurePlusOne(SGF, loc), orig, t);
+      assert(argType.getCategory() == mv.getType().getCategory());
+      if (argType.getASTType() != mv.getType().getASTType()) {
+        // Reabstract the value if necessary.
+        mv = SGF.emitOrigToSubstValue(loc, mv.ensurePlusOne(SGF, loc), orig, t);
+      }
     }
 
     if (parameters.isNoImplicitCopy && !argIsLoadable) {
@@ -538,7 +562,8 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
         auto eltAddrMV = cloner.clone(eltAddr);
         auto result = handleScalar(eltAddrMV, origPatternType,
                                    substComponentType, componentInit,
-                                   /*inout*/ false);
+                                   /*inout*/ false,
+                                   /*addressable*/ false);
         assert(result.isInContext() == (componentInit != nullptr));
         if (!result.isInContext())
           eltMVs.push_back(result);
@@ -571,7 +596,9 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
           auto eltAddrMV = cloner.clone(eltAddr);
 
           auto result = handleScalar(eltAddrMV, origPatternType, substEltType,
-                                     eltInit, /*inout*/ false);
+                                     eltInit,
+                                     /*inout*/ false,
+                                     /*addressable*/ false);
           assert(result.isInContext()); (void) result;
         });
       });
diff --git a/lib/Sema/TypeCheckType.cpp b/lib/Sema/TypeCheckType.cpp
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
diff --git a/test/SILGen/addressable_params.swift b/test/SILGen/addressable_params.swift
diff --git a/test/attr/attr_addressable.swift b/test/attr/attr_addressable.swift

Original file line number	Diff line number	Diff line change
`@@ -8576,7 +8576,7 @@ AnyFunctionType::Param ParamDecl::toFunctionParam(Type type) const {`
`8576`	`8576`	`auto flags = ParameterTypeFlags::fromParameterType(`
`8577`	`8577`	`type, isVariadic(), isAutoClosure(), isNonEphemeral(), getSpecifier(),`
`8578`	`8578`	`isIsolated(), /isNoDerivative/ false, isCompileTimeConst(),`
`8579`		`- isSending());`
	`8579`	`+ isSending(), isAddressable());`
`8580`	`8580`	`return AnyFunctionType::Param(type, label, flags, internalLabel);`
`8581`	`8581`	`}`
`8582`	`8582`