Fix various block issues, see commit message

src/block.zig

@@ -1,5 +1,6 @@
 const std = @import("std");
 const assert = std.debug.assert;
+const Allocator = std.mem.Allocator;
 const objc = @import("main.zig");
 
 // We have to use the raw C allocator for all heap allocation in here
@@ -15,6 +16,11 @@ const alloc = std.heap.raw_c_allocator;
 /// invocation-time arguments to the function. The Return param is the return
 /// type of the function.
 ///
+/// Within the CapturesArg, only `objc.c.id` values will be automatically
+/// memory managed (retained and released) when the block is copied.
+/// If you are passing through NSObjects, you should use the `objc.c.id`
+/// type and recreate a richer Zig type on the other side.
+///
 /// The function that must be implemented is available as the `Fn` field.
 /// The first argument to the function is always a pointer to the `Context`
 /// type (see field in the struct). This has the captured values.
@@ -53,37 +59,65 @@ pub fn Block(
         /// This is the block context sent as the first paramter to the function.
         pub const Context = BlockContext(Captures, InvokeFn);
 
-        /// The context for the block invocations.
-        context: *Context,
-
-        /// Create a new block. This is always created on the heap using the
-        /// libc allocator because the objc runtime expects `malloc` to be
-        /// used.
-        pub fn init(captures: Captures, func: *const Fn) !Self {
-            var ctx = try alloc.create(Context);
-            errdefer alloc.destroy(ctx);
-
-            const flags: BlockFlags = .{ .stret = @typeInfo(Return) == .@"struct" };
+        /// Create a new block context. The block context is what is passed
+        /// (by reference) to functions that request a block.
+        ///
+        /// Note that if the captures contain reference types (like
+        /// NSObject), they will NOT be retained/released UNTIL the block
+        /// is copied. A block copy happens automatically when the block
+        /// is copied to a function that expects a block in ObjC.
+        ///
+        /// If you want to manualy copy a block, you can use the `copy`
+        /// function but you must pair it with a `dispose` function. This
+        /// should only be done for blocks that are not passed to external
+        /// functions where the runtime will automatically copy them (C,
+        /// C++, ObjC, etc.).
+        pub fn init(captures: Captures, func: *const Fn) Context {
+            // The block starts as a stack-allocated block. We let the
+            // runtime copy it to the heap. It doesn't seem to be advisable
+            // to allocate it on the heap directly since the way refcounting
+            // is done and so on is all private API.
+            var ctx: Context = undefined;
             ctx.isa = NSConcreteStackBlock;
-            ctx.flags = @bitCast(flags);
+            ctx.flags = .{
+                .copy_dispose = true,
+                .stret = @typeInfo(Return) == .@"struct",
+                .signature = true,
+            };
             ctx.invoke = @ptrCast(func);
             ctx.descriptor = &descriptor;
             inline for (captures_info.fields) |field| {
                 @field(ctx, field.name) = @field(captures, field.name);
             }
 
-            return .{ .context = ctx };
-        }
-
-        pub fn deinit(self: *Self) void {
-            alloc.destroy(self.context);
-            self.* = undefined;
+            return ctx;
         }
 
         /// Invoke the block with the given arguments. The arguments are
         /// the arguments to pass to the function beyond the captured scope.
-        pub fn invoke(self: *const Self, args: anytype) Return {
-            return @call(.auto, self.context.invoke, .{self.context} ++ args);
+        pub fn invoke(ctx: *const Context, args: anytype) Return {
+            return @call(
+                .auto,
+                ctx.invoke,
+                .{ctx} ++ args,
+            );
+        }
+
+        /// Copies the given context by either literally copying it
+        /// to the heap or increasing the reference count. This must be
+        /// paired with a `release` call to release the block.
+        pub fn copy(ctx: *const Context) Allocator.Error!*Context {
+            const copied = _Block_copy(@ptrCast(@alignCast(ctx))) orelse
+                return error.OutOfMemory;
+            return @ptrCast(@alignCast(copied));
+        }
+
+        /// Release a copied block context. This must only be called on
+        /// contexts returned by the `copy` function. If you pass a block
+        /// context that was not copied, this will crash.
+        pub fn release(ctx: *const Context) void {
+            assert(@intFromPtr(ctx.isa) == @intFromPtr(NSConcreteMallocBlock));
+            _Block_release(@ptrCast(@alignCast(ctx)));
         }
 
         fn descCopyHelper(src: *anyopaque, dst: *anyopaque) callconv(.C) void {
@@ -92,9 +126,9 @@ pub fn Block(
             inline for (captures_info.fields) |field| {
                 if (field.type == objc.c.id) {
                     _Block_object_assign(
-                        @field(real_dst, field.name),
+                        @ptrCast(&@field(real_dst, field.name)),
                         @field(real_src, field.name),
-                        3,
+                        .object,
                     );
                 }
             }
@@ -104,7 +138,10 @@ pub fn Block(
             const real_src: *Context = @ptrCast(@alignCast(src));
             inline for (captures_info.fields) |field| {
                 if (field.type == objc.c.id) {
-                    _Block_object_dispose(@field(real_src, field.name), 3);
+                    _Block_object_dispose(
+                        @field(real_src, field.name),
+                        .object,
+                    );
                 }
             }
         }
@@ -146,7 +183,7 @@ fn BlockContext(comptime Captures: type, comptime InvokeFn: type) type {
     };
     fields[1] = .{
         .name = "flags",
-        .type = c_int,
+        .type = BlockFlags,
         .default_value_ptr = null,
         .is_comptime = false,
         .alignment = @alignOf(c_int),
@@ -201,9 +238,21 @@ fn BlockContext(comptime Captures: type, comptime InvokeFn: type) type {
 
 // Pointer to opaque instead of anyopaque: https://github.com/ziglang/zig/issues/18461
 const NSConcreteStackBlock = @extern(*opaque {}, .{ .name = "_NSConcreteStackBlock" });
+const NSConcreteMallocBlock = @extern(*opaque {}, .{ .name = "_NSConcreteMallocBlock" });
 
-extern "C" fn _Block_object_assign(dst: *anyopaque, src: *const anyopaque, flag: c_int) void;
-extern "C" fn _Block_object_dispose(src: *const anyopaque, flag: c_int) void;
+// https://github.com/llvm/llvm-project/blob/734d31a464e204db699c1cf9433494926deb2aa2/compiler-rt/lib/BlocksRuntime/Block_private.h#L101-L108
+const BlockFieldFlags = enum(c_int) {
+    object = 3, // BLOCK_FIELD_IS_OBJECT
+    block = 7, // BLOCK_FIELD_IS_BLOCK
+    byref = 8, // BLOCK_FIELD_IS_BYREF
+    weak = 16, // BLOCK_FIELD_IS_WEAK
+    byref_caller = 128, // BLOCK_BYREF_CALLER
+};
+
+extern "C" fn _Block_copy(src: *const anyopaque) callconv(.c) ?*anyopaque;
+extern "C" fn _Block_release(src: *const anyopaque) callconv(.c) void;
+extern "C" fn _Block_object_assign(dst: *anyopaque, src: *const anyopaque, flag: BlockFieldFlags) void;
+extern "C" fn _Block_object_dispose(src: *const anyopaque, flag: BlockFieldFlags) void;
 
 const Descriptor = extern struct {
     reserved: c_ulong = 0,
@@ -214,16 +263,16 @@ const Descriptor = extern struct {
 };
 
 const BlockFlags = packed struct(c_int) {
-    _unused: u22 = 0,
+    _unused: u23 = 0,
     noescape: bool = false,
-    _unused_2: bool = false,
-    copy_dispose: bool = true,
+    _unused_2: u1 = 0,
+    copy_dispose: bool = false,
     ctor: bool = false,
-    _unused_3: bool = false,
+    _unused_3: u1 = 0,
     global: bool = false,
-    stret: bool,
-    signature: bool = true,
-    _unused_4: u2 = 0,
+    stret: bool = false,
+    signature: bool = false,
+    _unused_4: u1 = 0,
 };
 
 test "Block" {
@@ -237,13 +286,44 @@ test "Block" {
         .y = 3,
     };
 
-    var block = try AddBlock.init(captures, (struct {
+    var block: AddBlock.Context = AddBlock.init(captures, (struct {
         fn addFn(block: *const AddBlock.Context) callconv(.C) i32 {
             return block.x + block.y;
         }
     }).addFn);
-    defer block.deinit();
 
-    const ret = block.invoke(.{});
+    const ret = AddBlock.invoke(&block, .{});
     try std.testing.expectEqual(@as(i32, 5), ret);
+
+    // Try copy and release
+    const copied = try AddBlock.copy(&block);
+    AddBlock.release(copied);
+}
+
+test "Block copy objc id" {
+    // Create an object, refcount 1
+    const NSObject = objc.getClass("NSObject").?;
+    const obj = NSObject.msgSend(
+        objc.Object,
+        objc.Sel.registerName("alloc"),
+        .{},
+    );
+    _ = obj.msgSend(objc.Object, objc.Sel.registerName("init"), .{});
+
+    const TestBlock = Block(struct {
+        id: objc.c.id,
+    }, .{}, i32);
+
+    var block = TestBlock.init(.{
+        .id = obj.value,
+    }, (struct {
+        fn addFn(block: *const TestBlock.Context) callconv(.C) i32 {
+            _ = block;
+            return 0;
+        }
+    }).addFn);
+
+    // Try copy and release
+    const copied = try TestBlock.copy(&block);
+    TestBlock.release(copied);
 }