Merge pull request #75282 from jckarter/addressable-params-1

jckarter · web-flow · commit a045d665bb5c · 2024-12-05T09:17:26.000-08:00
[WIP] Prototype an `@_addressable` attribute that puts an argument at a stable address.
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`