Merge pull request #28527 from slavapestov/owned-block-params

Fixes for __attribute__((ns_consumed)) block parameters

Author: slavapestov

lib/IRGen/GenClangType.cpp

@@ -585,7 +585,10 @@ clang::CanQualType GenClangType::visitSILFunctionType(CanSILFunctionType type) {
   }
 
   SmallVector<clang::QualType, 4> paramTypes;
+  SmallVector<clang::FunctionProtoType::ExtParameterInfo, 4> extParamInfos;
   for (auto paramTy : type->getParameters()) {
+    clang::FunctionProtoType::ExtParameterInfo extParamInfo;
+
     // Blocks should only take direct +0 parameters.
     switch (paramTy.getConvention()) {
     case ParameterConvention::Direct_Guaranteed:
@@ -594,7 +597,9 @@ clang::CanQualType GenClangType::visitSILFunctionType(CanSILFunctionType type) {
       break;
 
     case ParameterConvention::Direct_Owned:
-      llvm_unreachable("block takes owned parameter");
+      extParamInfo = extParamInfo.withIsConsumed(true);
+      break;
+
     case ParameterConvention::Indirect_In:
     case ParameterConvention::Indirect_In_Constant:
     case ParameterConvention::Indirect_Inout:
@@ -606,12 +611,16 @@ clang::CanQualType GenClangType::visitSILFunctionType(CanSILFunctionType type) {
                              paramTy.getArgumentType(IGM.getSILModule(), type));
     if (param.isNull())
       return clang::CanQualType();
+
     paramTypes.push_back(param);
+    extParamInfos.push_back(extParamInfo);
   }
 
   // Build the Clang function type.
-  clang::FunctionProtoType::ExtProtoInfo defaultEPI;
-  auto fnTy = clangCtx.getFunctionType(resultType, paramTypes, defaultEPI);
+  clang::FunctionProtoType::ExtProtoInfo extProtoInfo;
+  extProtoInfo.ExtParameterInfos = extParamInfos.begin();
+
+  auto fnTy = clangCtx.getFunctionType(resultType, paramTypes, extProtoInfo);
   clang::QualType ptrTy;
 
   switch (kind) {

lib/PrintAsObjC/DeclAndTypePrinter.cpp

@@ -1934,7 +1934,18 @@ class DeclAndTypePrinter::Implementation
     if (!FT->getParams().empty()) {
       interleave(FT->getParams(),
                  [this](const AnyFunctionType::Param &param) {
-                   print(param.getOldType(), OTK_None, param.getLabel(),
+                   switch (param.getValueOwnership()) {
+                   case ValueOwnership::Default:
+                   case ValueOwnership::Shared:
+                     break;
+                   case ValueOwnership::Owned:
+                     os << "SWIFT_RELEASES_ARGUMENT ";
+                     break;
+                   case ValueOwnership::InOut:
+                     llvm_unreachable("bad specifier");
+                   }
+
+                   print(param.getParameterType(), OTK_None, param.getLabel(),
                          IsFunctionParam);
                  },
                  [this] { os << ", "; });

lib/PrintAsObjC/PrintAsObjC.cpp

@@ -112,6 +112,11 @@ static void writePrologue(raw_ostream &out, ASTContext &ctx) {
          "#else\n"
          "# define SWIFT_NOESCAPE\n"
          "#endif\n"
+         "#if __has_attribute(ns_consumed)\n"
+         "# define SWIFT_RELEASES_ARGUMENT __attribute__((ns_consumed))\n"
+         "#else\n"
+         "# define SWIFT_RELEASES_ARGUMENT\n"
+         "#endif\n"
          "#if __has_attribute(warn_unused_result)\n"
          "# define SWIFT_WARN_UNUSED_RESULT __attribute__((warn_unused_result))\n"
          "#else\n"

test/IRGen/Inputs/usr/include/Gizmo.h

@@ -45,6 +45,7 @@ typedef long NSInteger;
 - (void) setFrame: (struct NSRect) rect;
 - (void) frob;
 - (void) test: (struct Fob) fob;
+- (void) perform: (void (^)(NS_CONSUMED Gizmo*)) block;
 + (void) runce;
 @end
 

test/IRGen/objc_block_consumed.swift

@@ -0,0 +1,19 @@
+
+// RUN: %empty-directory(%t)
+// RUN: %build-irgen-test-overlays
+// RUN: %target-swift-frontend(mock-sdk: -sdk %S/Inputs -I %t) -primary-file %s -emit-ir -disable-objc-attr-requires-foundation-module
+// RUN: %target-swift-frontend(mock-sdk: -sdk %S/Inputs -I %t) -primary-file %s -emit-silgen -disable-objc-attr-requires-foundation-module | %FileCheck %s
+
+// We want to test that IRGen doesn't assert on this code. SIL is the best place
+// to file check that the block parameter is actually +1.
+
+// REQUIRES: CPU=x86_64
+// REQUIRES: objc_interop
+
+import gizmo
+
+// CHECK-LABEL: sil hidden [ossa] @$s19objc_block_consumed24passBlockWithConsumedArgyySo5GizmoC_ADtF : $@convention(thin) (@guaranteed Gizmo, @guaranteed Gizmo) -> () {
+func passBlockWithConsumedArg(_ g: Gizmo, _ other: Gizmo) {
+  // CHECK: objc_method %0 : $Gizmo, #Gizmo.perform!1.foreign : (Gizmo) -> (((Gizmo?) -> ())?) -> (), $@convention(objc_method) (Optional<@convention(block) (@owned Optional<Gizmo>) -> ()>, Gizmo) -> ()
+  g.perform { other in }
+}

test/Interpreter/SDK/Inputs/objc_block_consumed.h

@@ -0,0 +1,5 @@
+#import <Foundation/Foundation.h>
+
+static inline void takesBlockWithConsumedArg(void (^ block)(NS_RELEASES_ARGUMENT NSObject *x), NSObject *x) {
+  block(x);
+}

test/Interpreter/SDK/objc_block_consumed.swift

@@ -0,0 +1,21 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -import-objc-header %S/Inputs/objc_block_consumed.h -o %t/main
+// RUN: %target-run %t/main
+
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import Foundation
+import StdlibUnittest
+
+class C : NSObject {
+  var tracked = LifetimeTracked(0)
+}
+
+var ObjCBlockConsumedTestSuite = TestSuite("ObjCBlockConsumed")
+
+ObjCBlockConsumedTestSuite.test("Test") {
+  takesBlockWithConsumedArg({ arg in }, C())
+}
+
+runAllTests()

test/PrintAsObjC/blocks.swift

@@ -113,6 +113,9 @@ typealias MyBlockWithNoescapeParam = (() -> ()) -> Int
     return input
   }
 
+  // CHECK-NEXT: - (void)blockWithConsumingArgument:(void (^ _Nonnull)(SWIFT_RELEASES_ARGUMENT NSObject * _Nonnull))block;
+  @objc func blockWithConsumingArgument(_ block: @escaping (__owned NSObject) -> ()) {}
+
   // CHECK-NEXT: @property (nonatomic, copy) NSInteger (^ _Nullable savedBlock)(NSInteger);
   @objc var savedBlock: ((Int) -> Int)?