diff --git a/LightRoute/LightRoute.swift b/LightRoute/LightRoute.swift index be265c3..03d4cf8 100644 --- a/LightRoute/LightRoute.swift +++ b/LightRoute/LightRoute.swift @@ -316,6 +316,8 @@ public final class TransitionPromise { } // Remove old link action then we can setup new transition action. self.postLinkAction = nil + + self.checkForPop(in: style) // Setup new transition action from transition case. self.postLintAction { [weak self] in @@ -395,7 +397,40 @@ public final class TransitionPromise { func postLintAction( _ completion: @escaping TransitionPostLinkAction) { self.postLinkAction = completion } - + + /// + /// This method check style for 'pop', if so then change destination view controller to correct from navigation stack. + /// - parameter style: Style from `TransitionStyle` class. + /// + private func checkForPop(in style: TransitionStyle) { + switch style { + case .navigationController(preferredStyle: let navStyle): + switch navStyle { + case .pop: + guard let dest = self.getDestinationFromNavigationStack() else { fatalError("[LightRoute]: Destination view controller was nil (not exist in navigation stack)") } + destination = dest + default: + break + } + + default: + break + } + } + + /// + /// This method found correct view controller from navigation stack. + /// - no parameter. + /// + private func getDestinationFromNavigationStack() -> UIViewController? { + guard let navController = root.navigationController else { + print("[LightRoute]: Transition error, navigation controller was nil.") + return nil + } + let destinationId = destination?.restorationIdentifier + + return navController.viewControllers.filter({ $0.restorationIdentifier == destinationId}).first + } }