Kunal Parmar avatar Kunal Parmar committed 183c769

Swizzling utilities.

Comments (0)

Files changed (3)

Sources/KPTNSObject+Foundation.h

 
 @interface NSObject (KPTNSObjectFoundation)
 
+// Swizzle a class method.
+//
+// Discussion
+//  Objective-C categories allow you to replace the behavior of any class. But
+//  often times, you dont want to replace, but instead want to add behavior.
+//  Using categories, you do not have access to the original implementation.
+//  Swizzling allows you to do that.
+//
+//  Example, you want to detect if there are memory leaks. You can swizzle the
+//  +alloc and -dealloc methods and keep track of a counter to determine if
+//  any instances were left behind.
+//
+//    @implementation MyClass (Swizzling)
+//
+//    + (id)swizzledAlloc {
+//      counter++;
+//
+//      // Call the original code.
+//      return [self swizzledAlloc];
+//    }
+//
+//    - (void)swizzledDealloc {
+//      counter--;
+//
+//      // Call the original code.
+//      [self swizzledDealloc];
+//    }
+//
+//    @end
+//
+//    int counter = 0;
+//    main() {
+//      [MyClass swizzleClassMethod:@selector(alloc)
+//                       withMethod:@selector(swizzledAlloc)
+//                            error:nil];
+//      [MyClass swizzleInstanceMethod:@selector(dealloc)
+//                          withMethod:@selector(swizzledDealloc)
+//                               error:nil];
+//
+//       ...
+//
+//      assert(counter == 0);
+//    }
+//
+//  NOTE: to call the original code, you call the swizzledAlloc method. This
+//  works because the implementation of swizzledAlloc method points to the
+//  original code after swizzling.
+//
+//  IMPORTANT:
+//  Swizzling does not work in a world where you have multiple classes deriving
+//  from one base class and you are trying to swizzle methods in them. You will
+//  run into strange bugs if you do that.
++ (BOOL)swizzleClassMethod:(SEL)originalSelector
+                withMethod:(SEL)newSelector
+                     error:(NSError **)error;
+
+// Swizzle an instance method.
+//
+// Discussion
+//  See |+swizzleClassMethod:withMethod:error:| for details.
++ (BOOL)swizzleInstanceMethod:(SEL)originalSelector
+                   withMethod:(SEL)newSelector
+                        error:(NSError **)error;
+
 // Safer way of doing +cancelPreviousPerformRequestsWithTarget:selector:object:
 // followed by -performSelector:withObject:afterDelay:.
 //

Sources/KPTNSObject+Foundation.m

 
 #import "KPTNSObject+Foundation.h"
 
+#import <objc/objc-runtime.h>
+
+static void MissingMethodError(Class klass, SEL selector, NSError **error)
+{
+  if (error == nil)
+  {
+    return;
+  }
+
+  NSString *description = [NSString stringWithFormat:@"Method %@ not implemented by class %@",
+                           NSStringFromSelector(selector),
+                           NSStringFromClass(klass)];
+  NSDictionary *userInfo = [NSDictionary dictionaryWithObject:description
+                                                       forKey:NSLocalizedDescriptionKey];
+  *error = [NSError errorWithDomain:NSCocoaErrorDomain
+                               code:-1
+                           userInfo:userInfo];
+}
+
+
 @implementation NSObject (KPTNSObjectFoundation)
 
++ (BOOL)swizzleClassMethod:(SEL)originalSelector
+                withMethod:(SEL)newSelector
+                     error:(NSError **)error
+{
+  return [object_getClass(self) swizzleInstanceMethod:originalSelector
+                                           withMethod:newSelector
+                                                error:error];
+}
+
++ (BOOL)swizzleInstanceMethod:(SEL)originalSelector
+                   withMethod:(SEL)newSelector
+                        error:(NSError **)error
+{
+  Method originalMethod = class_getInstanceMethod(self, originalSelector);
+  if (originalMethod == nil)
+  {
+    MissingMethodError(self, originalSelector, error);
+    return NO;
+  }
+
+  Method newMethod = class_getInstanceMethod(self, newSelector);
+  if (newMethod == nil)
+  {
+    MissingMethodError(self, newSelector, error);
+    return NO;
+  }
+
+  // There are two scenarios:
+  //  a. the method to swizzle i.e. |originalMethod| is not implemented by the
+  //     class, but is inherited from a superclass
+  //  b. the method to swizzle is implemented by the class
+  //
+  // For a), |originalMethod| is added to the class with the implementation of
+  // |newMethod| i.e. in the swizzled state and then the implementation of
+  // |newMethod| is replaced by that of |originalMethod|.
+  // NOTE: |class_addMethod| will add an override of a superclass's
+  // implementation and return YES, but will not replace an existing
+  // implementation in the class, returning NO.
+  //
+  // For b), the implementations of the methods are simply swapped.
+  if (class_addMethod(self,
+                      originalSelector,
+                      method_getImplementation(newMethod),
+                      method_getTypeEncoding(newMethod)))
+  {
+    class_replaceMethod(self,
+                        newSelector,
+                        method_getImplementation(originalMethod),
+                        method_getTypeEncoding(originalMethod));
+  }
+  else
+  {
+    method_exchangeImplementations(originalMethod, newMethod);
+  }
+
+  return YES;
+}
+
 - (void)coalescedPerformSelector:(SEL)aSelector
                       withObject:(id)anArgument
                       afterDelay:(NSTimeInterval)delay

Sources/KPTNSObject+FoundationTest.m

 
 #import "KPTNSObject+Foundation.h"
 
+@interface TestClassForSwizzlingBase : NSObject
+
++ (NSString *)inheritedClassMethod;
++ (NSString *)overridenClassMethod;
+- (NSString *)inheritedInstanceMethod;
+- (NSString *)overridenInstanceMethod;
+
+@end
+
+@implementation TestClassForSwizzlingBase
+
++ (NSString *)inheritedClassMethod
+{
+  return @"inheritedClassMethod";
+}
+
++ (NSString *)overridenClassMethod
+{
+  return nil;
+}
+
+- (NSString *)inheritedInstanceMethod
+{
+  return @"inheritedInstanceMethod";
+}
+
+- (NSString *)overridenInstanceMethod
+{
+  return nil;
+}
+
+@end
+
+
+@interface TestClassForSwizzling : TestClassForSwizzlingBase
+
++ (NSString *)overridenClassMethod;
+- (NSString *)overridenInstanceMethod;
+
+@end
+
+@implementation TestClassForSwizzling
+
++ (NSString *)overridenClassMethod
+{
+  return @"overridenClassMethod";
+}
+
+- (NSString *)overridenInstanceMethod
+{
+  return @"overridenInstanceMethod";
+}
+
+@end
+
+
+@interface TestClassForSwizzling (SwizzledMethod)
+
++ (NSString *)swizzledInheritedClassMethod;
++ (NSString *)swizzledOverridenClassMethod;
+- (NSString *)swizzledInheritedInstanceMethod;
+- (NSString *)swizzledOverridenInstanceMethod;
+
+@end
+
+@implementation TestClassForSwizzling (SwizzledMethod)
+
++ (NSString *)swizzledInheritedClassMethod
+{
+  return @"swizzledInheritedClassMethod";
+}
+
++ (NSString *)swizzledOverridenClassMethod
+{
+  return @"swizzledOverridenClassMethod";
+}
+
+- (NSString *)swizzledInheritedInstanceMethod
+{
+  return @"swizzledInheritedInstanceMethod";
+}
+
+- (NSString *)swizzledOverridenInstanceMethod
+{
+  return @"swizzledOverridenInstanceMethod";
+}
+
+@end
+
+
 @interface KPTNSObject_FoundationTest : SenTestCase
 @end
 
 @implementation KPTNSObject_FoundationTest
 
+- (void)testSwizzlingClassMethodOverridenByClass
+{
+  STAssertEqualObjects([TestClassForSwizzling overridenClassMethod],
+                       @"overridenClassMethod",
+                       @"Overriden class method not setup correctly");
+  NSError *error = nil;
+  STAssertTrue([TestClassForSwizzling swizzleClassMethod:@selector(overridenClassMethod)
+                                              withMethod:@selector(swizzledOverridenClassMethod)
+                                                   error:&error],
+               @"Failed to swizzle overriden class method; error %@", error);
+  STAssertEqualObjects([TestClassForSwizzling overridenClassMethod],
+                       @"swizzledOverridenClassMethod",
+                       @"Overriden class method returned wrong value after swizzling");
+  STAssertEqualObjects([TestClassForSwizzling swizzledOverridenClassMethod],
+                       @"overridenClassMethod",
+                       @"Overriden class method returned wrong value after swizzling");
+
+  STAssertNil([TestClassForSwizzlingBase overridenClassMethod],
+              @"Overriden class method was swizzled in base class");
+}
+
+- (void)testSwizzlingClassMethodInheritedByClass
+{
+  STAssertEqualObjects([TestClassForSwizzling inheritedClassMethod],
+                       @"inheritedClassMethod",
+                       @"Inherited class method not setup correctly");
+  NSError *error = nil;
+  STAssertTrue([TestClassForSwizzling swizzleClassMethod:@selector(inheritedClassMethod)
+                                              withMethod:@selector(swizzledInheritedClassMethod)
+                                                   error:&error],
+               @"Failed to swizzle inherited class method; error %@", error);
+  STAssertEqualObjects([TestClassForSwizzling inheritedClassMethod],
+                       @"swizzledInheritedClassMethod",
+                       @"Inherited class method returned wrong value after swizzling");
+  STAssertEqualObjects([TestClassForSwizzling swizzledInheritedClassMethod],
+                       @"inheritedClassMethod",
+                       @"Inherited class method returned wrong value after swizzling");
+
+  STAssertEqualObjects([TestClassForSwizzlingBase inheritedClassMethod],
+                       @"inheritedClassMethod",
+                       @"Inherited class method was swizzled in base class");
+}
+
+- (void)testSwizzlingUnimplementedClassMethodFails
+{
+  NSError *error = nil;
+  STAssertFalse([NSObject swizzleClassMethod:@selector(inheritedClassMethod)
+                                  withMethod:@selector(version)
+                                       error:nil],
+                @"Swizzling unimplemented class method succeeded");
+  STAssertFalse([NSObject swizzleClassMethod:@selector(inheritedClassMethod)
+                                  withMethod:@selector(version)
+                                       error:&error],
+                @"Swizzling unimplemented class method succeeded");
+  STAssertNotNil(error,
+                 @"Swizzling unimplemented class method did not return error");
+
+  error = nil;
+  STAssertFalse([NSObject swizzleClassMethod:@selector(version)
+                                  withMethod:@selector(inheritedClassMethod)
+                                       error:nil],
+                @"Swizzling with unimplemented class method succeeded");
+  STAssertFalse([NSObject swizzleClassMethod:@selector(version)
+                                  withMethod:@selector(inheritedClassMethod)
+                                       error:&error],
+                @"Swizzling with unimplemented class method succeeded");
+  STAssertNotNil(error,
+                 @"Swizzling with unimplemented class method did not return error");
+}
+
+- (void)testSwizzlingInstanceMethodOverridenByClass
+{
+  TestClassForSwizzling *testClass = [[[TestClassForSwizzling alloc] init] autorelease];
+  STAssertEqualObjects([testClass overridenInstanceMethod],
+                       @"overridenInstanceMethod",
+                       @"Overriden instance method not setup correctly");
+  NSError *error = nil;
+  STAssertTrue([TestClassForSwizzling swizzleInstanceMethod:@selector(overridenInstanceMethod)
+                                                 withMethod:@selector(swizzledOverridenInstanceMethod)
+                                                      error:&error],
+               @"Failed to swizzle overriden instance method; error %@", error);
+  STAssertEqualObjects([testClass overridenInstanceMethod],
+                       @"swizzledOverridenInstanceMethod",
+                       @"Overriden instance method returned wrong value after swizzling");
+  STAssertEqualObjects([testClass swizzledOverridenInstanceMethod],
+                       @"overridenInstanceMethod",
+                       @"Overriden instance method returned wrong value after swizzling");
+
+  TestClassForSwizzlingBase *testBaseClass = [[[TestClassForSwizzlingBase alloc] init] autorelease];
+  STAssertNil([testBaseClass overridenInstanceMethod],
+              @"Overriden instance method was swizzled in base class");
+}
+
+- (void)testSwizzlingInstanceMethodInheritedByClass
+{
+  TestClassForSwizzling *testClass = [[[TestClassForSwizzling alloc] init] autorelease];
+  STAssertEqualObjects([testClass inheritedInstanceMethod],
+                       @"inheritedInstanceMethod",
+                       @"Inherited instance method not setup correctly");
+  NSError *error = nil;
+  STAssertTrue([TestClassForSwizzling swizzleInstanceMethod:@selector(inheritedInstanceMethod)
+                                                 withMethod:@selector(swizzledInheritedInstanceMethod)
+                                                      error:&error],
+               @"Failed to swizzle overriden instance method; error %@", error);
+  STAssertEqualObjects([testClass inheritedInstanceMethod],
+                       @"swizzledInheritedInstanceMethod",
+                       @"Overriden instance method returned wrong value after swizzling");
+  STAssertEqualObjects([testClass swizzledInheritedInstanceMethod],
+                       @"inheritedInstanceMethod",
+                       @"Overriden instance method returned wrong value after swizzling");
+
+  TestClassForSwizzlingBase *testBaseClass = [[[TestClassForSwizzlingBase alloc] init] autorelease];
+  STAssertEqualObjects([testBaseClass inheritedInstanceMethod],
+                       @"inheritedInstanceMethod",
+                       @"Inherited instance method was swizzled in base class");
+}
+
+- (void)testSwizzlingUnimplementedInstanceMethodFails
+{
+  NSError *error = nil;
+  STAssertFalse([NSObject swizzleInstanceMethod:@selector(inheritedInstanceMethod)
+                                     withMethod:@selector(description)
+                                          error:nil],
+                @"Swizzling unimplemented instance method succeeded");
+  STAssertFalse([NSObject swizzleInstanceMethod:@selector(inheritedInstanceMethod)
+                                     withMethod:@selector(description)
+                                          error:&error],
+                @"Swizzling unimplemented instance method succeeded");
+  STAssertNotNil(error,
+                 @"Swizzling unimplemented instance method did not return error");
+
+  error = nil;
+  STAssertFalse([NSObject swizzleClassMethod:@selector(description)
+                                  withMethod:@selector(inheritedInstanceMethod)
+                                       error:nil],
+                @"Swizzling with unimplemented instance method succeeded");
+  STAssertFalse([NSObject swizzleClassMethod:@selector(description)
+                                  withMethod:@selector(inheritedInstanceMethod)
+                                       error:&error],
+                @"Swizzling with unimplemented instance method succeeded");
+  STAssertNotNil(error,
+                 @"Swizzling with unimplemented instance method did not return error");
+}
+
 - (void)testKPT_STRING_FOR_KEY
 {
-  STAssertEquals(KPT_STRING_FOR_KEY(retainCount, NSObject), @"retainCount",
+  STAssertEquals(KPT_STRING_FOR_KEY(retainCount, NSObject),
+                 @"retainCount",
                  @"retainCount key not returned");
 
   // Cannot test failure since this is compile time check.
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.