objc_setAssociatedObject in a category sets for all subclasses

Question

I have a custom container view controller that manages the view hierarchy of my app. I know that every controller is some sort of child of this container controller. I thought it would be nice to have a category on UIViewController that would allow me to access the container controller, no matter where I am in the hierarchy.

This involves a recursive walk up the controller hierarchy, so I thought it would be good to try and only do that walk once per controller. So with objc_setAssociatedObject, I set the container once I've found it and set a flag so that I know whether or not I need to walk the hierarchy on subsequent calls (I planned to invalidate that if the viewcontroller ever moved, but that's probably overkill, and I didn't get that far).

Anyway, that works fine for the most part except that my flag for whether or not the hierarchy has been walked seems to be attached to UIViewController, and not specific subclasses of UIViewController.

I swizzled +load to try to set default values on my associated objects to no avail.

Any ideas? How to I get associated objects in a category to associate with the subclasses of the class the category is defined on?

Here's my code, for good measure.

#import "UIViewController+LMPullMenuContainer.h"
#import <objc/runtime.h>

static char const * const CachedKey = "__LM__CachedBoolPullMenuAssociatedObjectKey";
static char const * const PullMenuKey = "__LM__PullMenuAssociatedObjectKey";

@implementation UIViewController (LMPullMenuContainer)

+ (void)load {
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        SEL initSelector = @selector(initWithCoder:);
        SEL pullViewInitSelector =  @selector(init__LM__Swizzled__WithCoder:);
        Method originalMethod = class_getInstanceMethod(self, initSelector);
        Method newMethod = class_getInstanceMethod(self, pullViewInitSelector);

        BOOL methodAdded = class_addMethod([self class],
                                           initSelector,
                                           method_getImplementation(newMethod),
                                           method_getTypeEncoding(newMethod));

        if (methodAdded) {
            class_replaceMethod([self class],
                                pullViewInitSelector,
                                method_getImplementation(originalMethod),
                                method_getTypeEncoding(originalMethod));
        } else {
            method_exchangeImplementations(originalMethod, newMethod);
        }
    });
}

- (instancetype)init__LM__Swizzled__WithCoder:(NSCoder *)coder {
    self = [self init__LM__Swizzled__WithCoder:coder];
    if (self != nil)
    {
        objc_setAssociatedObject(self, CachedKey, [NSNumber numberWithBool:NO], OBJC_ASSOCIATION_RETAIN_NONATOMIC);
        objc_setAssociatedObject(self, PullMenuKey, nil, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
    }
    return self;
}

- (LMPullMenuContainerViewController*)pullMenuContainerController {
    BOOL isCached = [objc_getAssociatedObject(self, CachedKey) boolValue];
    if (isCached) {
        return objc_getAssociatedObject(self, PullMenuKey);
    } else {
        return [self pullMenuParentOf:self];
    }
}

- (LMPullMenuContainerViewController *)pullMenuParentOf:(UIViewController *)controller {
    if (controller.parentViewController) {
        if ([controller.parentViewController isKindOfClass:[LMPullMenuContainerViewController class]]) {
            objc_setAssociatedObject(self, CachedKey, [NSNumber numberWithBool:YES], OBJC_ASSOCIATION_RETAIN_NONATOMIC);
            objc_setAssociatedObject(self, PullMenuKey, controller.parentViewController, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
            return (LMPullMenuContainerViewController *)(controller.parentViewController);
        } else {
            return [self pullMenuParentOf:controller.parentViewController];
        }
    } else {
        objc_setAssociatedObject(self, CachedKey, [NSNumber numberWithBool:YES], OBJC_ASSOCIATION_RETAIN_NONATOMIC);
        objc_setAssociatedObject(self, PullMenuKey, nil, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
        return nil;
    }
}

For now I've resigned to setting the property manually where necessary.

Answer 1

As it happens, the above code works just fine. My container controller was loading all of the controllers it manages when it was first initialized and not when the controllers were first displayed, so to me it looked as though the flag had been set before it should have been.