1 files changed, 34 insertions, 9 deletions

diff --git a/test/test-imp-forwarding.m b/test/test-imp-forwarding.m
index e7bff27..0beb24a 100644
--- a/test/test-imp-forwarding.m
+++ b/test/test-imp-forwarding.m
@@ -1,8 +1,14 @@
 #include "../lib/objc.c"
 #include <objc/runtime.h>
 #include <stdio.h>
+#include <assert.h>
 #import <Foundation/Foundation.h>
 
+static void *what_to_call(void *a, void *b) {
+    printf("what_to_call: %p %p\n", a, b);
+    return a;
+}
+
 static void imp1(id self, SEL sel, int a, int b) {
     NSLog(@"imp1: self=%@ sel=%s a=%d b=%d\n", self, sel_getName(sel), a, b);
 }
@@ -13,7 +19,7 @@ struct big {
 
 static struct big imp2(id self, SEL sel, int a, int b) {
     NSLog(@"imp2: self=%@ sel=%s a=%d b=%d\n", self, sel_getName(sel), a, b);
-    return (struct big) {{0}};
+    return (struct big) {{4}};
 }
 
 struct big (^test)(id, int) = ^(id self, int a) {
@@ -22,14 +28,33 @@ struct big (^test)(id, int) = ^(id self, int a) {
 };
 
 int main() {
-    IMP testi = imp_implementationWithBlock(test);
-    ((struct big (*)(id, SEL, int)) testi)(@"test", @selector(dumb), 5);
-    IMP old = (IMP) imp1;
     SEL sel = @selector(some);
-    struct temp_block_literal temp_block = get_temp_block(&old, sel);
-    IMP new = imp_implementationWithBlock((id) &temp_block);
-    ((void (*)(id, SEL, int, int)) new)(@"foo", sel, 1, 2);
-    old = (IMP) imp2;
-    struct big big = ((struct big (*)(id, SEL, int, int)) new)(@"bar", sel, 1, 2);
+    IMP imp;
+    assert(!get_trampoline(what_to_call, imp1, (void *) 0x123, &imp));
+    printf("imp = %p\n", imp);
+    ((void (*)(id, SEL, int, int)) imp)(@"foo", sel, 1, 2);
+    free_trampoline(imp);
+    assert(!get_trampoline(what_to_call, imp2, (void *) 0x123, &imp));
+    struct big big = ((struct big (*)(id, SEL, int, int)) imp)(@"bar", sel, 1, 2);
     printf("out? %d\n", big.x[0]);
+
+    /* test alloc/free */
+    enum { n = 10000 };
+    static void *imps[n];
+    for(size_t i = 0; i < n; i++) {
+        assert(!get_trampoline(what_to_call, imp1, (void *) 0x123, &imps[i]));
+    }
+    for(size_t i = 0; i < n; i++) {
+        assert(imps[i]);
+        if (arc4random() & 1) {
+            free_trampoline(imps[i]);
+            imps[i] = NULL;
+        }
+    }
+    for(size_t i = 0; i < n; i++) {
+        if (imps[i]) {
+            free_trampoline(imps[i]);
+            imps[i] = NULL;
+        }
+    }
 }