From 5217176c534982d17c70bc18b9068e255686f6ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wilfredo=20Vel=C3=A1zquez-Rodr=C3=ADguez?=
 <will.velazquez@flexibits.com>
Date: Wed, 30 Sep 2026 13:48:21 -0400
Subject: [PATCH] Fix memory leak of block retained by another block

---
 Test/BlockCapture_arc.m | 58 +++++++++++++++++++++++++++++++++++++++++
 Test/CMakeLists.txt     |  1 +
 arc.mm                  | 14 ++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 Test/BlockCapture_arc.m

diff --git a/Test/BlockCapture_arc.m b/Test/BlockCapture_arc.m
new file mode 100644
index 00000000..551e085a
--- /dev/null
+++ b/Test/BlockCapture_arc.m
@@ -0,0 +1,58 @@
+#include "Test.h"
+
+// Regression test: objc_retain must return the same pointer it was called
+// with, regardless of whether the argument is a stack or heap block.  Clang's
+// ARC codegen declares llvm.objc.retain with the `returned` attribute; at -O2+
+// the optimizer relies on that to replace uses of the retain result with the
+// original pointer.  A previous version of libobjc2's objc_retain forwarded
+// unconditionally to Block_copy, which for stack blocks allocates a heap copy
+// and returns a different pointer.  The optimizer discarded the heap-copy
+// pointer and later released the original stack pointer (a no-op for stack
+// blocks), so the heap copy and every strong reference it captured leaked.
+
+static int liveCount;
+
+@interface Capture : Test @end
+@implementation Capture
++ (id)new
+{
+	liveCount++;
+	return [super new];
+}
+- (void)dealloc
+{
+	liveCount--;
+}
+- (int)value
+{
+	return 42;
+}
+@end
+
+// A method (not a plain function) so message dispatch prevents inlining and
+// forces the compiler to treat the block parameter as escapable.
+@interface Runner : Test @end
+@implementation Runner
+- (void)runWith:(int (^)(void))b
+{
+	// Capturing `b` into another block requires ARC to retain the block
+	// value — this is the exact site where the bug fires.
+	int (^wrapper)(void) = ^{ return b(); };
+	(void)wrapper();
+}
+@end
+
+int main(void)
+{
+	Runner *r = [Runner new];
+	for (int i = 0; i < 100; i++)
+	{
+		@autoreleasepool
+		{
+			Capture *c = [Capture new];
+			[r runWith:^{ return [c value]; }];
+		}
+	}
+	assert(liveCount == 0);
+	return 0;
+}
diff --git a/Test/CMakeLists.txt b/Test/CMakeLists.txt
index b88dda39..f2c0544c 100644
--- a/Test/CMakeLists.txt
+++ b/Test/CMakeLists.txt
@@ -17,6 +17,7 @@ set(TESTS
 	AllocatePair.m
 	AssociatedObject.m
 	AssociatedObject2.m
+	BlockCapture_arc.m
 	BlockTest_arc.m
 	ConstantString.m
 	Category.m
diff --git a/arc.mm b/arc.mm
index 10d4acb6..2a49b6be 100644
--- a/arc.mm
+++ b/arc.mm
@@ -496,6 +496,20 @@ static inline id retain(id obj, BOOL isWeak)
 	Class cls = obj->isa;
 	if (UNLIKELY(objc_test_class_flag(cls, objc_class_flag_is_block)))
 	{
+		// objc_retain is declared with the `returned` attribute, so callers
+		// (and the LLVM ARC optimizer) assume the return value equals the
+		// input.  Block_copy honours that for heap blocks (it just bumps the
+		// refcount) but violates it for stack blocks, where it allocates a
+		// new heap copy with a different address.  Under -O2 the optimizer
+		// then discards the returned heap pointer and later releases the
+		// original stack pointer — which is a no-op — leaking the heap copy
+		// and everything it captured.  Return the input for stack blocks;
+		// any real escape goes through _Block_object_assign / objc_retainBlock,
+		// which do the Block_copy correctly.
+		if (cls == static_cast<void*>(&_NSConcreteStackBlock))
+		{
+			return obj;
+		}
 		return Block_copy(obj);
 	}
 	if (objc_test_class_flag(cls, objc_class_flag_fast_arc))