Merge Multiple subclasses hooking

JsonLeecoder · Jun 4, 2015 · d8a43f0 · d8a43f0
2 parents 28485bd + 66fa42f
commit d8a43f0
Show file tree

Hide file tree

Showing 3 changed files with 140 additions and 30 deletions.
diff --git 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
@@ -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<NSCopying>)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,31 +653,69 @@ 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)));
 }
 
 @end
 
 @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
@@ -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<AspectInfo> 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<AspectInfo> 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<AspectInfo> 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<AspectInfo> 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");
 }