UIViewController+Additions.h

//
//  UIViewController+Additions.h
//  SFSCategoryKit
//
//  Created by Dalton Claybrook on 4/6/15.
//  Copyright (c) 2015 Dalton Claybrook. All rights reserved.
//

#import <UIKit/UIKit.h>

NS_ASSUME_NONNULL_BEGIN

@interface UIViewController (Additions)

- (void)performOnViewLoad:(void(^)())block
NS_SWIFT_NAME(performOnViewLoad(block:));

@end

NS_ASSUME_NONNULL_END

UIViewController+Additions.m

//
//  UIViewController+Additions.m
//  SFSCategoryKit
//
//  Created by Dalton Claybrook on 4/6/15.
//  Copyright (c) 2015 Dalton Claybrook, LLC. All rights reserved.
//

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

static const void * AdditionsViewDidLoadKey = &AdditionsViewDidLoadKey;

@implementation UIViewController (SFSAdditions)

+ (void)load
{
    method_exchangeImplementations(class_getInstanceMethod(self, @selector(viewDidLoad)), class_getInstanceMethod(self, @selector(dlc_viewDidLoad)));
}

#pragma mark - Public

- (void)performOnViewLoad:(void (^)())block
{
    if (!block)
    {
        NSAssert(NO, @"Must use a block: %@", NSStringFromSelector(_cmd));
        return;
    }
    
    if (self.isViewLoaded)
    {
        block();
    }
    else
    {
        NSMutableArray *array = [self dlc_arrayForKey:AdditionsViewDidLoadKey];
        [array addObject:[block copy]];
    }
}

#pragma mark - Private

- (void)dlc_viewDidLoad
{
    [self dlc_viewDidLoad];
    [self dlc_executeBlocksForKey:AdditionsViewDidLoadKey];
}

- (void)dlc_executeBlocksForKey:(const void *)key
{
    NSArray *blocks = objc_getAssociatedObject(self, key);
    objc_setAssociatedObject(self, key, nil, 0);
    
    for (id blockObj in blocks)
    {
        void (^voidBlock)(void) = blockObj;
        voidBlock();
    }
}

- (NSMutableArray *)dlc_arrayForKey:(const void *)key
{
    NSMutableArray *array = objc_getAssociatedObject(self, key);
    if (!array)
    {
        array = [NSMutableArray array];
        objc_setAssociatedObject(self, key, array, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
    }
    return array;
}

@end