add support for groups in ZKSwizzleInterface

alexzielenski · Jun 12, 2015 · 6d0947f · 6d0947f
1 parent ce1fb4a
commit 6d0947f
Show file tree

Hide file tree

Showing 3 changed files with 150 additions and 29 deletions.
diff --git a/ZKSwizzle/ZKSwizzle.h b/ZKSwizzle/ZKSwizzle.h
@@ -8,30 +8,26 @@
 
 #import <Foundation/Foundation.h>
 #import <objc/runtime.h>
-#include <sys/cdefs.h>
+#import <sys/cdefs.h>
 
 // This is a class for streamlining swizzling. Simply create a new class of any name you want and
-
 // Example:
 /*
- 
  @interface ZKHookClass : NSObject
- 
  - (NSString *)description; // hooks -description on NSObject
  - (void)addedMethod; // all subclasses of NSObject now respond to -addedMethod
- 
  @end
  
  @implementation ZKHookClass
- 
  ...
- 
  @end
  
  [ZKSwizzle swizzleClass:ZKClass(ZKHookClass) forClass:ZKClass(destination)];
- 
  */
 
+#ifndef ZKSWIZZLE_DEFS
+#define ZKSWIZZLE_DEFS
+
 // CRAZY MACROS FOR DYNAMIC PROTOTYPE CREATION
 #define VA_NUM_ARGS(...) VA_NUM_ARGS_IMPL(0, ## __VA_ARGS__, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5 ,4 ,3 ,2, 1, 0)
 #define VA_NUM_ARGS_IMPL(_0, _1,_2,_3,_4,_5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20 ,N,...) N
@@ -62,7 +58,6 @@
 #define INVOKE(MACRO, NUMBER, ...) CAT(MACRO, NUMBER)(__VA_ARGS__)
 #define WRAP_LIST(...) INVOKE(WRAP, VA_NUM_ARGS(__VA_ARGS__), __VA_ARGS__)
 
-
 // Gets the a class with the name CLASS
 #define ZKClass(CLASS) objc_getClass(#CLASS)
 
@@ -71,39 +66,51 @@
 #define ZKHookIvar(OBJECT, TYPE, NAME) (*(TYPE *)ZKIvarPointer(OBJECT, NAME))
 #else
 #define ZKHookIvar(OBJECT, TYPE, NAME) \
-_Pragma("clang diagnostic push") \
-_Pragma("clang diagnostic ignored \"-Wignored-attributes\"") \
-(*(__unsafe_unretained TYPE *)ZKIvarPointer(OBJECT, NAME)) \
-_Pragma("clang diagnostic pop")
+    _Pragma("clang diagnostic push") \
+    _Pragma("clang diagnostic ignored \"-Wignored-attributes\"") \
+    (*(__unsafe_unretained TYPE *)ZKIvarPointer(OBJECT, NAME)) \
+    _Pragma("clang diagnostic pop")
 #endif
 // returns the original implementation of the swizzled function or null or not found
 #define ZKOrig(TYPE, ...) ((TYPE (*)(id, SEL WRAP_LIST(__VA_ARGS__)))(ZKOriginalImplementation(self, _cmd, __PRETTY_FUNCTION__)))(self, _cmd, ##__VA_ARGS__)
 
 // returns the original implementation of the superclass of the object swizzled
 #define ZKSuper(TYPE, ...) ((TYPE (*)(id, SEL WRAP_LIST(__VA_ARGS__)))(ZKSuperImplementation(self, _cmd, __PRETTY_FUNCTION__)))(self, _cmd, ##__VA_ARGS__)
 
-// Ripped off from MobileSubstrate
+#define _ZKSwizzleInterfaceConditionally(CLASS_NAME, TARGET_CLASS, SUPERCLASS, GROUP, IMMEDIATELY) \
+    @interface _$ ## CLASS_NAME : SUPERCLASS @end \
+    @implementation _$ ## CLASS_NAME \
+    + (void)initialize {} \
+    @end \
+    @interface CLASS_NAME : _$ ## CLASS_NAME @end \
+    @implementation CLASS_NAME (ZKSWIZZLE) \
+    + (void)load { \
+        _$ZKRegisterInterface(self, #GROUP);\
+        if (IMMEDIATELY) { \
+            [self _ZK_unconditionallySwizzle]; \
+        } \
+    } \
+    + (void)_ZK_unconditionallySwizzle { \
+        ZKSwizzle(CLASS_NAME, TARGET_CLASS); \
+    } \
+@end
+
 // Bootstraps your swizzling class so that it requires no setup
 // outside of this macro call
 // If you override +load you must call ZKSwizzle(CLASS_NAME, TARGET_CLASS)
 // yourself, otherwise the swizzling would not take place
 #define ZKSwizzleInterface(CLASS_NAME, TARGET_CLASS, SUPERCLASS) \
-@interface _$ ## CLASS_NAME : SUPERCLASS @end \
-@implementation _$ ## CLASS_NAME \
-+ (void)initialize {} \
-@end \
-@interface CLASS_NAME : _$ ## CLASS_NAME @end \
-@implementation CLASS_NAME (ZKSWIZZLE) \
-+ (void)load { \
-ZKSwizzle(CLASS_NAME, TARGET_CLASS); \
-} \
-@end
+    _ZKSwizzleInterfaceConditionally(CLASS_NAME, TARGET_CLASS, SUPERCLASS, ZK_UNGROUPED, YES)
 
-// thanks OBJC_OLD_DISPATCH_PROTOTYPES=0
-typedef id (*ZKIMP)(id, SEL, ...);
+// Same as ZKSwizzleInterface, except
+#define ZKSwizzleInterfaceGroup(CLASS_NAME, TARGET_CLASS, SUPER_CLASS, GROUP) \
+    _ZKSwizzleInterfaceConditionally(CLASS_NAME, TARGET_CLASS, SUPER_CLASS, GROUP, NO)
 
 __BEGIN_DECLS
 
+// Make sure to cast this before you use it
+typedef id (*ZKIMP)(id, SEL, ...);
+
 // returns a pointer to the instance variable "name" on the object
 void *ZKIvarPointer(id self, const char *name);
 // returns the original implementation of a method with selector "sel" of an object hooked by the methods below
@@ -116,9 +123,14 @@ ZKIMP ZKSuperImplementation(id object, SEL sel, const char *info);
 #define ZKSwizzle(src, dst) _ZKSwizzle(ZKClass(src), ZKClass(dst))
 BOOL _ZKSwizzle(Class src, Class dest);
 
+#define ZKSwizzleGroup(NAME) _ZKSwizzleGroup(#NAME)
+void _$ZKRegisterInterface(Class cls, const char *groupName);
+BOOL _ZKSwizzleGroup(const char *groupName);
+
 // Calls above method with the superclass of source for desination
 #define ZKSwizzleClass(src) _ZKSwizzleClass(ZKClass(src))
 BOOL _ZKSwizzleClass(Class cls);
 
 __END_DECLS
+#endif
 
diff --git a/ZKSwizzle/ZKSwizzle.m b/ZKSwizzle/ZKSwizzle.m
@@ -9,6 +9,10 @@
 #import "ZKSwizzle.h"
 static NSMutableDictionary *classTable;
 
+@interface NSObject (ZKSwizzle)
++ (void)_ZK_unconditionallySwizzle;
+@end
+
 void *ZKIvarPointer(id self, const char *name) {
     Ivar ivar = class_getInstanceVariable(object_getClass(self), name);
     return ivar == NULL ? NULL : (__bridge void *)self + ivar_getOffset(ivar);
@@ -156,7 +160,7 @@ BOOL _ZKSwizzle(Class src, Class dest) {
 
     if ([classTable objectForKey:NSStringFromClass(src)]) {
         [NSException raise:@"Invalid Argument"
-                    format:@"This source class was already swizzled with another"];
+                    format:@"This source class (%@) was already swizzled with another, (%@)", NSStringFromClass(src), classTable[NSStringFromClass(src)]];
         return NO;
     }
 
@@ -173,6 +177,12 @@ BOOL _ZKSwizzleClass(Class cls) {
 }
 
 static BOOL enumerateMethods(Class destination, Class source) {
+#if OBJC_API_VERSION < 2
+    [NSException raise:@"Unsupported feature" format:@"ZKSwizzle is only available in objc 2.0"];
+    return NO;
+
+#else
+
     unsigned int methodCount;
     Method *methodList = class_copyMethodList(source, &methodCount);
     BOOL success = YES;
@@ -181,6 +191,11 @@ static BOOL enumerateMethods(Class destination, Class source) {
         SEL selector  = method_getName(method);
         NSString *methodName = NSStringFromSelector(selector);
 
+        // Don't do anything with the unconditional swizzle
+        if (sel_isEqual(selector, @selector(_ZK_unconditionallySwizzle))) {
+            continue;
+        }
+
         // We only swizzle methods that are implemented
         if (class_respondsToSelector(destination, selector)) {
             Method originalMethod = class_getInstanceMethod(destination, selector);
@@ -230,5 +245,47 @@ static BOOL enumerateMethods(Class destination, Class source) {
 
     free(propertyList);
     free(methodList);
+    return success;
+#endif
+}
+
+// Options were to use a group class and traverse its subclasses
+// or to create a groups dictionary
+// This works because +load on NSObject is called before attribute((constructor))
+static NSMutableDictionary *groups = nil;
+void _$ZKRegisterInterface(Class cls, const char *groupName) {
+    if (!groups)
+        groups = [NSMutableDictionary dictionary];
+
+    NSString *groupString = @(groupName);
+    NSMutableArray *groupList = groups[groupString];
+    if (!groupList) {
+        groupList = [NSMutableArray array];
+        groups[groupString] = groupList;
+    }
+
+    [groupList addObject:NSStringFromClass(cls)];
+}
+
+BOOL _ZKSwizzleGroup(const char *groupName) {
+    NSArray *groupList = groups[@(groupName)];
+    if (!groupList) {
+        [NSException raise:@"Invalid Argument" format:@"ZKSwizzle: There is no group by the name of %s", groupName];
+        return NO;
+    }
+
+    BOOL success = YES;
+    for (NSString *className in groupList) {
+        Class cls = NSClassFromString(className);
+        if (cls == NULL)
+            continue;
+
+        if (class_respondsToSelector(object_getClass(cls), @selector(_ZK_unconditionallySwizzle))) {
+            [cls _ZK_unconditionallySwizzle];
+        } else {
+            success = NO;
+        }
+    }
+
     return success;
 }
diff --git a/ZKSwizzleTests/ZKTests.m b/ZKSwizzleTests/ZKTests.m
@@ -122,18 +122,64 @@ - (NSString *)description {
 
 @end
 
+@interface GroupClass : NSObject
++ (NSString *)classMethod;
+- (NSString *)instanceMethod;
+@end
+
+@implementation GroupClass
+
++ (NSString *)classMethod {
+    return @"classMethod";
+}
+
+- (NSString *)instanceMethod {
+    return @"instanceMethod";
+}
+
+@end
+
+// Swizzled Group
+ZKSwizzleInterfaceGroup(GroupSwizzle, GroupClass, NSObject, Yosemite)
+@implementation GroupSwizzle
+
++ (NSString *)classMethod {
+    return @"swizzled";
+}
+
+- (NSString *)instanceMethod {
+    return @"swizzled";
+}
+
+@end
+
+// Unswizzled Group – unused
+ZKSwizzleInterfaceGroup(GroupSwizzle2, GroupClass, NSObject, Mavericks)
+@implementation GroupSwizzle2
+
++ (NSString *)classMethod {
+    return @"swizzled2";
+}
+
+- (NSString *)instanceMethod {
+    return @"swizzled2";
+}
+
+@end
+
+
 @interface ZKTests : XCTestCase
 @end
 
 @implementation ZKTests
 
 - (void)setUp {
     [super setUp];
-    _ZKSwizzleClass(ZKClass(ZKSwizzlerClass));
-    _ZKSwizzleClass(ZKClass(ZKSwizzlerClass2));
 }
 
 - (void)testExample {
+    _ZKSwizzleClass(ZKClass(ZKSwizzlerClass));
+    _ZKSwizzleClass(ZKClass(ZKSwizzlerClass2));
     ZKOriginalClass *instance = [[ ZKOriginalClass alloc] init];
     XCTAssertEqualObjects([ ZKOriginalClass classMethod], @"replaced", @"replacing class methods");
     XCTAssertEqualObjects([instance instanceMethod], @"replaced", @"replacing instance methods");
@@ -147,4 +193,10 @@ - (void)testExample {
     XCTAssertEqualObjects([[[NewClass alloc] init] description], @"DummyClass", @"ZKSuper outside of swizzling");
 }
 
+- (void)testGroups {
+    ZKSwizzleGroup(Yosemite);
+    XCTAssertEqualObjects([[[GroupClass alloc] init] instanceMethod], @"swizzled");
+    XCTAssertEqualObjects([GroupClass classMethod], @"swizzled");
+}
+
 @end