PSPDFUIKitMainThreadGuard.m

#import "VMUIKitMainThreadGuard.h"

#import <objc/runtime.h>
#import <objc/message.h>
#import <UIKit/UIKit.h>

/**
 * This installs a small guard that checks for the most common threading-errors in UIKit.
 * This won't really slow down performance but still only is compiled in DEBUG versions.
 * @note No private API is used here.
 * @see https://gist.github.com/steipete/5664345
 */
@implementation VMUIKitMainThreadGuard

#ifdef DEBUG

#define PSPDFAssert(expression, ...) \
    do { if (!(expression)) { \
             NSLog(@"%@", [NSString stringWithFormat:@"Assertion failure: %s in %s on line %s:%d. %@", #expression, __PRETTY_FUNCTION__, __FILE__, __LINE__, [NSString stringWithFormat:@"" __VA_ARGS__]]); \
             abort(); } } while (0)

// http://www.mikeash.com/pyblog/friday-qa-2010-01-29-method-replacement-for-fun-and-profit.html
BOOL PSPDFReplaceMethodWithBlock(Class c, SEL origSEL, SEL newSEL, id block)
{
    PSPDFAssert(c && origSEL && newSEL && block);
    if ([c instancesRespondToSelector:newSEL]) {
        return YES; // Selector already implemented, skip silently.

    }
    Method origMethod = class_getInstanceMethod(c, origSEL);

    // Add the new method.
    IMP impl = imp_implementationWithBlock(block);
    if (!class_addMethod(c, newSEL, impl, method_getTypeEncoding(origMethod))) {
        NSLog(@"Failed to add method: %@ on %@", NSStringFromSelector(newSEL), c);
        return NO;
    } else {
        Method newMethod = class_getInstanceMethod(c, newSEL);

        // If original doesn't implement the method we want to swizzle, create it.
        if (class_addMethod(c, origSEL, method_getImplementation(newMethod), method_getTypeEncoding(origMethod))) {
            class_replaceMethod(c, newSEL, method_getImplementation(origMethod), method_getTypeEncoding(newMethod));
        } else {
            method_exchangeImplementations(origMethod, newMethod);
        }
    }
    return YES;
}

SEL _PSPDFPrefixedSelector(SEL selector)
{
    return NSSelectorFromString([NSString stringWithFormat:@"pspdf_%@", NSStringFromSelector(selector)]);
}

void PSPDFAssertIfNotMainThread(void)
{
    PSPDFAssert(NSThread.isMainThread, @"\nERROR: All calls to UIKit need to happen on the main thread. You have a bug in your code. Use dispatch_async(dispatch_get_main_queue(), ^{ ... }); if you're unsure what thread you're in.\n\nBreak on PSPDFAssertIfNotMainThread to find out where.\n\nStacktrace: %@", NSThread.callStackSymbols);
}

__attribute__((constructor)) static void PSPDFUIKitMainThreadGuard(void)
{
    @autoreleasepool {
        for (NSString *selStr in @[@"setNeedsLayout", @"setNeedsDisplay", @"setNeedsDisplayInRect:"]) {
            SEL selector = NSSelectorFromString(selStr);
            SEL newSelector = NSSelectorFromString([NSString stringWithFormat:@"pspdf_%@", selStr]);
            if ([selStr hasSuffix:@":"]) {
                PSPDFReplaceMethodWithBlock(UIView.class, selector, newSelector, ^(__unsafe_unretained UIView *_self, CGRect r) {
                                                // Check for window, since *some* UIKit methods are indeed thread safe.
                                                // https://developer.apple.com/library/ios/#releasenotes/General/WhatsNewIniPhoneOS/Articles/iPhoneOS4.html
                                                /*
                                                   Drawing to a graphics context in UIKit is now thread-safe. Specifically:
                                                   - The routines used to access and manipulate the graphics context can now correctly handle contexts residing on different threads.
                                                   - String and image drawing is now thread-safe.
                                                   - Using color and font objects in multiple threads is now safe to do.
                                                 */
                                                if (_self.window) {
                                                    PSPDFAssertIfNotMainThread();
                                                }
                                                ((void ( *)(id, SEL, CGRect))objc_msgSend)(_self, newSelector, r);
                                            });
            } else {
                PSPDFReplaceMethodWithBlock(UIView.class, selector, newSelector, ^(__unsafe_unretained UIView *_self) {
                                                if (_self.window) {
                                                    if (!NSThread.isMainThread) {
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
                                                        dispatch_queue_t queue = dispatch_get_current_queue();
#pragma clang diagnostic pop
                                                        // iOS 8 layouts the MFMailComposeController in a background thread on an UIKit queue.
                                                        // https://github.com/PSPDFKit/PSPDFKit/issues/1423
                                                        if (!queue || !strstr(dispatch_queue_get_label(queue), "UIKit")) {
                                                            PSPDFAssertIfNotMainThread();
                                                        }
                                                    }
                                                }
                                                ((void ( *)(id, SEL))objc_msgSend)(_self, newSelector);
                                            });
            }
        }
    }
}
#endif // DEBUG

@end