 }
 
 private struct RouteHandlerAssociatedKey {
-    static var viewController = "viewController"
+    static var navigationController = "navigationController"
 }
 
 extension RouteHandler {
     public var navigationController: UINavigationController? {
-        return viewController?.navigationController
+        return objc_getAssociatedObject(self, &RouteHandlerAssociatedKey.navigationController) as? UINavigationController
     }
 
-    fileprivate var viewController: UIViewController? {
-        get { return objc_getAssociatedObject(self, &RouteHandlerAssociatedKey.viewController) as? UIViewController }
-        set { objc_setAssociatedObject(self, &RouteHandlerAssociatedKey.viewController, newValue, .OBJC_ASSOCIATION_ASSIGN) }
+    fileprivate func setNavigationController(_ navigationController: UINavigationController?) {
+        objc_setAssociatedObject(self, &RouteHandlerAssociatedKey.navigationController, navigationController, .OBJC_ASSOCIATION_ASSIGN)
     }
 }
 


extension XCTableViewDataSource {
    public var router: Router {
        guard let router = tableView?.viewController?.router else {
            #if DEBUG
            if isDebuggerAttached {
                fatalError("Datasource doesn't have a view controller.")
            }
            #endif
            return Router(viewController: UIViewController())
        }

        return router

}

public class Router {
    private weak var navigationController: UINavigationController?
    private var routeHandlers: [String: Any] = [:]

    fileprivate init(navigationController: UINavigationController?) {
        self.navigationController = navigationController
    }

    public func register<T: RouteHandler>(_ handler: @autoclosure () -> T) -> T {
        let key = "\(T.self)"

        guard let existingHandler = routeHandlers[key] as? T else {
            var handler = handler()
            handler.setNavigationController(navigationController)
            routeHandlers[key] = handler
            return handler
        }

        return existingHandler
    }
}

// MARK: Extensions

extension UIViewController {
    public var router: Router {
        guard let navigationController = navigationController else {
            if let navController = self as? UINavigationController {
                return navController._navigationControllerRouter
            }

            fatalError("Router requires a navigation controller.")
        }

        return navigationController._navigationControllerRouter
    }
}

extension UINavigationController {
    private struct AssociatedKey {
        static var router = "router"
    }

    fileprivate var _navigationControllerRouter: Router {
        get {
            let router: Router

            if let existingRouter: Router = associatedObject(&AssociatedKey.router) {
                router = existingRouter
            } else {
                router = Router(navigationController: self)
                self._navigationControllerRouter = router
            }

            return router
