threadguard.m

// Borrowed from PSPDFKit: https://gist.github.com/steipete/5664345

#define ENABLE_UIKIT_THREAD_ASSERTS (0)

#if ENABLE_UIKIT_THREAD_ASSERTS && defined(DEBUG)

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

#import "CRLogger.h"

static BOOL CRReplaceMethodWithBlock(Class c, SEL origSEL, SEL newSEL, id block)
{
    NSCParameterAssert(c && origSEL && newSEL && block);
    
    if ([c instancesRespondToSelector:newSEL])
        return YES; // Selector already implemented, skip silently.

    // Add the new method.
    Method origMethod = class_getInstanceMethod(c, origSEL);
    IMP imp = imp_implementationWithBlock(block);
    if (!class_addMethod(c, newSEL, imp, method_getTypeEncoding(origMethod))) {
        CRLogError(@"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;
}

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

static void CRAssertIfNotMainThread()
{
    NSCAssert(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\n"
                 "Break on CRAssertIfNotMainThread to find out where.\n\n"
                 "Stacktrace: %@", NSThread.callStackSymbols);
}

__attribute__((constructor)) static void CRUIKitMainThreadGuard()
{
    @autoreleasepool {
        for (NSString *selStr in @[@"setNeedsLayout", @"setNeedsDisplay", @"setNeedsDisplayInRect:"]) {
            SEL selector = NSSelectorFromString(selStr);
            SEL newSelector = NSSelectorFromString([NSString stringWithFormat:@"pspdf_%@", selStr]);
            if ([selStr hasSuffix:@":"]) {
                CRReplaceMethodWithBlock(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) CRAssertIfNotMainThread();
                    ((void ( *)(id, SEL, CGRect))objc_msgSend)(self, newSelector, r);
                });
            } else {
                CRReplaceMethodWithBlock(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")) {
                                CRAssertIfNotMainThread();
                            }
                        }
                    }
                    ((void ( *)(id, SEL))objc_msgSend)(self, newSelector);
                });
            }
        }
    }
}

#endif