[libc++] Add support for blocks in std::function

rdar://14390808

Author: ldionne

libcxx/include/functional

@@ -508,6 +508,10 @@ POLICY:  For non-variadic implementations, the number of arguments is limited
 
 #include <__functional_base>
 
+#if defined(_LIBCPP_HAS_EXTENSION_BLOCKS) && !defined(_LIBCPP_HAS_OBJC_ARC)
+#include <Block.h>
+#endif
+
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #pragma GCC system_header
 #endif
@@ -1484,6 +1488,12 @@ template <class _Fp>
 _LIBCPP_INLINE_VISIBILITY
 bool __not_null(function<_Fp> const& __f) { return !!__f; }
 
+#ifdef _LIBCPP_HAS_EXTENSION_BLOCKS
+template <class _Rp, class ..._Args>
+_LIBCPP_INLINE_VISIBILITY
+bool __not_null(_Rp (^__p)(_Args...)) { return __p; }
+#endif
+
 } // namespace __function
 
 #ifndef _LIBCPP_CXX03_LANG
@@ -2245,6 +2255,72 @@ template <class _Rp, class... _ArgTypes> class __policy_func<_Rp(_ArgTypes...)>
 #endif // _LIBCPP_NO_RTTI
 };
 
+#if defined(_LIBCPP_HAS_EXTENSION_BLOCKS) && !defined(_LIBCPP_HAS_OBJC_ARC)
+
+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)
+        : __f_(__f ? Block_copy(__f) : (__block_type)0)
+    { }
+
+    // [TODO] add && to save on a retain
+
+    _LIBCPP_INLINE_VISIBILITY
+    explicit __func(__block_type __f, const _Alloc& /* unused */)
+        : __f_(__f ? Block_copy(__f) : (__block_type)0)
+    { }
+
+    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 (__p) __func(__f_);
+    }
+
+    virtual void destroy() _NOEXCEPT {
+        if (__f_)
+            Block_release(__f_);
+        __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.");
+    }
+
+    virtual _Rp operator()(_ArgTypes&& ... __arg) {
+        return __invoke(__f_, _VSTD::forward<_ArgTypes>(__arg)...);
+    }
+
+#ifndef _LIBCPP_NO_RTTI
+    virtual const void* target(type_info const& __ti) const _NOEXCEPT {
+        if (__ti == typeid(__func::__block_type))
+            return &__f_;
+        return (const void*)nullptr;
+    }
+
+    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
+
 }  // __function
 
 template<class _Rp, class ..._ArgTypes>

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

@@ -0,0 +1,142 @@
+//===----------------------------------------------------------------------===//
+//
+// 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++98, c++03
+// REQUIRES: has-fblocks
+
+// RUN: %{build} -fblocks
+// RUN: %{run}
+
+#include <functional>
+#include <cstdlib>
+#include <cassert>
+
+#include "test_macros.h"
+#include "count_new.h"
+
+
+struct A {
+  static int count;
+  int id_;
+  explicit A(int id) { ++count; id_ = id; }
+  A(const A &a) { id_ = a.id_; ++count; }
+  ~A() { id_ = -1; --count; }
+  int operator()() const { return -1; }
+  int operator()(int i) const { return i; }
+  int operator()(int, int) const { return -2; }
+  int operator()(int, int, int) const { return -3; }
+  int id() const { return id_; }
+};
+
+int A::count = 0;
+
+int g(int) { return 0; }
+
+int main(int, char**)
+{
+    // swap
+    {
+        std::function<int(int)> f1 = g;
+        std::function<int(int)> f2 = ^(int x) { return x + 1; };
+        assert(globalMemCounter.checkOutstandingNewEq(0));
+        assert(*f1.target<int(*)(int)>() == g);
+        assert(*f2.target<int(^)(int)>() != 0);
+        swap(f1, f2);
+        assert(globalMemCounter.checkOutstandingNewEq(0));
+        assert(*f1.target<int(^)(int)>() != 0);
+        assert(*f2.target<int(*)(int)>() == g);
+    }
+
+    // operator bool
+    {
+        std::function<int(int)> f;
+        assert(!f);
+        f = ^(int x) { return x+1; };
+        assert(f);
+    }
+
+    // operator()
+    {
+        std::function<int ()> r1(^{ return 4; });
+        assert(r1() == 4);
+    }
+    {
+        __block bool called = false;
+        std::function<void ()> r1(^{ called = true; });
+        r1();
+        assert(called);
+    }
+    {
+        __block int param = 0;
+        std::function<void (int)> r1(^(int x){ param = x; });
+        r1(4);
+        assert(param == 4);
+    }
+    {
+        std::function<int (int)> r1(^(int x){ return x + 4; });
+        assert(r1(3) == 7);
+    }
+    {
+        __block int param1 = 0;
+        __block int param2 = 0;
+        std::function<void (int, int)> r1(^(int x, int y){ param1 = x; param2 = y; });
+        r1(3, 4);
+        assert(param1 == 3);
+        assert(param2 == 4);
+    }
+    {
+        std::function<int (int, int)> r1(^(int x, int y){ return x + y; });
+        assert(r1(3, 4) == 7);
+    }
+
+    // swap
+    {
+        std::function<int(int)> f1 = A(999);
+        std::function<int(int)> f2 = ^(int x) { return x + 1; };
+        assert(A::count == 1);
+        assert(globalMemCounter.checkOutstandingNewEq(1));
+        assert(f1.target<A>()->id() == 999);
+        assert((*f2.target<int(^)(int)>())(13) == 14);
+        f1.swap(f2);
+        assert(A::count == 1);
+        assert(globalMemCounter.checkOutstandingNewEq(1));
+        assert((*f1.target<int(^)(int)>())(13) == 14);
+        assert(f2.target<A>()->id() == 999);
+    }
+    assert(globalMemCounter.checkOutstandingNewEq(0));
+    assert(A::count == 0);
+
+    // operator== and operator!=
+    {
+        std::function<int(int)> f;
+        assert(f == nullptr);
+        assert(nullptr == f);
+        f = ^(int x) { return x + 1; };
+        assert(f != nullptr);
+        assert(nullptr != f);
+    }
+
+    // target
+    {
+        int (^block)(int) = Block_copy(^(int x) { return x + 1; });
+        std::function<int(int)> f = block;
+        assert(*f.target<int(^)(int)>() == block);
+        assert(f.target<int(*)(int)>() == 0);
+        Block_release(block);
+    }
+
+    // target_type
+    {
+        std::function<int(int)> f = ^(int x) { return x + 1; };
+        assert(f.target_type() == typeid(int(^)(int)));
+    }
+
+    return 0;
+}