Programmatically calling segues without implementing `prepareForSegue:sender:`

README.md

UIViewController+SegueBlocks

This is a category which adds methods on UIViewController to provide a handler at the spot that you programmatically perform a segue. This avoids a mess of if, else if, else statements in prepareForSegue:sender:. In fact, you don't even need to implement the method at all if you don't need do anything before the segues called from actions in your storyboard.

Usage

NSNumber *blogPostID = @1;
[self performSegueWithIdentifier:@"BlogPost" handler:^(BlogPostVC *blogPostVC) {
    blogPostVC.blogPostID = blogPostID;
}];

Implementation

The category provides two methods which take blocks – performSegueWithIdentifier:handler: and performSegueWithIdentifier:handler:shouldSegue: – which construct an internal container object, SegueHandler, and then call performSegueWithIdentifier:sender:.

The category's load method swizzles shouldPerformSegueWithIdentifier:sender: and prepareForSegue:sender: to provide implementations that use blocks provided via the SegueHandler sender object. The custom code only executes if the sender is an instance of SegueHandler, otherwise it defaults to the original implementation.

Because implementations of prepareForSegue:sender: rarely call [super prepareForSegue:sender:], not only is UIViewController swizzled but also all of its subclasses. This ensures that the custom code will get called.

UIViewController+SegueBlocks.h

//
//  UIViewController+SegueBlocks.h
//  SegueBlocks
//
//  Created by Peter Robinett on 2015-02-04.
//  Copyright (c) 2015 Bubble Foundry. All rights reserved.
//

#import <UIKit/UIKit.h>

@interface UIViewController (SegueBlocks)

- (void)performSegueWithIdentifier:(NSString *)identifier handler:(void (^)(id destinationViewController))handlerBlock;

- (void)performSegueWithIdentifier:(NSString *)identifier handler:(void (^)(id destinationViewController))handlerBlock shouldSegue:(BOOL (^)())shouldSegueBlock;

@end

UIViewController+SegueBlocks.m

//
//  UIViewController+SegueBlocks.m
//  SegueBlocks
//
//  Created by Peter Robinett on 2015-02-04.
//  Copyright (c) 2015 Bubble Foundry. All rights reserved.
//

#import "UIViewController+SegueBlocks.h"

#import <objc/runtime.h>

#pragma mark - SegueHandler

@interface SegueHandler : NSObject

@property (copy) void (^handlerBlock)(UIViewController *destinationViewController);
@property (copy) BOOL (^shouldSegueBlock)();

@end

@implementation SegueHandler

@end

#pragma mark - SegueBlocks

@implementation UIViewController (SegueBlocks)

- (void)performSegueWithIdentifier:(NSString *)identifier handler:(void (^)(id destinationViewController))handlerBlock
{
    [self performSegueWithIdentifier:identifier handler:handlerBlock shouldSegue:nil];
}

- (void)performSegueWithIdentifier:(NSString *)identifier handler:(void (^)(id destinationViewController))handlerBlock shouldSegue:(BOOL (^)())shouldSegueBlock
{
    SegueHandler *handler = [SegueHandler new];
    handler.handlerBlock = handlerBlock;
    handler.shouldSegueBlock = shouldSegueBlock;

    [self performSegueWithIdentifier:identifier sender:handler];
}

#pragma mark - Private

// from http://www.cocoawithlove.com/2010/01/getting-subclasses-of-objective-c-class.html
NSArray *ClassGetSubclasses(Class parentClass)
{
    int numClasses = objc_getClassList(NULL, 0);
    Class *classes = NULL;
 
    // classes = malloc(sizeof(Class) * numClasses);
    // from http://stackoverflow.com/a/8731509/46768
    classes = (__unsafe_unretained Class *)malloc(sizeof(Class) * numClasses);
    numClasses = objc_getClassList(classes, numClasses);
     
    NSMutableArray *result = [NSMutableArray array];
    for (NSInteger i = 0; i < numClasses; i++)
    {
        Class superClass = classes[i];
        do
        {
            superClass = class_getSuperclass(superClass);
        } while(superClass && superClass != parentClass);
         
        if (superClass == nil)
        {
            continue;
        }
         
        [result addObject:classes[i]];
    }
 
    free(classes);
     
    return result;
}

/*
  An interesting feature of +load is that it's special-cased by the runtime to be invoked in categories which implement it as well as the main class. This means that if you implement +load in a class and in a category on that class, both will be called. This probably goes against everything you know about how categories work, but that's because +load is not a normal method. This feature means that +load is an excellent place to do evil things like method swizzling.
  - https://www.mikeash.com/pyblog/friday-qa-2009-05-22-objective-c-class-loading-and-initialization.html
*/
// code from http://nshipster.com/method-swizzling/
+ (void)load {
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        Class class = [self class];
        // add to UIViewController
        [self swizzleClassWithSegueMethods:class];

        // now add to any subclasses, since they normally don't call `[super prepareForSegue:sender:`.
        NSArray *subclasses = ClassGetSubclasses(class);
        for (Class subclass in subclasses) {
            [self swizzleClassWithSegueMethods:subclass];
        }

    });
}

+ (void)swizzleClassWithSegueMethods:(Class)class
{
    // swizzle shouldPerformSegueWithIdentifier:sender:
    SEL originalShouldSelector = @selector(shouldPerformSegueWithIdentifier:sender:);
    SEL swizzledShouldSelector = @selector(bftypedsegues_shouldPerformSegueWithIdentifier:sender:);

    Method originalShouldMethod = class_getInstanceMethod(class, originalShouldSelector);
    Method swizzledShouldMethod = class_getInstanceMethod(class, swizzledShouldSelector);

    BOOL didAddShouldMethod = class_addMethod(class, originalShouldSelector, method_getImplementation(swizzledShouldMethod), method_getTypeEncoding(swizzledShouldMethod));

    if (didAddShouldMethod) {
        class_replaceMethod(class, swizzledShouldSelector, method_getImplementation(originalShouldMethod), method_getTypeEncoding(originalShouldMethod));
    }
    else {
        method_exchangeImplementations(originalShouldMethod, swizzledShouldMethod);
    }

    // swizzle prepareForSegue:sender:
    SEL originalSelector = @selector(prepareForSegue:sender:);
    SEL swizzledSelector = @selector(bftypedsegues_prepareForSegue:sender:);

    Method originalMethod = class_getInstanceMethod(class, originalSelector);
    Method swizzledMethod = class_getInstanceMethod(class, swizzledSelector);

    BOOL didAddMethod = class_addMethod(class, originalSelector, method_getImplementation(swizzledMethod), method_getTypeEncoding(swizzledMethod));

    if (didAddMethod) {
        class_replaceMethod(class, swizzledSelector, method_getImplementation(originalMethod), method_getTypeEncoding(originalMethod));
    }
    else {
        method_exchangeImplementations(originalMethod, swizzledMethod);
    }
}

#pragma mark - Method Swizzling

- (void)bftypedsegues_prepareForSegue:(UIStoryboardSegue *)segue sender:(id)sender
{
    // handle things
    if ([sender isKindOfClass:[SegueHandler class]]) {
        SegueHandler *handler = sender;

        if (handler.handlerBlock != nil) {
            handler.handlerBlock(segue.destinationViewController);
        }
    }
    else {
        // call the original implementation
        [self bftypedsegues_prepareForSegue:segue sender:sender];
    }
}

- (BOOL)bftypedsegues_shouldPerformSegueWithIdentifier:(NSString *)identifier sender:(id)sender
{
    // handle things
    if ([sender isKindOfClass:[SegueHandler class]]) {
        SegueHandler *handler = sender;

        if (handler.shouldSegueBlock != nil) {
            return handler.shouldSegueBlock();
        }
        else {
            return YES;
        }
    }
    else {
        // call the original implementation
        return [self bftypedsegues_shouldPerformSegueWithIdentifier:identifier sender:sender];
    }
}

@end