Set tabBarController as associated object to UIViewControllers.

robbdimitrov · May 8, 2014 · eb74eaf · eb74eaf
1 parent 189fb87
commit eb74eaf
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 4 deletions.
diff --git a/Example/Classes/Controllers/RDVSecondViewController.m b/Example/Classes/Controllers/RDVSecondViewController.m
@@ -22,6 +22,8 @@
 // THE SOFTWARE.
 
 #import "RDVSecondViewController.h"
+#import "RDVTabBarController.h"
+#import "RDVTabBarItem.h"
 
 @implementation RDVSecondViewController
 

diff --git a/RDVTabBarController/RDVTabBarController.m b/RDVTabBarController/RDVTabBarController.m
@@ -23,6 +23,13 @@
 
 #import "RDVTabBarController.h"
 #import "RDVTabBarItem.h"
+#import <objc/runtime.h>
+
+@interface UIViewController (RDVTabBarControllerItemInternal)
+
+- (void)rdv_setTabBarController:(RDVTabBarController *)tabBarController;
+
+@end
 
 @interface RDVTabBarController () {
     UIView *_contentView;
@@ -134,10 +141,15 @@ - (void)setViewControllers:(NSArray *)viewControllers {
             RDVTabBarItem *tabBarItem = [[RDVTabBarItem alloc] init];
             [tabBarItem setTitle:viewController.title];
             [tabBarItems addObject:tabBarItem];
+            [viewController rdv_setTabBarController:self];
         }
 
         [[self tabBar] setItems:tabBarItems];
     } else {
+        for (UIViewController *viewController in _viewControllers) {
+            [viewController rdv_setTabBarController:nil];
+        }
+
         _viewControllers = nil;
     }
 }
@@ -269,14 +281,24 @@ - (void)tabBar:(RDVTabBar *)tabBar didSelectItemAtIndex:(NSInteger)index {
 
 #pragma mark - UIViewController+RDVTabBarControllerItem
 
+@implementation UIViewController (RDVTabBarControllerItemInternal)
+
+- (void)rdv_setTabBarController:(RDVTabBarController *)tabBarController {
+    objc_setAssociatedObject(self, @selector(rdv_tabBarController), tabBarController, OBJC_ASSOCIATION_ASSIGN);
+}
+
+@end
+
 @implementation UIViewController (RDVTabBarControllerItem)
 
 - (RDVTabBarController *)rdv_tabBarController {
-    UIViewController *parentViewController = [self parentViewController];
-    while (parentViewController && ![parentViewController isKindOfClass:[RDVTabBarController class]]) {
-        parentViewController = [parentViewController parentViewController];
+    RDVTabBarController *tabBarController = objc_getAssociatedObject(self, @selector(rdv_tabBarController));
+
+    if (!tabBarController && self.parentViewController) {
+        tabBarController = [self.parentViewController rdv_tabBarController];
     }
-    return (RDVTabBarController *)parentViewController;
+
+    return tabBarController;
 }
 
 - (RDVTabBarItem *)rdv_tabBarItem {