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

[ldionne]

Summary

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

rdar://100907096

[ldionne]

I was waiting to confirm some stuff internally, I think this is the way to go and not a Clang bug.

[ahatanak]

LGTM

libcxx/include/__functional/function.h

_LIBCPP_INLINE_VISIBILITY const _Tp* target() const _NOEXCEPT
         if (__policy_->__clone) // Out of line storage.
             return reinterpret_cast<const _Tp*>(__buf_.__large);
         else
             return reinterpret_cast<const _Tp*>(&__buf_.__small);
     }
 #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 *);
 
 template<class _Rp1, class ..._ArgTypes1, class _Alloc, class _Rp, class ..._ArgTypes>
 class __func<_Rp1(^)(_ArgTypes1...), _Alloc, _Rp(_ArgTypes...)>
     : public  __base<_Rp(_ArgTypes...)>
 {
     typedef _Rp1(^__block_type)(_ArgTypes1...);
     __block_type __f_;
 
 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 {
         _LIBCPP_ASSERT(false,
             "Block pointers are just pointers, so they should always fit into "
             "std::function's small buffer optimization. This function should "
             "never be invoked.");
         return nullptr;
     }
 
     virtual void __clone(__base<_Rp(_ArgTypes...)>* __p) const {
         ::new ((void*)__p) __func(__f_);
     }
 
     virtual void destroy() _NOEXCEPT {
+#ifndef _LIBCPP_HAS_OBJC_ARC
         if (__f_)
             _Block_release(__f_);
+#endif
         __f_ = 0;
     }
 
     virtual void destroy_deallocate() _NOEXCEPT {
         _LIBCPP_ASSERT(false,
             "Block pointers are just pointers, so they should always fit into "
             "std::function's small buffer optimization. This function should "
             "never be invoked.");
#ifndef _LIBCPP_NO_RTTI
     }
 
     virtual const std::type_info& target_type() const _NOEXCEPT {
         return typeid(__func::__block_type);
     }
 #endif // _LIBCPP_NO_RTTI
 };
 
-#endif // _LIBCPP_HAS_EXTENSION_BLOCKS && !_LIBCPP_HAS_OBJC_ARC
+#endif // _LIBCPP_HAS_EXTENSION_BLOCKS
 
 } // namespace __function
 
 template<class _Rp, class ..._ArgTypes>
 class _LIBCPP_TEMPLATE_VIS function<_Rp(_ArgTypes...)>
     : public __function::__maybe_derive_from_unary_function<_Rp(_ArgTypes...)>,
       public __function::__maybe_derive_from_binary_function<_Rp(_ArgTypes...)>
 {

libcxx/test/libcxx/utilities/function.objects/func.blocks.arc.pass.mm

+//===----------------------------------------------------------------------===//
+//
+// 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;
+}

libcxx/test/libcxx/utilities/function.objects/func.blocks.pass.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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 the "blocks" extension
 
 // UNSUPPORTED: c++03
 
 // This test requires the Blocks runtime, which is (only?) available
 // on Darwin out-of-the-box.
 // REQUIRES: has-fblocks && darwin
 
-// RUN: %{build} -fblocks
-// RUN: %{run}
+// ADDITIONAL_COMPILE_FLAGS: -fblocks
 
 #include <functional>
 #include <cstdlib>
 #include <cassert>
 
 #include <Block.h>
 
 #include "test_macros.h"