Copy a method IMP for multiple method swizzles - objective-c

Question

I have a class set up that ideally will read the methods of any class passed in and then map all of them to on single selector at runtime before forwarding them off to their original selector.
This does work right now, but I can only do it for one method at a time. The issue seems to be that once I swizzle the first method, my IMP to catch and forward the method has now been swapped with that other methods IMP. Any further attempts at this screw up because they use newly swapped IMP to replace the others.
1)So I have MethodA, MethodB, and CustomCatchAllMethod.
2)I swap MethodA with CustomCatchAllMEthod. MethodA->CustomCatchAllMethod, CustomCatchAllMethod->MethodA
3)Now I try to swap to MethodB with CustomCatchAllMethod as well, but since CustomCatchAllMethod now = MethodA, MethodB becomes MethodA and MethodA->MethodB.
So how do I get/copy a new instance of my IMP for each new selector I want to intercept?
Here's a rough mockup of the above flow:
void swizzle(Class classImCopying, SEL orig){
SEL new = #selector(catchAll:);
Method origMethod = class_getInstanceMethod(classImCopying, orig);
Method newMethod = class_getInstanceMethod(catchAllClass,new);
method_exchangeImplementations(origMethod, newMethod);
}
//In some method elsewhere
//I want this to make both methodA and methodB point to catchAll:
swizzle(someClass, #selector(methodA:));
swizzle(someClass, #selector(methodB:));

Answer 1

That common method-swizzling pattern only works when you want to intercept one method with one other. In your case you are basically moving the implementation for catchAll: around instead of inserting it everywhere.
To properly to this you'd have to use:
IMP imp = method_getImplementation(newMethod);
method_setImplementation(origMethod, imp);
This leaves you with one problem though: how to forward to the original implementation?
That is what the original pattern used exchangeImplementations for.
In your case you could:
keep a table of the original IMPs around or
rename the original methods with some common prefix, so you can build a call to them from catchAll:
Note that you can only handle methods of the same arity when you want to forward everything through the same method.

Answer 2

You can capture original IMP with block, get block's IMP and set it as implementation of method.
Method method = class_getInstanceMethod(class, setterSelector);
SEL selector = method_getName(method);
IMP originalImp = method_getImplementation(method);
id(^block)(id self, id arg) = ^id(id self, id arg) {
return ((id(*)(id, SEL, id))originalImp)(self, selector, arg);
};
IMP newImp = imp_implementationWithBlock(block);
method_setImplementation(method, newImp);