[libc++] Fix std::function's handling of blocks under Objc ARC

Previously, some uses of std::function with blocks would crash when ARC was enabled.

rdar://100907096

Differential Revision: https://reviews.llvm.org/D135706

(cherry picked from commit 0e4802bf45952b1120c52d4d1bf6bfa2800fd102)

3 files changed, 102 insertions, 4 deletions

diff --git a/libcxx/include/__functional/function.h b/libcxx/include/__functional/function.h
index db3af6e24101..55b607f3f804 100644
--- a/libcxx/include/__functional/function.h
+++ b/libcxx/include/__functional/function.h
@@ -883,7 +883,7 @@ template <class _Rp, class... _ArgTypes> class __policy_func<_Rp(_ArgTypes...)>
 #endif // _LIBCPP_NO_RTTI
 };
 
-#if defined(_LIBCPP_HAS_BLOCKS_RUNTIME) && !defined(_LIBCPP_HAS_OBJC_ARC)
+#if defined(_LIBCPP_HAS_BLOCKS_RUNTIME)
 
 extern "C" void *_Block_copy(const void *);
 extern "C" void _Block_release(const void *);
@@ -898,14 +898,22 @@ class __func<_Rp1(^)(_ArgTypes1...), _Alloc, _Rp(_ArgTypes...)>
 public:
     _LIBCPP_INLINE_VISIBILITY
     explicit __func(__block_type const& __f)
+#ifdef _LIBCPP_HAS_OBJC_ARC
+        : __f_(__f)
+#else
         : __f_(reinterpret_cast<__block_type>(__f ? _Block_copy(__f) : nullptr))
+#endif
     { }
 
     // [TODO] add && to save on a retain
 
     _LIBCPP_INLINE_VISIBILITY
     explicit __func(__block_type __f, const _Alloc& /* unused */)
+#ifdef _LIBCPP_HAS_OBJC_ARC
+        : __f_(__f)
+#else
         : __f_(reinterpret_cast<__block_type>(__f ? _Block_copy(__f) : nullptr))
+#endif
     { }
 
     virtual __base<_Rp(_ArgTypes...)>* __clone() const {
@@ -921,8 +929,10 @@ public:
     }
 
     virtual void destroy() _NOEXCEPT {
+#ifndef _LIBCPP_HAS_OBJC_ARC
         if (__f_)
             _Block_release(__f_);
+#endif
         __f_ = 0;
     }
 
@@ -950,7 +960,7 @@ public:
 #endif // _LIBCPP_NO_RTTI
 };
 
-#endif // _LIBCPP_HAS_EXTENSION_BLOCKS && !_LIBCPP_HAS_OBJC_ARC
+#endif // _LIBCPP_HAS_EXTENSION_BLOCKS
 
 } // namespace __function
 
diff --git a/libcxx/test/libcxx/utilities/function.objects/func.blocks.arc.pass.mm b/libcxx/test/libcxx/utilities/function.objects/func.blocks.arc.pass.mm
new file mode 100644
index 000000000000..186fe22e6e47
--- /dev/null
+++ b/libcxx/test/libcxx/utilities/function.objects/func.blocks.arc.pass.mm
@@ -0,0 +1,89 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// std::function support for "blocks" when ARC is enabled
+
+// UNSUPPORTED: c++03
+
+// This test requires the Blocks runtime, which is (only?) available on Darwin
+// out-of-the-box.
+// REQUIRES: has-fblocks && darwin
+
+// ADDITIONAL_COMPILE_FLAGS: -fblocks -fobjc-arc
+
+#include <functional>
+
+#include <cassert>
+#include <cstddef>
+#include <string>
+
+struct Foo {
+  Foo() = default;
+  Foo(std::size_t (^bl)()) : f(bl) {}
+
+  std::function<int()> f;
+};
+
+Foo Factory(std::size_t (^bl)()) {
+  Foo result(bl);
+  return result;
+}
+
+Foo Factory2() {
+  auto hello = std::string("Hello world");
+  return Factory(^() {
+    return hello.size();
+  });
+}
+
+Foo AssignmentFactory(std::size_t (^bl)()) {
+  Foo result;
+  result.f = bl;
+  return result;
+}
+
+Foo AssignmentFactory2() {
+  auto hello = std::string("Hello world");
+  return AssignmentFactory(^() {
+    return hello.size();
+  });
+}
+
+int main(int, char **) {
+  // Case 1, works
+  {
+    auto hello = std::string("Hello world");
+    auto f = AssignmentFactory(^() {
+      return hello.size();
+    });
+    assert(f.f() == 11);
+  }
+
+  // Case 2, works
+  {
+    auto f = AssignmentFactory2();
+    assert(f.f() == 11);
+  }
+
+  // Case 3, works
+  {
+    auto hello = std::string("Hello world");
+    auto f = Factory(^() {
+      return hello.size();
+    });
+    assert(f.f() == 11);
+  }
+
+  // Case 4, used to crash under ARC
+  {
+    auto f = Factory2();
+    assert(f.f() == 11);
+  }
+
+  return 0;
+}
diff --git a/libcxx/test/libcxx/utilities/function.objects/func.blocks.sh.cpp b/libcxx/test/libcxx/utilities/function.objects/func.blocks.pass.cpp
index ecebc7c9800f..b95b6ebb534a 100644
--- a/libcxx/test/libcxx/utilities/function.objects/func.blocks.sh.cpp
+++ b/libcxx/test/libcxx/utilities/function.objects/func.blocks.pass.cpp
@@ -14,8 +14,7 @@
 // on Darwin out-of-the-box.
 // REQUIRES: has-fblocks && darwin
 
-// RUN: %{build} -fblocks
-// RUN: %{run}
+// ADDITIONAL_COMPILE_FLAGS: -fblocks
 
 #include <functional>
 #include <cstdlib>