TimeOutMobileTeam · May 31, 2013 18:39 · May 28, 2013 · May 28, 2013 · May 28, 2013 · May 28, 2013
diff --git a/PSPDFUIKitMainThreadGuard.m b/PSPDFUIKitMainThreadGuard.m
@@ -48,18 +48,16 @@ static void PSPDFAssertIfNotMainThread(void) {
     @autoreleasepool {
         for (NSString *selector in @[PROPERTY(setNeedsLayout), PROPERTY(setNeedsDisplay), PROPERTY(setNeedsDisplayInRect:)]) {
             SEL newSelector = NSSelectorFromString([NSString stringWithFormat:@"pspdf_%@", selector]);
-            for (Class theClass in @[UIView.class, CALayer.class]) {
-                if ([selector hasSuffix:@":"]) {
-                    PSPDFReplaceMethod(theClass, NSSelectorFromString(selector), newSelector, imp_implementationWithBlock(^(UIView *_self, CGRect r) {
-                        PSPDFAssertIfNotMainThread();
-                        ((void ( *)(id, SEL, CGRect))objc_msgSend)(_self, newSelector, r);
-                    }));
-                }else {
-                    PSPDFReplaceMethod(theClass, NSSelectorFromString(selector), newSelector, imp_implementationWithBlock(^(UIView *_self) {
-                        PSPDFAssertIfNotMainThread();
-                        ((void ( *)(id, SEL))objc_msgSend)(_self, newSelector);
-                    }));
-                }
+            if ([selector hasSuffix:@":"]) {
+                PSPDFReplaceMethod(UIView.class, NSSelectorFromString(selector), newSelector, imp_implementationWithBlock(^(UIView *_self, CGRect r) {
+                    PSPDFAssertIfNotMainThread();
+                    ((void ( *)(id, SEL, CGRect))objc_msgSend)(_self, newSelector, r);
+                }));
+            }else {
+                PSPDFReplaceMethod(UIView.class, NSSelectorFromString(selector), newSelector, imp_implementationWithBlock(^(UIView *_self) {
+                    PSPDFAssertIfNotMainThread();
+                    ((void ( *)(id, SEL))objc_msgSend)(_self, newSelector);
+                }));
             }
         }
     }

diff --git a/PSPDFUIKitMainThreadGuard.m b/PSPDFUIKitMainThreadGuard.m
@@ -4,6 +4,16 @@
 
 #define PROPERTY(propName) NSStringFromSelector(@selector(propName))
 
+// A better assert. NSAssert is too runtime dependant, and assert() doesn't log.
+// http://www.mikeash.com/pyblog/friday-qa-2013-05-03-proper-use-of-asserts.html
+// Accepts both:
+// - PSPDFAssert(x > 0);
+// - PSPDFAssert(y > 3, @"Bad value for y");
+#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)
+
 // You should only use this in debug builds. It doesn't use private API, but I wouldn't ship it.
 #ifdef DEBUG
 

diff --git a/PSPDFUIKitMainThreadGuard.m b/PSPDFUIKitMainThreadGuard.m
@@ -2,6 +2,8 @@
 #import <objc/message.h>
 #import <QuartzCore/QuartzCore.h>
 
+#define PROPERTY(propName) NSStringFromSelector(@selector(propName))
+
 // You should only use this in debug builds. It doesn't use private API, but I wouldn't ship it.
 #ifdef DEBUG
 

diff --git a/PSPDFUIKitMainThreadGuard.m b/PSPDFUIKitMainThreadGuard.m
@@ -52,4 +52,5 @@ static void PSPDFAssertIfNotMainThread(void) {
         }
     }
 }
+
 #endif
diff --git a/PSPDFUIKitMainThreadGuard.m b/PSPDFUIKitMainThreadGuard.m
@@ -2,6 +2,9 @@
 #import <objc/message.h>
 #import <QuartzCore/QuartzCore.h>
 
+// You should only use this in debug builds. It doesn't use private API, but I wouldn't ship it.
+#ifdef DEBUG
+
 static void PSPDFSwizzleMethod(Class c, SEL orig, SEL new) {
     Method origMethod = class_getInstanceMethod(c, orig);
     Method newMethod = class_getInstanceMethod(c, new);
@@ -48,4 +51,5 @@ static void PSPDFAssertIfNotMainThread(void) {
             }
         }
     }
-}
+}
+#endif
diff --git a/PSPDFUIKitMainThreadGuard.m b/PSPDFUIKitMainThreadGuard.m
@@ -0,0 +1,51 @@
+#import <objc/runtime.h>
+#import <objc/message.h>
+#import <QuartzCore/QuartzCore.h>
+
+static void PSPDFSwizzleMethod(Class c, SEL orig, SEL new) {
+    Method origMethod = class_getInstanceMethod(c, orig);
+    Method newMethod = class_getInstanceMethod(c, new);
+    if (class_addMethod(c, orig, method_getImplementation(newMethod), method_getTypeEncoding(newMethod))) {
+        class_replaceMethod(c, new, method_getImplementation(origMethod), method_getTypeEncoding(origMethod));
+    }else {
+        method_exchangeImplementations(origMethod, newMethod);
+    }
+}
+
+void PSPDFReplaceMethod(Class c, SEL orig, SEL newSel, IMP impl) {
+    Method method = class_getInstanceMethod(c, orig);
+    if (!class_addMethod(c, newSel, impl, method_getTypeEncoding(method))) {
+        PSPDFLogError(@"Failed to add method: %@ on %@", NSStringFromSelector(newSel), c);
+    }else PSPDFSwizzleMethod(c, orig, newSel);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////
+#pragma mark - Tracks down calls to UIKit from a Thread other than Main
+
+static 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]);
+}
+
+// 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 of PSPDFKit.
+// @note No private API is used here.
+__attribute__((constructor)) static void PSPDFUIKitMainThreadGuard(void) {
+    @autoreleasepool {
+        for (NSString *selector in @[PROPERTY(setNeedsLayout), PROPERTY(setNeedsDisplay), PROPERTY(setNeedsDisplayInRect:)]) {
+            SEL newSelector = NSSelectorFromString([NSString stringWithFormat:@"pspdf_%@", selector]);
+            for (Class theClass in @[UIView.class, CALayer.class]) {
+                if ([selector hasSuffix:@":"]) {
+                    PSPDFReplaceMethod(theClass, NSSelectorFromString(selector), newSelector, imp_implementationWithBlock(^(UIView *_self, CGRect r) {
+                        PSPDFAssertIfNotMainThread();
+                        ((void ( *)(id, SEL, CGRect))objc_msgSend)(_self, newSelector, r);
+                    }));
+                }else {
+                    PSPDFReplaceMethod(theClass, NSSelectorFromString(selector), newSelector, imp_implementationWithBlock(^(UIView *_self) {
+                        PSPDFAssertIfNotMainThread();
+                        ((void ( *)(id, SEL))objc_msgSend)(_self, newSelector);
+                    }));
+                }
+            }
+        }
+    }
+}
-Original file line number
+Diff line change
@@ -52,4 +52,5 @@ static void PSPDFAssertIfNotMainThread(void) {
             }
         }
     }
     #endif