AK: Support storing blocks in AK::Function

This has two slightly different implementations for ARC and non-ARC
compiler modes. The main idea is to store a block pointer as our
closure and use either ARC magic or BlockRuntime methods to manage
the memory for the block. Things are complicated by the fact that
we don't yet force-enable swift, so we can't count on the swift.org
llvm fork being our compiler toolchain. The patch adds some CMake
checks and ifdefs to still support environments without support
for blocks or ARC.

AK/Function.h

@@ -2,6 +2,7 @@
  * Copyright (C) 2016 Apple Inc. All rights reserved.
  * Copyright (c) 2021, Gunnar Beutner <gbeutner@serenityos.org>
  * Copyright (c) 2018-2023, Andreas Kling <andreas@ladybird.org>
+ * Copyright (c) 2025, Andrew Kaster <andrew@ladybird.org>
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
@@ -34,8 +35,14 @@
 #include <AK/ScopeGuard.h>
 #include <AK/Span.h>
 #include <AK/StdLibExtras.h>
+#include <AK/TypeCasts.h>
 #include <AK/Types.h>
 
+// BlockRuntime methods for Objective-C block closure support.
+extern "C" void* _Block_copy(void const*);
+extern "C" void _Block_release(void const*);
+extern "C" size_t Block_size(void const*);
+
 namespace AK {
 
 // These annotations are used to avoid capturing a variable with local storage in a lambda that outlives it
@@ -48,6 +55,17 @@ namespace AK {
 #    define IGNORE_USE_IN_ESCAPING_LAMBDA
 #endif
 
+namespace Detail {
+#ifdef AK_HAS_OBJC_ARC
+inline constexpr bool HaveObjcArc = true;
+#else
+inline constexpr bool HaveObjcArc = false;
+#endif
+
+// validated in TestFunction.mm
+inline constexpr size_t block_layout_size = 32;
+}
+
 template<typename>
 class Function;
 
@@ -84,7 +102,7 @@ public:
         if (!m_size)
             return {};
         if (auto* wrapper = callable_wrapper())
-            return ReadonlyBytes { wrapper, m_size };
+            return ReadonlyBytes { wrapper->raw_callable(), m_size };
         return {};
     }
 
@@ -102,6 +120,13 @@ public:
         init_with_callable(move(f), CallableKind::FunctionPointer);
     }
 
+    template<typename BlockType>
+    Function(BlockType b)
+    requires((IsBlockClosure<BlockType> && IsCallableWithArguments<BlockType, Out, In...>))
+    {
+        init_with_callable(move(b), CallableKind::Block);
+    }
+
     Function(Function&& other)
     {
         move_from(move(other));
@@ -141,6 +166,15 @@ public:
         return *this;
     }
 
+    template<typename BlockType>
+    Function& operator=(BlockType&& block)
+    requires((IsBlockClosure<BlockType> && IsCallableWithArguments<BlockType, Out, In...>))
+    {
+        clear();
+        init_with_callable(static_cast<RemoveCVReference<BlockType>>(block), CallableKind::Block);
+        return *this;
+    }
+
     Function& operator=(nullptr_t)
     {
         clear();
@@ -160,6 +194,7 @@ private:
     enum class CallableKind {
         FunctionPointer,
         FunctionObject,
+        Block,
     };
 
     class CallableWrapperBase {
@@ -169,6 +204,7 @@ private:
         virtual Out call(In...) = 0;
         virtual void destroy() = 0;
         virtual void init_and_swap(u8*, size_t) = 0;
+        virtual void const* raw_callable() const = 0;
     };
 
     template<typename CallableType>
@@ -189,7 +225,15 @@ private:
 
         void destroy() final override
         {
-            delete this;
+            if constexpr (IsBlockClosure<CallableType>) {
+                if constexpr (Detail::HaveObjcArc)
+                    m_callable = nullptr;
+                else
+                    _Block_release(m_callable);
+            } else {
+                // This code is a bit too clever for gcc. Pinky promise we're only deleting heap objects.
+                AK_IGNORE_DIAGNOSTIC("-Wfree-nonheap-object", delete this);
+            }
         }
 
         // NOLINTNEXTLINE(readability-non-const-parameter) False positive; destination is used in a placement new expression
@@ -199,6 +243,14 @@ private:
             new (destination) CallableWrapper { move(m_callable) };
         }
 
+        void const* raw_callable() const final override
+        {
+            if constexpr (IsBlockClosure<CallableType>)
+                return static_cast<u8 const*>(bridge_cast<void>(m_callable)) + Detail::block_layout_size;
+            else
+                return &m_callable;
+        }
+
     private:
         CallableType m_callable;
     };
@@ -207,6 +259,7 @@ private:
         NullPointer,
         Inline,
         Outline,
+        Block,
     };
 
     CallableWrapperBase* callable_wrapper() const
@@ -215,6 +268,7 @@ private:
         case FunctionKind::NullPointer:
             return nullptr;
         case FunctionKind::Inline:
+        case FunctionKind::Block:
             return bit_cast<CallableWrapperBase*>(&m_storage);
         case FunctionKind::Outline:
             return *bit_cast<CallableWrapperBase**>(&m_storage);
@@ -234,12 +288,22 @@ private:
         }
         m_deferred_clear = false;
         auto* wrapper = callable_wrapper();
-        if (m_kind == FunctionKind::Inline) {
+        switch (m_kind) {
+        case FunctionKind::Inline:
             VERIFY(wrapper);
             wrapper->~CallableWrapperBase();
-        } else if (m_kind == FunctionKind::Outline) {
+            break;
+        case FunctionKind::Outline:
             VERIFY(wrapper);
             wrapper->destroy();
+            break;
+        case FunctionKind::Block:
+            VERIFY(wrapper);
+            wrapper->destroy();
+            wrapper->~CallableWrapperBase();
+            break;
+        case FunctionKind::NullPointer:
+            break;
         }
         m_kind = FunctionKind::NullPointer;
     }
@@ -256,18 +320,33 @@ private:
         }
         VERIFY(m_call_nesting_level == 0);
         using WrapperType = CallableWrapper<Callable>;
+        if (callable_kind == CallableKind::FunctionObject)
+            m_size = sizeof(Callable);
+        else
+            m_size = 0;
+        if constexpr (IsBlockClosure<Callable>) {
+            auto block_size = Block_size(bridge_cast<void>(callable));
+            VERIFY(block_size >= Detail::block_layout_size);
+            m_size = block_size - Detail::block_layout_size;
+        }
+
         if constexpr (alignof(Callable) > inline_alignment || sizeof(WrapperType) > inline_capacity) {
             *bit_cast<CallableWrapperBase**>(&m_storage) = new WrapperType(forward<Callable>(callable));
             m_kind = FunctionKind::Outline;
         } else {
             static_assert(sizeof(WrapperType) <= inline_capacity);
-            new (m_storage) WrapperType(forward<Callable>(callable));
-            m_kind = FunctionKind::Inline;
+            if constexpr (IsBlockClosure<Callable>) {
+                if constexpr (Detail::HaveObjcArc) {
+                    new (m_storage) WrapperType(forward<Callable>(callable));
+                } else {
+                    new (m_storage) WrapperType(reinterpret_cast<Callable>(_Block_copy(callable)));
+                }
+                m_kind = FunctionKind::Block;
+            } else {
+                new (m_storage) WrapperType(forward<Callable>(callable));
+                m_kind = FunctionKind::Inline;
+            }
         }
-        if (callable_kind == CallableKind::FunctionObject)
-            m_size = sizeof(WrapperType);
-        else
-            m_size = 0;
     }
 
     void move_from(Function&& other)
@@ -279,8 +358,9 @@ private:
         case FunctionKind::NullPointer:
             break;
         case FunctionKind::Inline:
+        case FunctionKind::Block:
             other_wrapper->init_and_swap(m_storage, inline_capacity);
-            m_kind = FunctionKind::Inline;
+            m_kind = other.m_kind;
             break;
         case FunctionKind::Outline:
             *bit_cast<CallableWrapperBase**>(&m_storage) = other_wrapper;

AK/StdLibExtraDetails.h

@@ -142,6 +142,27 @@ inline constexpr bool IsFunction<Ret(Args...) const volatile&&> = true;
 template<class Ret, class... Args>
 inline constexpr bool IsFunction<Ret(Args..., ...) const volatile&&> = true;
 
+template<class>
+inline constexpr bool IsBlockClosure = false;
+#ifdef AK_HAS_BLOCKS
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^)(Args...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^)(Args..., ...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^const)(Args...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^const)(Args..., ...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^volatile)(Args...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^volatile)(Args..., ...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^const volatile)(Args...)> = true;
+template<class Ret, class... Args>
+inline constexpr bool IsBlockClosure<Ret (^const volatile)(Args..., ...)> = true;
+#endif
+
 template<class T>
 inline constexpr bool IsRvalueReference = false;
 template<class T>
@@ -641,6 +662,7 @@ using AK::Detail::InvokeResult;
 using AK::Detail::IsArithmetic;
 using AK::Detail::IsAssignable;
 using AK::Detail::IsBaseOf;
+using AK::Detail::IsBlockClosure;
 using AK::Detail::IsClass;
 using AK::Detail::IsConst;
 using AK::Detail::IsConstructible;

Meta/CMake/common_options.cmake

@@ -45,3 +45,22 @@ serenity_option(ENABLE_STD_STACKTRACE OFF CACHE BOOL "Force use of std::stacktra
 if (ENABLE_SWIFT)
     include(${CMAKE_CURRENT_LIST_DIR}/Swift/swift-settings.cmake)
 endif()
+
+include(CheckCXXSourceCompiles)
+set(BLOCKS_REQUIRED_LIBRARIES "")
+if (NOT APPLE)
+    find_package(BlocksRuntime)
+    if (BlocksRuntime_FOUND)
+        set(BLOCKS_REQUIRED_LIBRARIES BlocksRuntime::BlocksRuntime)
+        set(CMAKE_REQUIRED_LIBRARIES BlocksRuntime::BlocksRuntime)
+    endif()
+endif()
+check_cxx_source_compiles([=[
+    int main() { __block int x = 0; auto b = ^{++x;}; b(); }
+]=] CXX_COMPILER_SUPPORTS_BLOCKS)
+
+set(CMAKE_REQUIRED_FLAGS "-fobjc-arc")
+check_cxx_source_compiles([=[
+    int main() { auto b = ^{}; auto __weak w = b; w(); }
+]=] CXX_COMPILER_SUPPORTS_OBJC_ARC)
+unset(CMAKE_REQUIRED_FLAGS)

Tests/AK/CMakeLists.txt

@@ -88,6 +88,16 @@ foreach(source IN LISTS AK_TEST_SOURCES)
     serenity_test("${source}" AK)
 endforeach()
 
+if (CXX_COMPILER_SUPPORTS_BLOCKS)
+    serenity_test(TestFunction.mm AK NAME TestFunction)
+    target_link_libraries(TestFunction PRIVATE ${BLOCKS_REQUIRED_LIBRARIES})
+endif()
+if (CXX_COMPILER_SUPPORTS_OBJC_ARC)
+    serenity_test(TestFunction.mm AK NAME TestFunctionArc)
+    target_compile_options(TestFunctionArc PRIVATE -fobjc-arc)
+    target_link_libraries(TestFunction PRIVATE ${BLOCKS_REQUIRED_LIBRARIES})
+endif()
+
 target_link_libraries(TestString PRIVATE LibUnicode)
 
 if (ENABLE_SWIFT)

Tests/AK/TestFunction.mm

@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) 2025, Andrew Kaster <andrew@ladybird.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <LibTest/TestCase.h>
+
+#include <AK/ByteReader.h>
+#include <AK/Function.h>
+#include <AK/Platform.h>
+
+TEST_CASE(SimpleBlock)
+{
+    auto b = ^{ };
+
+    static_assert(IsBlockClosure<decltype(b)>);
+
+    auto f = Function<void()>(b);
+
+    f();
+}
+
+TEST_CASE(BlockCaptureInt)
+{
+    __block int x = 0;
+    auto b = ^{
+        x = 2;
+    };
+    auto f = Function<void()>(b);
+
+    f();
+
+    EXPECT_EQ(x, 2);
+}
+
+TEST_CASE(BlockCaptureString)
+{
+    __block String s = "hello"_string;
+    auto b = ^{
+        s = "world"_string;
+    };
+    auto f = Function<void()>(b);
+
+    f();
+
+    EXPECT_EQ(s, "world"_string);
+}
+
+TEST_CASE(BlockCaptureLongStringAndInt)
+{
+    __block String s = "hello, world, this is a long string to avoid small string optimization"_string;
+    __block int x = 0;
+    auto b = ^{
+        s = "world, hello, this is a long string to avoid small string optimization"_string;
+        x = 2;
+    };
+    auto f = Function<void()>(b);
+
+    f();
+
+    EXPECT_EQ(s, "world, hello, this is a long string to avoid small string optimization"_string);
+    EXPECT_EQ(x, 2);
+}
+
+// Struct definitions from llvm-project/compiler-rt/lib/BlocksRuntime/Block_private.h @ d0177670a0e59e9d9719386f85bb78de0929407c
+struct Block_descriptor {
+    unsigned long int reserved;
+    unsigned long int size;
+    void (*copy)(void* dst, void* src);
+    void (*dispose)(void*);
+};
+
+struct Block_layout {
+    void* isa;
+    int flags;
+    int reserved;
+    void (*invoke)(void*, ...);
+    struct Block_descriptor* descriptor;
+    /* Imported variables. */
+};
+// This check is super important for proper tracking of block closure captures
+static_assert(sizeof(Block_layout) == AK::Detail::block_layout_size);
+
+TEST_CASE(BlockPointerCaptures)
+{
+    int x = 0;
+    int* p = &x;
+
+    auto b = ^{
+        *p = 2;
+    };
+
+    auto f = Function<void()>(b);
+    auto span = f.raw_capture_range();
+
+    int* captured_p = ByteReader::load_pointer<int>(span.data());
+    EXPECT_EQ(captured_p, p);
+
+    f();
+
+    EXPECT_EQ(x, 2);
+}
+
+TEST_CASE(AssignBlock)
+{
+    auto b = ^{ };
+
+    auto f = Function<void()>(b);
+
+    auto b2 = ^{ };
+
+    f = b2;
+
+    f();
+
+    f = b;
+
+    f();
+}
+
+#ifdef AK_HAS_OBJC_ARC
+TEST_CASE(AssignWeakBlock)
+{
+    __block int count = 0;
+    Function<void()> f;
+
+    {
+        auto b = ^{ ++count; };
+        f = b;
+    }
+    f();
+    EXPECT_EQ(count, 1);
+
+    {
+        auto b = ^{ ++count; };
+        auto const __weak weak_b = b;
+        f = weak_b;
+        f();
+        EXPECT_EQ(count, 2);
+    }
+    f();
+    EXPECT_EQ(count, 3);
+}
+#endif