From f9f2e4e313aa0bbb7cdd733bff742e6bb1567abf Mon Sep 17 00:00:00 2001
From: theraven <theraven@72102866-910b-0410-8b05-ffd578937521>
Date: Thu, 13 Oct 2011 10:28:31 +0000
Subject: [PATCH] Added a function for getting the type encoding for a block
 returned by imp_implementationFromBlock().

---
 Test/BlockImpTest.m | 11 ++++++++---
 block_to_imp.c      | 34 ++++++++++++++++++++++++++++++++++
 objc/runtime.h      |  7 +++++++
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/Test/BlockImpTest.m b/Test/BlockImpTest.m
index c6eda71..f43a040 100644
--- a/Test/BlockImpTest.m
+++ b/Test/BlockImpTest.m
@@ -2,6 +2,7 @@
 #include <objc/blocks_runtime.h>
 #include <assert.h>
 #include <stdio.h>
+#include <stdlib.h>
 
 struct big
 {
@@ -23,7 +24,10 @@ int main(void)
 		return b; };
 	blk = Block_copy(blk);
 	IMP imp = imp_implementationWithBlock(blk);
-	class_addMethod((objc_getMetaClass("Foo")), @selector(count:), imp, "i@:i");
+	char *type = block_copyIMPTypeEncoding_np(blk);
+	assert(NULL != type);
+	class_addMethod((objc_getMetaClass("Foo")), @selector(count:), imp, type);
+	free(type);
 	assert(2 == [Foo count: 2]);
 	assert(4 == [Foo count: 2]);
 	assert(6 == [Foo count: 2]);
@@ -36,9 +40,10 @@ int main(void)
 	};
 	imp = imp_implementationWithBlock(blk);
 	assert(imp && "Can't make sret IMP");
-	char *type;
-	asprintf(&type, "%s@:", @encode(struct big));
+	type = block_copyIMPTypeEncoding_np(blk);
+	assert(NULL != type);
 	class_addMethod((objc_getMetaClass("Foo")), @selector(sret), imp, type);
+	free(type);
 	struct big s = [Foo sret];
 	assert(s.a == 1);
 	assert(s.b == 2);
diff --git a/block_to_imp.c b/block_to_imp.c
index 19803c9..66f8dfa 100644
--- a/block_to_imp.c
+++ b/block_to_imp.c
@@ -2,6 +2,8 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <assert.h>
+#include <ctype.h>
 #include <unistd.h>
 #include <sys/mman.h>
 #include "objc/runtime.h"
@@ -123,3 +125,35 @@ BOOL imp_removeBlock(IMP anImp)
 	Block_release(((void**)anImp) - 1);
 	return YES;
 }
+
+PRIVATE size_t lengthOfTypeEncoding(const char *types);
+
+char *block_copyIMPTypeEncoding_np(void*block)
+{
+	char *buffer = strdup(block_getType_np(block));
+	if (NULL == buffer) { return NULL; }
+	char *replace = buffer;
+	// Skip the return type
+	replace += lengthOfTypeEncoding(replace);
+	while (isdigit(*replace)) { replace++; }
+	// The first argument type should be @? (block), and we need to transform
+	// it to @, so we have to delete the ?.  Assert here because this isn't a
+	// block encoding at all if the first argument is not a block, and since we
+	// got it from block_getType_np(), this means something is badly wrong.
+	assert('@' == *replace);
+	replace++;
+	assert('?' == *replace);
+	// Use strlen(replace) not replace+1, because we want to copy the NULL
+	// terminator as well.
+	memmove(replace, replace+1, strlen(replace));
+	// The next argument should be an object, and we want to replace it with a
+	// selector
+	while (isdigit(*replace)) { replace++; }
+	if ('@' != *replace)
+	{
+		free(buffer);
+		return NULL;
+	}
+	*replace = ':';
+	return buffer;
+}
diff --git a/objc/runtime.h b/objc/runtime.h
index 61078e5..86c67ed 100644
--- a/objc/runtime.h
+++ b/objc/runtime.h
@@ -923,6 +923,13 @@ void objc_removeAssociatedObjects(id object);
  * arguments as the method.
  */
 IMP imp_implementationWithBlock(void *block);
+/**
+ * Returns the type encoding of an IMP that would be returned by passing the
+ * block to imp_implementationWithBlock().  Returns NULL if this is not a valid
+ * block encoding for transforming to an IMP (it must take id as its first
+ * argument).  The caller is responsible for freeing the returned value.
+ */
+char *block_copyIMPTypeEncoding_np(void*block);
 /**
  * Returns the block that was used in an IMP created by
  * imp_implementationWithBlock().  The result of calling this function with any