Aspects 1.3.0. Optional, automatic hook de-registration.

Aspects.h

@@ -7,10 +7,12 @@
 
 #import <Foundation/Foundation.h>
 
-typedef NS_ENUM(NSUInteger, AspectPosition) {
-    AspectPositionBefore,  /// Called before the original implementation.
-    AspectPositionInstead, /// Will replace the original implementation.
-    AspectPositionAfter    /// Called after the original implementation.
+typedef NS_OPTIONS(NSUInteger, AspectOptions) {
+    AspectPositionAfter   = 0,            /// Called after the original implementation (default)
+    AspectPositionInstead = 1,            /// Will replace the original implementation.
+    AspectPositionBefore  = 2,            /// Called before the original implementation.
+
+    AspectOptionAutomaticRemoval = 1 << 3 /// Will remove the hook after the first execution.
 };
 
 /// Opaque Aspect Token that allows to deregister the hook.
@@ -30,26 +32,28 @@ typedef NS_ENUM(NSUInteger, AspectPosition) {
  */
 @interface NSObject (Aspects)
 
-/// Adds a block of code before/instead/after the current `selector` for a specific object.
+/// Adds a block of code before/instead/after the current `selector` for a specific class.
 /// If you choose `AspectPositionInstead`, the `arguments` array will contain the original invocation as last argument.
+/// @note Hooking static methods is not supported.
 /// @return A token which allows to later deregister the aspect.
-- (id<Aspect>)aspect_hookSelector:(SEL)selector
-                       atPosition:(AspectPosition)position
-                        withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block
++ (id<Aspect>)aspect_hookSelector:(SEL)selector
+                      withOptions:(AspectOptions)options
+                       usingBlock:(void (^)(id instance, NSArray *args))block
                             error:(NSError **)error;
 
-/// Hooks a selector class-wide.
-+ (id<Aspect>)aspect_hookSelector:(SEL)selector
-                       atPosition:(AspectPosition)position
-                        withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block
+/// Adds a block of code before/instead/after the current `selector` for a specific instance.
+- (id<Aspect>)aspect_hookSelector:(SEL)selector
+                      withOptions:(AspectOptions)options
+                       usingBlock:(void (^)(id instance, NSArray *args))block
                             error:(NSError **)error;
 
 @end
 
 
 typedef NS_ENUM(NSUInteger, AspectsErrorCode) {
     AspectsErrorSelectorBlacklisted,                   /// Selectors like release, retain, autorelease are blacklisted.
-    AspectsErrorSelectorDeallocPosition,               /// When hooking dealloc, AspectPositionInstead is not allowed.
+    AspectsErrorDoesNotRespondToSelector,              /// Selector could not be found.
+    AspectsErrorSelectorDeallocPosition,               /// When hooking dealloc, only AspectPositionBefore is allowed.
     AspectsErrorSelectorAlreadyHookedInClassHierarchy, /// Statically hooking the same method in subclasses is not allowed.
     AspectsErrorFailedToAllocateClassPair,             /// The runtime failed creating a class pair.
     AspectsErrorRemoveObjectAlreadyDeallocated = 100   /// (for removing) The object hooked is already deallocated.

Aspects.m

@@ -12,18 +12,20 @@
 
 #define AspectLog(...)
 //#define AspectLog(...) do { NSLog(__VA_ARGS__); }while(0)
+#define AspectLogError(...) do { NSLog(__VA_ARGS__); }while(0)
 
 // Tracks a single aspect.
 @interface AspectIdentifier : NSObject
-- (id)initWithSelector:(SEL)selector object:(id)object block:(id)block;
+- (id)initWithSelector:(SEL)selector object:(id)object options:(AspectOptions)options block:(id)block;
 @property (nonatomic, assign) SEL selector;
 @property (nonatomic, strong) id block;
 @property (nonatomic, weak) id object;
+@property (nonatomic, assign) AspectOptions options;
 @end
 
 // Tracks all aspects for an object/class.
 @interface AspectsContainer : NSObject
-- (void)addAspect:(AspectIdentifier *)aspect atPosition:(AspectPosition)injectPosition;
+- (void)addAspect:(AspectIdentifier *)aspect withOptions:(AspectOptions)injectPosition;
 - (BOOL)removeAspect:(id)aspect;
 - (BOOL)hasAspects;
 @property (atomic, copy) NSArray *beforeAspects;
@@ -42,8 +44,12 @@ @interface NSInvocation (Aspects)
 - (NSArray *)aspects_arguments;
 @end
 
+typedef void(^AspectBlock)(id instance, NSArray *arguments);
+
+#define AspectPositionFilter 0x07
+
 #define AspectError(errorCode, errorDescription) do { \
-AspectLog(@"Aspects: %@", errorDescription); \
+AspectLogError(@"Aspects: %@", errorDescription); \
 if (error) { *error = [NSError errorWithDomain:AspectsErrorDomain code:errorCode userInfo:@{NSLocalizedDescriptionKey: errorDescription}]; }}while(0)
 
 NSString *const AspectsErrorDomain = @"AspectsErrorDomain";
@@ -55,34 +61,35 @@ @implementation NSObject (Aspects)
 ///////////////////////////////////////////////////////////////////////////////////////////
 #pragma mark - Public Aspects API
 
-+ (id)aspect_hookSelector:(SEL)selector
-               atPosition:(AspectPosition)position
-                withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block
-                    error:(NSError **)error {
-    return aspect_add((id)self, selector, position, block, error);
++ (id<Aspect>)aspect_hookSelector:(SEL)selector
+                      withOptions:(AspectOptions)options
+                       usingBlock:(AspectBlock)block
+                            error:(NSError **)error {
+    return aspect_add((id)self, selector, options, block, error);
 }
 
-- (id)aspect_hookSelector:(SEL)selector
-               atPosition:(AspectPosition)position
-                withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block
-                    error:(NSError **)error {
-    return aspect_add(self, selector, position, block, error);
+/// @return A token which allows to later deregister the aspect.
+- (id<Aspect>)aspect_hookSelector:(SEL)selector
+                      withOptions:(AspectOptions)options
+                       usingBlock:(AspectBlock)block
+                            error:(NSError **)error {
+    return aspect_add(self, selector, options, block, error);
 }
 
 ///////////////////////////////////////////////////////////////////////////////////////////
 #pragma mark - Private Helper
 
-static id aspect_add(id self, SEL selector, AspectPosition position, void (^block)(__unsafe_unretained id object, NSArray *arguments), NSError **error) {
+static id aspect_add(id self, SEL selector, AspectOptions options, AspectBlock block, NSError **error) {
     NSCParameterAssert(self);
     NSCParameterAssert(selector);
     NSCParameterAssert(block);
 
     __block AspectIdentifier *identifier = nil;
     aspect_performLocked(^{
-        if (aspect_isSelectorAllowedAndTrack(self, selector, position, error)) {
+        if (aspect_isSelectorAllowedAndTrack(self, selector, options, error)) {
             AspectsContainer *aspectContainer = aspect_getContainerForObject(self, selector);
-            identifier = [[AspectIdentifier alloc] initWithSelector:selector object:self block:block];;
-            [aspectContainer addAspect:identifier atPosition:position];
+            identifier = [[AspectIdentifier alloc] initWithSelector:selector object:self options:options block:block];
+            [aspectContainer addAspect:identifier withOptions:options];
 
             // Modify the class to allow message interception.
             aspect_prepareClassAndHookSelector(self, selector, error);
@@ -337,7 +344,13 @@ static void aspect_undoSwizzleClassInPlace(Class klass) {
 #pragma mark - Aspect Invoke Point
 
 // This is a macro so we get a cleaner stack trace.
-#define aspect_invoke(aspects, arguments) for (AspectIdentifier *aspect in aspects) {((void (^)(id, NSArray *))aspect.block)(self, arguments); }
+#define aspect_invoke(aspects, arguments) \
+for (AspectIdentifier *aspect in aspects) {\
+    ((void (^)(id, NSArray *))aspect.block)(self, arguments);\
+    if (aspect.options & AspectOptionAutomaticRemoval) { \
+        aspectsToRemove = [aspectsToRemove?:@[] arrayByAddingObject:aspect]; \
+    } \
+}
 
 // This is the swizzled forwardInvocation: method.
 static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL selector, NSInvocation *invocation) {
@@ -346,6 +359,7 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
 	SEL aliasSelector = aspect_aliasForSelector(invocation.selector);
     AspectsContainer *objectContainer = objc_getAssociatedObject(self, aliasSelector);
     AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector);
+    NSArray *aspectsToRemove = nil;
 
     // Before hooks.
     NSArray *arguments = nil;
@@ -380,16 +394,16 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
 
     // If no hooks are installed, call original implementation (usually to throw an exception)
     if (!respondsToAlias) {
-        SEL aspectsForwardInvocationSEL = NSSelectorFromString(AspectsForwardInvocationSelectorName);
-        if ([self respondsToSelector:aspectsForwardInvocationSEL]) {
-#pragma clang diagnostic push
-#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
-            [self performSelector:aspectsForwardInvocationSEL withObject:invocation];
-#pragma clang diagnostic pop
+        SEL originalForwardInvocationSEL = NSSelectorFromString(AspectsForwardInvocationSelectorName);
+        if ([self respondsToSelector:originalForwardInvocationSEL]) {
+            ((void( *)(id, SEL, NSInvocation *))objc_msgSend)(self, originalForwardInvocationSEL, invocation);
         }else {
             [self doesNotRecognizeSelector:invocation.selector];
         }
     }
+
+    // Remove any hooks that are queued for deregistration.
+    [aspectsToRemove makeObjectsPerformSelector:@selector(remove)];
 }
 #undef aspect_invoke
 
@@ -437,7 +451,7 @@ static void aspect_destroyContainerForObject(id<NSObject> self, SEL selector) {
     return swizzledClassesDict;
 }
 
-static BOOL aspect_isSelectorAllowedAndTrack(id self, SEL selector, AspectPosition position, NSError **error) {
+static BOOL aspect_isSelectorAllowedAndTrack(NSObject *self, SEL selector, AspectOptions options, NSError **error) {
     static NSSet *disallowedSelectorList;
     static dispatch_once_t pred;
     dispatch_once(&pred, ^{
@@ -447,15 +461,22 @@ static BOOL aspect_isSelectorAllowedAndTrack(id self, SEL selector, AspectPositi
     // Check against the blacklist.
     NSString *selectorName = NSStringFromSelector(selector);
     if ([disallowedSelectorList containsObject:selectorName]) {
-        NSString *errorDescription = [NSString stringWithFormat:@"Selector `%@` is blacklisted.", selectorName];
+        NSString *errorDescription = [NSString stringWithFormat:@"Selector %@ is blacklisted.", selectorName];
         AspectError(AspectsErrorSelectorBlacklisted, errorDescription);
         return NO;
     }
 
     // Additional checks.
-    if ([selectorName isEqualToString:@"dealloc"] && position == AspectPositionInstead) {
-        NSString *errorDescription = @"dealloc can not be replaced. Use AspectPositionBefore.";
-        AspectError(AspectsErrorSelectorDeallocPosition, errorDescription);
+    AspectOptions position = options&AspectPositionFilter;
+    if ([selectorName isEqualToString:@"dealloc"] && position != AspectPositionBefore) {
+        NSString *errorDesc = @"AspectPositionBefore is the only valid position when hooking dealloc.";
+        AspectError(AspectsErrorSelectorDeallocPosition, errorDesc);
+        return NO;
+    }
+
+    if (![self respondsToSelector:selector] && ![self.class instancesRespondToSelector:selector]) {
+        NSString *errorDesc = [NSString stringWithFormat:@"Unable to find selector -[%@ %@].", NSStringFromClass(self.class), selectorName];
+        AspectError(AspectsErrorDoesNotRespondToSelector, errorDesc);
         return NO;
     }
 
@@ -624,19 +645,20 @@ - (NSArray *)aspects_arguments {
 
 @implementation AspectIdentifier
 
-- (id)initWithSelector:(SEL)selector object:(id)object block:(id)block {
+- (id)initWithSelector:(SEL)selector object:(id)object options:(AspectOptions)options block:(id)block {
     NSCParameterAssert(block);
     NSCParameterAssert(selector);
     if (self = [super init]) {
         _selector = selector;
         _block = block;
+        _options = options;
         _object = object; // weak
     }
     return self;
 }
 
 - (NSString *)description {
-    return [NSString stringWithFormat:@"<%@: %p, SEL:%@ object:%@ block:%@>", self.class, self, NSStringFromSelector(self.selector), self.object, self.block];
+    return [NSString stringWithFormat:@"<%@: %p, SEL:%@ object:%@ options:%tu block:%@>", self.class, self, NSStringFromSelector(self.selector), self.object, self.options, self.block];
 }
 
 - (BOOL)remove {
@@ -654,8 +676,9 @@ - (BOOL)hasAspects {
     return self.beforeAspects.count > 0 || self.insteadAspects.count > 0 || self.afterAspects.count > 0;
 }
 
-- (void)addAspect:(AspectIdentifier *)aspect atPosition:(AspectPosition)injectPosition {
-    switch (injectPosition) {
+- (void)addAspect:(AspectIdentifier *)aspect withOptions:(AspectOptions)options {
+    NSUInteger position = options&AspectPositionFilter;
+    switch (position) {
         case AspectPositionBefore:  self.beforeAspects  = [(self.beforeAspects ?:@[]) arrayByAddingObject:aspect]; break;
         case AspectPositionInstead: self.insteadAspects = [(self.insteadAspects?:@[]) arrayByAddingObject:aspect]; break;
         case AspectPositionAfter:   self.afterAspects   = [(self.afterAspects  ?:@[]) arrayByAddingObject:aspect]; break;

Aspects.podspec

@@ -1,6 +1,6 @@
 Pod::Spec.new do |s|
   s.name         = "Aspects"
-  s.version      = "1.2.0"
+  s.version      = "1.3.0"
   s.summary      = "Delightful, simple library for aspect oriented programming."
   s.homepage     = "https://github.com/steipete/Aspects"
   s.license      = { :type => 'MIT', :file => 'LICENSE' }

AspectsDemo/AspectsDemo.xcodeproj/project.pbxproj

@@ -20,6 +20,8 @@
 		78573F1819155A2E000D3B00 /* InfoPlist.strings in Resources */ = {isa = PBXBuildFile; fileRef = 78573F1619155A2E000D3B00 /* InfoPlist.strings */; };
 		78573F1A19155A2E000D3B00 /* AspectsDemoTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 78573F1919155A2E000D3B00 /* AspectsDemoTests.m */; };
 		78573F2519155A74000D3B00 /* Aspects.m in Sources */ = {isa = PBXBuildFile; fileRef = 78573F2319155A74000D3B00 /* Aspects.m */; };
+		78D7D77119177C8E002EB314 /* AspectsViewController.m in Sources */ = {isa = PBXBuildFile; fileRef = 78D7D76F19177C8E002EB314 /* AspectsViewController.m */; };
+		78D7D77219177C8E002EB314 /* AspectsViewController.xib in Resources */ = {isa = PBXBuildFile; fileRef = 78D7D77019177C8E002EB314 /* AspectsViewController.xib */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXContainerItemProxy section */
@@ -51,6 +53,9 @@
 		78573F1919155A2E000D3B00 /* AspectsDemoTests.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = AspectsDemoTests.m; sourceTree = "<group>"; };
 		78573F2319155A74000D3B00 /* Aspects.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; name = Aspects.m; path = ../../Aspects.m; sourceTree = "<group>"; };
 		78573F2419155A74000D3B00 /* Aspects.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = Aspects.h; path = ../../Aspects.h; sourceTree = "<group>"; };
+		78D7D76E19177C8E002EB314 /* AspectsViewController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = AspectsViewController.h; sourceTree = "<group>"; };
+		78D7D76F19177C8E002EB314 /* AspectsViewController.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = AspectsViewController.m; sourceTree = "<group>"; };
+		78D7D77019177C8E002EB314 /* AspectsViewController.xib */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = file.xib; path = AspectsViewController.xib; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -116,6 +121,9 @@
 				78573F0419155A2E000D3B00 /* AspectsAppDelegate.m */,
 				78573F0619155A2E000D3B00 /* Images.xcassets */,
 				78573EFB19155A2E000D3B00 /* Supporting Files */,
+				78D7D76E19177C8E002EB314 /* AspectsViewController.h */,
+				78D7D76F19177C8E002EB314 /* AspectsViewController.m */,
+				78D7D77019177C8E002EB314 /* AspectsViewController.xib */,
 			);
 			path = AspectsDemo;
 			sourceTree = "<group>";
@@ -227,6 +235,7 @@
 			files = (
 				78573EFF19155A2E000D3B00 /* InfoPlist.strings in Resources */,
 				78573F0719155A2E000D3B00 /* Images.xcassets in Resources */,
+				78D7D77219177C8E002EB314 /* AspectsViewController.xib in Resources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -245,6 +254,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				78D7D77119177C8E002EB314 /* AspectsViewController.m in Sources */,
 				78573F0519155A2E000D3B00 /* AspectsAppDelegate.m in Sources */,
 				78573F2519155A74000D3B00 /* Aspects.m in Sources */,
 				78573F0119155A2E000D3B00 /* main.m in Sources */,

AspectsDemo/AspectsDemo/AspectsAppDelegate.h

@@ -6,7 +6,7 @@
 //  Copyright (c) 2014 PSPDFKit GmbH. All rights reserved.
 //
 
-#import <UIKit/UIKit.h>
+@import UIKit;
 
 @interface AspectsAppDelegate : UIResponder <UIApplicationDelegate>
 

AspectsDemo/AspectsDemo/AspectsAppDelegate.m

@@ -7,20 +7,29 @@
 //
 
 #import "AspectsAppDelegate.h"
+#import "AspectsViewController.h"
 #import "Aspects.h"
 
 @implementation AspectsAppDelegate
 
 - (BOOL)application:(UIApplication *)application didFinishLaunchingWithOptions:(NSDictionary *)launchOptions {
+    AspectsViewController *aspectsController = [AspectsViewController new];
     self.window = [[UIWindow alloc] initWithFrame:[[UIScreen mainScreen] bounds]];
-
-//    [self.window aspect_hookSelector:@selector(makeKeyAndVisible) atPosition:AspectPositionBefore withBlock:^(id object, NSArray *arguments) {
-//        NSLog(@"We're about to call -[UIWindow makeKeyAndVisible].");
-//    }];
-
-    // Override point for customization after application launch.
     self.window.backgroundColor = [UIColor whiteColor];
+    self.window.rootViewController = [[UINavigationController alloc] initWithRootViewController:aspectsController];
     [self.window makeKeyAndVisible];
+
+    // Ignore hooks when we are testing.
+    if (!NSClassFromString(@"XCTestCase")) {
+        [aspectsController aspect_hookSelector:@selector(buttonPressed:) withOptions:0 usingBlock:^(id instance, NSArray *arguments) {
+            NSLog(@"Button was pressed by: %@", arguments.firstObject);
+        } error:NULL];
+
+        [aspectsController aspect_hookSelector:@selector(viewWillLayoutSubviews) withOptions:0 usingBlock:^(id instance, NSArray *arguments) {
+            NSLog(@"Controller is layouting!");
+        } error:NULL];
+    }
+
     return YES;
 }
 

AspectsDemo/AspectsDemo/AspectsViewController.h

@@ -0,0 +1,15 @@
+//
+//  AspectsViewController.h
+//  AspectsDemo
+//
+//  Created by Peter Steinberger on 05/05/14.
+//  Copyright (c) 2014 PSPDFKit GmbH. All rights reserved.
+//
+
+#import <UIKit/UIKit.h>
+
+@interface AspectsViewController : UIViewController
+
+- (IBAction)buttonPressed:(id)sender;
+
+@end