Merge branch 'master' of github.com:tokorom/SegueContext

tokorom · Jul 18, 2015 · 9dfe048 · 9dfe048
2 parents 4665ea8 + ff02326
commit 9dfe048
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 25 deletions.
diff --git a/Source/SegueContext.swift b/Source/SegueContext.swift
@@ -234,31 +234,18 @@ extension UIViewController {
     }
 
     public func performSegueWithIdentifier(identifier: String, sender: AnyObject? = nil, context: Any?) -> SegueContext {
-        let customContext: Context
-        if let context = context as? Context {
-            customContext = context
-        } else {
-            customContext = Context(object: context)
-        }
-        customContext.segueIdentifier = identifier
-        self.sendCustomContext = customContext
-
-        self.performSegueWithIdentifier(identifier, sender: sender)
-
-        return SegueContext()
+        return self.performSegueWithIdentifier(identifier, sender: sender, context: context, callback: nil)
     }
 
     public func performSegueWithIdentifier(identifier: String, sender: AnyObject? = nil, callback: Any?) -> SegueContext {
-        let customContextForCallback = Context(callback: callback)
-        customContextForCallback.segueIdentifier = identifier
-        self.sendCustomContextForCallback = customContextForCallback
-
-        self.performSegueWithIdentifier(identifier, sender: sender)
-
-        return SegueContext()
+        return self.performSegueWithIdentifier(identifier, sender: sender, context: nil, callback: callback)
     }
 
     public func performSegueWithIdentifier(identifier: String, sender: AnyObject? = nil, context: Any?, callback: Any?) -> SegueContext {
+#if !DISABLE_SWIZZLING
+        self.replacePrepareForSegueIfNeeded()
+#endif
+
         let customContext: Context
         if let context = context as? Context {
             customContext = context
@@ -419,10 +406,36 @@ extension UIViewController {
         }
     }
 
-    /* For connection to the Objective-C */
-    func swc_prepareForSegue(segue: UIStoryboardSegue!, sender: AnyObject?) {
-        if let segue = segue,
-               destination = segue.destinationViewController as? UIViewController,
+}
+
+#if !DISABLE_SWIZZLING
+
+// MARK: - Swizzling
+
+var swc_swizzled_already: UInt8 = 0
+
+extension UIViewController {
+
+    class func replacePrepareForSegueIfNeeded() {
+        if nil == objc_getAssociatedObject(self, &swc_swizzled_already) {
+            objc_setAssociatedObject(self, &swc_swizzled_already, NSNumber(bool: true), objc_AssociationPolicy(OBJC_ASSOCIATION_RETAIN_NONATOMIC))
+            let original = class_getInstanceMethod(self, "prepareForSegue:sender:")
+            let replaced = class_getInstanceMethod(self, "swc_wrapped_prepareForSegue:sender:")
+            method_exchangeImplementations(original, replaced)
+        }
+    }
+
+    func replacePrepareForSegueIfNeeded() {
+        self.dynamicType.replacePrepareForSegueIfNeeded()
+    }
+
+    func swc_wrapped_prepareForSegue(segue: UIStoryboardSegue, sender: AnyObject?) {
+        self.swc_wrapped_prepareForSegue(segue, sender: sender)
+        self.swc_prepareForSegue(segue, sender: sender)
+    }
+
+    func swc_prepareForSegue(segue: UIStoryboardSegue, sender: AnyObject?) {
+        if let destination = segue.destinationViewController as? UIViewController,
                source = segue.sourceViewController as? UIViewController,
                customContext = source.sendCustomContext
         {
@@ -433,8 +446,7 @@ extension UIViewController {
             destination.configureCustomContext(customContext)
             source.sendCustomContext = nil
         }
-        if let segue = segue,
-               destination = segue.destinationViewController as? UIViewController,
+        if let destination = segue.destinationViewController as? UIViewController,
                source = segue.sourceViewController as? UIViewController,
                customContextForCallback = source.sendCustomContextForCallback
         {
@@ -448,3 +460,5 @@ extension UIViewController {
     }
 
 }
+
+#endif
diff --git a/Source/UIViewController+SegueContext.m b/Source/UIViewController+SegueContext.m
@@ -12,6 +12,8 @@ - (void)swc_prepareForSegue:(UIStoryboardSegue *)segue sender:(id)sender;
 
 @implementation UIViewController (SegueContext)
 
+#ifdef DISABLE_SWIZZLING
+
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wobjc-protocol-method-implementation"
 
@@ -22,4 +24,6 @@ - (void)prepareForSegue:(UIStoryboardSegue *)segue sender:(id)sender
 
 #pragma clang diagnostic pop
 
+#endif
+
 @end