diff --git a/.travis.yml b/.travis.yml index 301070b..f827d1b 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,5 +1,5 @@ language: objective-c script: - xcodebuild -project AspectsDemo/AspectsDemo.xcodeproj -scheme AspectsDemo -sdk iphonesimulator ONLY_ACTIVE_ARCH=NO test -- xcodebuild -project AspectsDemo/AspectsDemo.xcodeproj -scheme AspectsDemo -sdk iphonesimulator -destination 'platform=iOS Simulator,name=iPad Retina (64-bit),OS=7.1' test +- xcodebuild -project AspectsDemo/AspectsDemo.xcodeproj -scheme AspectsDemo -sdk iphonesimulator -destination 'platform=iOS Simulator,name=iPad Retina (64-bit),OS=8.1' test - xcodebuild -project AspectsDemoOSX/AspectsDemoOSX.xcodeproj -scheme AspectsDemoOSX test diff --git a/Aspects.m b/Aspects.m index d921198..c907066 100644 --- a/Aspects.m +++ b/Aspects.m @@ -66,10 +66,15 @@ - (BOOL)hasAspects; @end @interface AspectTracker : NSObject -- (id)initWithTrackedClass:(Class)trackedClass parent:(AspectTracker *)parent; +- (id)initWithTrackedClass:(Class)trackedClass; @property (nonatomic, strong) Class trackedClass; +@property (nonatomic, readonly) NSString *trackedClassName; @property (nonatomic, strong) NSMutableSet *selectorNames; -@property (nonatomic, weak) AspectTracker *parentEntry; +@property (nonatomic, strong) NSMutableDictionary *selectorNamesToSubclassTrackers; +- (void)addSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName; +- (void)removeSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName; +- (BOOL)subclassHasHookedSelectorName:(NSString *)selectorName; +- (NSSet *)subclassTrackersHookingSelectorName:(NSString *)selectorName; @end @interface NSInvocation (Aspects) @@ -594,40 +599,50 @@ static BOOL aspect_isSelectorAllowedAndTrack(NSObject *self, SEL selector, Aspec Class klass = [self class]; NSMutableDictionary *swizzledClassesDict = aspect_getSwizzledClassesDict(); Class currentClass = [self class]; + + AspectTracker *tracker = swizzledClassesDict[currentClass]; + if ([tracker subclassHasHookedSelectorName:selectorName]) { + NSSet *subclassTracker = [tracker subclassTrackersHookingSelectorName:selectorName]; + NSSet *subclassNames = [subclassTracker valueForKey:@"trackedClassName"]; + NSString *errorDescription = [NSString stringWithFormat:@"Error: %@ already hooked subclasses: %@. A method can only be hooked once per class hierarchy.", selectorName, subclassNames]; + AspectError(AspectErrorSelectorAlreadyHookedInClassHierarchy, errorDescription); + return NO; + } + do { - AspectTracker *tracker = swizzledClassesDict[currentClass]; + tracker = swizzledClassesDict[currentClass]; if ([tracker.selectorNames containsObject:selectorName]) { - - // Find the topmost class for the log. - if (tracker.parentEntry) { - AspectTracker *topmostEntry = tracker.parentEntry; - while (topmostEntry.parentEntry) { - topmostEntry = topmostEntry.parentEntry; - } - NSString *errorDescription = [NSString stringWithFormat:@"Error: %@ already hooked in %@. A method can only be hooked once per class hierarchy.", selectorName, NSStringFromClass(topmostEntry.trackedClass)]; - AspectError(AspectErrorSelectorAlreadyHookedInClassHierarchy, errorDescription); - return NO; - }else if (klass == currentClass) { + if (klass == currentClass) { // Already modified and topmost! return YES; } + NSString *errorDescription = [NSString stringWithFormat:@"Error: %@ already hooked in %@. A method can only be hooked once per class hierarchy.", selectorName, NSStringFromClass(currentClass)]; + AspectError(AspectErrorSelectorAlreadyHookedInClassHierarchy, errorDescription); + return NO; } - }while ((currentClass = class_getSuperclass(currentClass))); + } while ((currentClass = class_getSuperclass(currentClass))); // Add the selector as being modified. currentClass = klass; - AspectTracker *parentTracker = nil; + AspectTracker *subclassTracker = nil; do { - AspectTracker *tracker = swizzledClassesDict[currentClass]; + tracker = swizzledClassesDict[currentClass]; if (!tracker) { - tracker = [[AspectTracker alloc] initWithTrackedClass:currentClass parent:parentTracker]; + tracker = [[AspectTracker alloc] initWithTrackedClass:currentClass]; swizzledClassesDict[(id)currentClass] = tracker; } - [tracker.selectorNames addObject:selectorName]; + if (subclassTracker) { + [tracker addSubclassTracker:subclassTracker hookingSelectorName:selectorName]; + } else { + [tracker.selectorNames addObject:selectorName]; + } + // All superclasses get marked as having a subclass that is modified. - parentTracker = tracker; + subclassTracker = tracker; }while ((currentClass = class_getSuperclass(currentClass))); - } + } else { + return YES; + } return YES; } @@ -638,14 +653,18 @@ static void aspect_deregisterTrackedSelector(id self, SEL selector) { NSMutableDictionary *swizzledClassesDict = aspect_getSwizzledClassesDict(); NSString *selectorName = NSStringFromSelector(selector); Class currentClass = [self class]; + AspectTracker *subclassTracker = nil; do { AspectTracker *tracker = swizzledClassesDict[currentClass]; - if (tracker) { + if (subclassTracker) { + [tracker removeSubclassTracker:subclassTracker hookingSelectorName:selectorName]; + } else { [tracker.selectorNames removeObject:selectorName]; - if (tracker.selectorNames.count == 0) { - [swizzledClassesDict removeObjectForKey:tracker]; - } } + if (tracker.selectorNames.count == 0 && tracker.selectorNamesToSubclassTrackers) { + [swizzledClassesDict removeObjectForKey:currentClass]; + } + subclassTracker = tracker; }while ((currentClass = class_getSuperclass(currentClass))); } @@ -653,16 +672,50 @@ static void aspect_deregisterTrackedSelector(id self, SEL selector) { @implementation AspectTracker -- (id)initWithTrackedClass:(Class)trackedClass parent:(AspectTracker *)parent { +- (id)initWithTrackedClass:(Class)trackedClass { if (self = [super init]) { _trackedClass = trackedClass; - _parentEntry = parent; _selectorNames = [NSMutableSet new]; + _selectorNamesToSubclassTrackers = [NSMutableDictionary new]; } return self; } + +- (BOOL)subclassHasHookedSelectorName:(NSString *)selectorName { + return self.selectorNamesToSubclassTrackers[selectorName] != nil; +} + +- (void)addSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName { + NSMutableSet *trackerSet = self.selectorNamesToSubclassTrackers[selectorName]; + if (!trackerSet) { + trackerSet = [NSMutableSet new]; + self.selectorNamesToSubclassTrackers[selectorName] = trackerSet; + } + [trackerSet addObject:subclassTracker]; +} +- (void)removeSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName { + NSMutableSet *trackerSet = self.selectorNamesToSubclassTrackers[selectorName]; + [trackerSet removeObject:subclassTracker]; + if (trackerSet.count == 0) { + [self.selectorNamesToSubclassTrackers removeObjectForKey:selectorName]; + } +} +- (NSSet *)subclassTrackersHookingSelectorName:(NSString *)selectorName { + NSMutableSet *hookingSubclassTrackers = [NSMutableSet new]; + for (AspectTracker *tracker in self.selectorNamesToSubclassTrackers[selectorName]) { + if ([tracker.selectorNames containsObject:selectorName]) { + [hookingSubclassTrackers addObject:tracker]; + } + [hookingSubclassTrackers unionSet:[tracker subclassTrackersHookingSelectorName:selectorName]]; + } + return hookingSubclassTrackers; +} +- (NSString *)trackedClassName { + return NSStringFromClass(self.trackedClass); +} + - (NSString *)description { - return [NSString stringWithFormat:@"<%@: %@, trackedClass: %@, selectorNames:%@, parent:%p>", self.class, self, NSStringFromClass(self.trackedClass), self.selectorNames, self.parentEntry]; + return [NSString stringWithFormat:@"<%@: %@, trackedClass: %@, selectorNames:%@, subclass selector names: %@>", self.class, self, NSStringFromClass(self.trackedClass), self.selectorNames, self.selectorNamesToSubclassTrackers.allKeys]; } @end diff --git a/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m b/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m index 72b1276..2bc32fc 100644 --- a/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m +++ b/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m @@ -629,6 +629,9 @@ @implementation A - (void)foo { NSLog(@"%s", __PRETTY_FUNCTION__); } +- (void)bar { + NSLog(@"%s", __PRETTY_FUNCTION__); +} @end @interface B : A @end @@ -638,6 +641,20 @@ - (void)foo { NSLog(@"%s", __PRETTY_FUNCTION__); [super foo]; } +- (void)bar { + NSLog(@"%s", __PRETTY_FUNCTION__); + [super bar]; +} +@end + +@interface C : NSObject +- (void)foo; +@end + +@implementation C +- (void)foo { + NSLog(@"%s", __PRETTY_FUNCTION__); +} @end @interface AspectsSelectorTests : XCTestCase @end @@ -666,6 +683,8 @@ @implementation AspectsSelectorTests - (void)testSelectorMangling2 { __block BOOL A_aspect_called = NO; __block BOOL B_aspect_called = NO; + __block BOOL C_aspect_called = NO; + id aspectToken1 = [A aspect_hookSelector:@selector(foo) withOptions:AspectPositionBefore usingBlock:^(id info) { NSLog(@"before -[A foo]"); A_aspect_called = YES; @@ -678,13 +697,51 @@ - (void)testSelectorMangling2 { } error:NULL]; XCTAssertNil(aspectToken2, @"Must not return a token"); + // a sibling and it's subclasses should be able to hook the same selector + id aspectToken3 = [C aspect_hookSelector:@selector(foo) withOptions:AspectPositionBefore usingBlock:^(id info) { + NSLog(@"before -[C foo]"); + C_aspect_called = YES; + } error:NULL]; + XCTAssertNotNil(aspectToken3, @"Must return a token"); + B *b = [B new]; [b foo]; // TODO: A is not yet called, we can't detect the target IMP for an invocation. XCTAssertTrue(A_aspect_called, @"A aspect should be called"); XCTAssertFalse(B_aspect_called, @"B aspect should not be called"); - + XCTAssertFalse(C_aspect_called, @"C aspect should not be called"); + + C *c = [C new]; + [c foo]; + XCTAssertTrue(C_aspect_called, @"C aspect should be called"); + + XCTAssertTrue([aspectToken1 remove], @"Must be able to deregister"); + XCTAssertTrue([aspectToken3 remove], @"Must be able to deregister"); +} +- (void)testSelectorMangling3 { + __block BOOL A_aspect_called = NO; + __block BOOL B_aspect_called = NO; + + id aspectToken1 = [B aspect_hookSelector:@selector(bar) withOptions:AspectPositionBefore usingBlock:^(id info) { + NSLog(@"before -[B bar]"); + B_aspect_called = YES; + } error:NULL]; + XCTAssertNotNil(aspectToken1, @"Must return a token"); + + // if a subclass already hooks this selector we shouldn't be able to hook it in a superclass + id aspectToken2 = [A aspect_hookSelector:@selector(bar) withOptions:AspectPositionBefore usingBlock:^(id info) { + NSLog(@"before -[A bar]"); + A_aspect_called = YES; + } error:NULL]; + XCTAssertNil(aspectToken2, @"Must not return a token"); + + B *b = [B new]; + [b bar]; + + XCTAssertFalse(A_aspect_called, @"A aspect should not be called"); + XCTAssertTrue(B_aspect_called, @"B aspect should be called"); + XCTAssertTrue([aspectToken1 remove], @"Must be able to deregister"); }