Allocate more block trampolines when all sets are full

Test/BlockImpTest.m

@@ -21,6 +21,18 @@ +(int)count: (int)i;
 +(struct big)sret;
 @end
 
+// Makes an IMP from a block that is still on the stack.
+__attribute__((noinline))
+static IMP stackBlockIMP(int value)
+{
+	int (^blk)(id) = ^(id self) { return value; };
+	IMP imp = imp_implementationWithBlock(blk);
+	// The runtime must keep its own copy: this frame is gone by the time the
+	// IMP is called.
+	assert(imp_getBlock(imp) != NULL);
+	assert(imp_getBlock(imp) != blk);
+	return imp;
+}
 
 int main(void)
 {
@@ -67,5 +79,12 @@ int main(void)
 	assert(s.e == 5);
 	imp_removeBlock(imp);
 	assert(imp_getBlock(imp) != blk);
+
+	IMP stackImp = stackBlockIMP(7);
+	IMP stackImp2 = stackBlockIMP(42);
+	assert(7 == ((int(*)(id,SEL))stackImp)(cls, @selector(count:)));
+	assert(42 == ((int(*)(id,SEL))stackImp2)(cls, @selector(count:)));
+	assert(imp_removeBlock(stackImp));
+	assert(imp_removeBlock(stackImp2));
 	return 0;
 }

Test/CMakeLists.txt

@@ -57,7 +57,7 @@ set(TESTS
 )
 
 if (EMBEDDED_BLOCKS_RUNTIME)
-	list(APPEND TESTS BlockImpTest.m)
+	list(APPEND TESTS BlockImpTest.m ManyBlockImps.m)
 endif()
 
 set(ENABLE_ALL_OBJC_ARC_TESTS On)

Test/ManyBlockImps.m

@@ -0,0 +1,187 @@
+#include "../objc/runtime.h"
+#include "../objc/blocks_runtime.h"
+#include <assert.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+/*
+ * Sets and unsets a few thousand blocks as IMPs.  One set of trampolines holds
+ * one page worth of them (128 on AArch64 and 256 on x86-64 with 4 KiB pages),
+ * so the runtime needs many sets and must reuse the trampolines that are freed
+ * in any of them.
+ */
+
+#define COUNT 5000
+
+struct big
+{
+	int a, b, c, d, e;
+};
+
+/*
+ * Makes an IMP from a heap copy of a block, and leaves the IMP with the only
+ * reference to that copy.
+ */
+static IMP heapBlockIMP(id stackBlock)
+{
+	id block = Block_copy(stackBlock);
+	IMP imp = imp_implementationWithBlock(block);
+	assert(imp != NULL);
+	assert(imp_getBlock(imp) == block);
+	Block_release(block);
+	return imp;
+}
+
+/*
+ * Returns an IMP that returns value, or a struct big that starts with value.
+ */
+static IMP makeIMP(int value, BOOL sret)
+{
+	if (sret)
+	{
+		struct big (^block)(id) = ^(id self)
+		{
+			struct big b = { value, value + 1, value + 2, value + 3, value + 4 };
+			return b;
+		};
+		return heapBlockIMP(block);
+	}
+	int (^block)(id) = ^(id self) { return value; };
+	return heapBlockIMP(block);
+}
+
+static int checkBig(struct big b)
+{
+	assert(b.b == b.a + 1);
+	assert(b.c == b.a + 2);
+	assert(b.d == b.a + 3);
+	assert(b.e == b.a + 4);
+	return b.a;
+}
+
+static void checkIMP(IMP imp, int value, BOOL sret)
+{
+	id block = imp_getBlock(imp);
+	assert(block != nil);
+	if (sret)
+	{
+		assert(checkBig(((struct big(*)(id, SEL))imp)(nil, @selector(big))) == value);
+		assert(checkBig(((struct big(^)(id))block)(nil)) == value);
+	}
+	else
+	{
+		assert(((int(*)(id, SEL))imp)(nil, @selector(value)) == value);
+		assert(((int(^)(id))block)(nil) == value);
+	}
+}
+
+static void removeIMP(IMP imp)
+{
+	assert(imp_removeBlock(imp));
+	assert(imp_getBlock(imp) == nil);
+	// Removing it twice must not put the trampoline on the free list twice.
+	assert(!imp_removeBlock(imp));
+}
+
+static int compareIMPs(const void *a, const void *b)
+{
+	uintptr_t x = (uintptr_t)*(const IMP *)a;
+	uintptr_t y = (uintptr_t)*(const IMP *)b;
+	return (x > y) - (x < y);
+}
+
+static void assertDistinct(IMP *imps)
+{
+	IMP *sorted = calloc(COUNT, sizeof(IMP));
+	memcpy(sorted, imps, COUNT * sizeof(IMP));
+	qsort(sorted, COUNT, sizeof(IMP), compareIMPs);
+	for (int i=1 ; i<COUNT ; i++)
+	{
+		assert(sorted[i-1] != sorted[i]);
+	}
+	free(sorted);
+}
+
+static void test(BOOL sret)
+{
+	IMP *imps = calloc(COUNT, sizeof(IMP));
+	int *values = calloc(COUNT, sizeof(int));
+	IMP *freed = calloc(COUNT, sizeof(IMP));
+	int freedCount = 0;
+
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		values[i] = i;
+		imps[i] = makeIMP(values[i], sret);
+	}
+	assertDistinct(imps);
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		checkIMP(imps[i], values[i], sret);
+	}
+
+	// Remove every other IMP, so that every set has free trampolines.
+	for (int i=1 ; i<COUNT ; i+=2)
+	{
+		removeIMP(imps[i]);
+		freed[freedCount++] = imps[i];
+	}
+	for (int i=0 ; i<COUNT ; i+=2)
+	{
+		checkIMP(imps[i], values[i], sret);
+	}
+
+	// Make as many new IMPs.  Only the newest set can still have trampolines
+	// that were never used, so most of them must reuse freed trampolines.
+	qsort(freed, freedCount, sizeof(IMP), compareIMPs);
+	int reused = 0;
+	for (int i=1 ; i<COUNT ; i+=2)
+	{
+		values[i] = COUNT + i;
+		imps[i] = makeIMP(values[i], sret);
+		if (bsearch(&imps[i], freed, freedCount, sizeof(IMP), compareIMPs))
+		{
+			reused++;
+		}
+	}
+	assert(reused >= freedCount / 2);
+	assertDistinct(imps);
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		checkIMP(imps[i], values[i], sret);
+	}
+
+	// Remove them all, then do it all again with every trampoline free.
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		removeIMP(imps[i]);
+	}
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		values[i] = 2 * COUNT + i;
+		imps[i] = makeIMP(values[i], sret);
+	}
+	assertDistinct(imps);
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		checkIMP(imps[i], values[i], sret);
+	}
+	for (int i=0 ; i<COUNT ; i++)
+	{
+		removeIMP(imps[i]);
+	}
+
+	free(imps);
+	free(values);
+	free(freed);
+}
+
+int main(void)
+{
+	test(NO);
+	// Blocks that return structures in memory use their own trampolines on
+	// targets that pass the return slot as the first argument (e.g. x86-64).
+	test(YES);
+	return 0;
+}

block_to_imp.c

@@ -267,11 +267,24 @@ static id invalid(id self, SEL _cmd)
 static struct trampoline_set *alloc_trampolines(char *start, char *end)
 {
 	struct trampoline_set *metadata = calloc(1, sizeof(struct trampoline_set));
+	if (metadata == NULL)
+	{
+		return NULL;
+	}
 #if _WIN32
 	metadata->region = VirtualAlloc(NULL, trampoline_region_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
 #else
 	metadata->region = mmap(NULL, trampoline_region_size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
+	if (metadata->region == MAP_FAILED)
+	{
+		metadata->region = NULL;
+	}
 #endif
+	if (metadata->region == NULL)
+	{
+		free(metadata);
+		return NULL;
+	}
 	struct block_header *headers_start = REGION_HEADERS_START(metadata);
 	char *rx_buffer_start = REGION_RX_BUFFER_START(metadata);
 	for (int i=0 ; i<trampoline_header_per_page ; i++)
@@ -316,36 +329,49 @@ IMP imp_implementationWithBlock(id block)
 	// If we don't have a trampoline intrinsic for this architecture, return a
 	// null IMP.
 	if (0 >= trampolineSize) { return 0; }
-	block = Block_copy(block);
-	// Allocate some trampolines if this is the first time that we need to do this.
-	if (*setptr == NULL)
+	struct trampoline_set *set;
+	for (set=*setptr ; set!=NULL ; set=set->next)
 	{
-		*setptr = alloc_trampolines(start, end);
+		if (set->first_free != -1)
+		{
+			break;
+		}
 	}
-	for (struct trampoline_set *set=*setptr ; set!=NULL ; set=set->next)
+	// If this is the first time that we need trampolines, or all of the sets
+	// are full, then allocate a new set.  Put it at the start of the list, so
+	// that the next allocations find a free trampoline straight away, but keep
+	// the older sets in the list: imp_getBlock() and imp_removeBlock() must
+	// still find their IMPs, and trampolines freed in them must be reused.
+	if (set == NULL)
 	{
-		if (set->first_free != -1)
+		set = alloc_trampolines(start, end);
+		if (set == NULL)
 		{
-			int i = set->first_free;
-			struct block_header *headers_start = REGION_HEADERS_START(set);
-			char *rx_buffer_start = REGION_RX_BUFFER_START(set);
-			struct block_header *h = &headers_start[i];
-			struct block_header *next = h->block;
-			set->first_free = next ? (next - headers_start) : -1;
-			assert(set->first_free < trampoline_header_per_page);
-			assert(set->first_free >= -1);
-			h->fnptr = (void(*)(void))b->invoke;
-			h->block = b;
-			uintptr_t addr = (uintptr_t)&rx_buffer_start[i*sizeof(struct block_header)];
-#if (__ARM_ARCH_ISA_THUMB == 2)
-			// If the trampoline is Thumb-2 code, then we must set the low bit
-			// to 1 so that b[l]x instructions put the CPU in the correct mode.
-			addr |= 1;
-#endif
-			return (IMP)addr;
+			return 0;
 		}
+		set->next = *setptr;
+		*setptr = set;
 	}
-	UNREACHABLE("Failed to allocate block");
+	int i = set->first_free;
+	struct block_header *headers_start = REGION_HEADERS_START(set);
+	char *rx_buffer_start = REGION_RX_BUFFER_START(set);
+	struct block_header *h = &headers_start[i];
+	struct block_header *next = h->block;
+	set->first_free = next ? (next - headers_start) : -1;
+	assert(set->first_free < (int)trampoline_header_per_page);
+	assert(set->first_free >= -1);
+	// Store the copy, not the block that we were passed: if that one is on the
+	// stack, then it is gone once the caller returns.
+	b = Block_copy(b);
+	h->fnptr = (void(*)(void))b->invoke;
+	h->block = b;
+	uintptr_t addr = (uintptr_t)&rx_buffer_start[i*sizeof(struct block_header)];
+#if (__ARM_ARCH_ISA_THUMB == 2)
+	// If the trampoline is Thumb-2 code, then we must set the low bit
+	// to 1 so that b[l]x instructions put the CPU in the correct mode.
+	addr |= 1;
+#endif
+	return (IMP)addr;
 }
 
 static int indexForIMP(IMP anIMP, struct trampoline_set **setptr)
@@ -379,7 +405,13 @@ id imp_getBlock(IMP anImp)
 	{
 		return NULL;
 	}
-	return REGION_HEADERS_START(set)[idx].block;
+	struct block_header *h = &REGION_HEADERS_START(set)[idx];
+	// A free trampoline's block field is a free list link, not a block.
+	if (h->fnptr == (void(*)(void))invalid)
+	{
+		return NULL;
+	}
+	return h->block;
 }
 
 BOOL imp_removeBlock(IMP anImp)
@@ -398,6 +430,11 @@ BOOL imp_removeBlock(IMP anImp)
 	}
 	struct block_header *header_start = REGION_HEADERS_START(set);
 	struct block_header *h = &header_start[idx];
+	// Removing a free trampoline again would corrupt the free list.
+	if (h->fnptr == (void(*)(void))invalid)
+	{
+		return NO;
+	}
 	Block_release(h->block);
 	h->fnptr = (void(*)(void))invalid;
 	h->block = set->first_free == -1 ? NULL : &header_start[set->first_free];