From 6bdbc2e2456c698af634bd20670841ccbff23f3a Mon Sep 17 00:00:00 2001
From: Domas Nutautas <domas@trafi.com>
Date: Thu, 21 May 2020 21:38:09 +0300
Subject: [PATCH] Fix for unwinding in tabs with navigation stacks

---
 Sources/StoryFlow/Flow/Helpers/Unwinding.swift | 16 ++++++++--------
 StoryFlowTests/ImplicitFlowTests.swift         |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/Sources/StoryFlow/Flow/Helpers/Unwinding.swift b/Sources/StoryFlow/Flow/Helpers/Unwinding.swift
index ab63430..e2f85f8 100644
--- a/Sources/StoryFlow/Flow/Helpers/Unwinding.swift
+++ b/Sources/StoryFlow/Flow/Helpers/Unwinding.swift
@@ -12,9 +12,9 @@ extension UIViewController {
         Thread.onMain {
             if canHandle(updateType, from: source) {
                 return self
-            } else if let vc = navStack?.first(where: { $0.canHandle(updateType, from: source) }) {
+            } else if let vc = navStacks.first(where: { $0.canHandle(updateType, from: source) }) {
                 return vc
-            } else if let vc = tabs?.first(where: { $0.canHandle(updateType, from: source) }) {
+            } else if let vc = tabs.first(where: { $0.canHandle(updateType, from: source) }) {
                 return vc
             } else {
                 return (parent ?? presentingViewController)?.unwindVc(for: updateType, from: source)
@@ -47,16 +47,16 @@ private extension UIViewController {
         }
     }
 
-    func firstChild<T>(_ t: T.Type) -> T? {
-        self as? T ?? children.first { $0.firstChild(t) != nil }?.firstChild(t)
+    func allChild<T>(_ t: T.Type) -> [T] {
+        [self as? T].compactMap { $0 } + children.flatMap { $0.allChild(t) }
     }
 
-    var navStack: [UIViewController]? {
-        firstChild(UINavigationController.self)?.viewControllers.reversed()
+    var navStacks: [UIViewController] {
+        allChild(UINavigationController.self).flatMap { $0.viewControllers.reversed() }
     }
 
-    var tabs: [UIViewController]? {
-        firstChild(UITabBarController.self)?.viewControllers
+    var tabs: [UIViewController] {
+        allChild(UITabBarController.self).flatMap { $0.viewControllers ?? [] }
     }
 
     func isVc(for updateType: Any.Type) -> Bool {
diff --git a/StoryFlowTests/ImplicitFlowTests.swift b/StoryFlowTests/ImplicitFlowTests.swift
index 50887f2..82a10d8 100644
--- a/StoryFlowTests/ImplicitFlowTests.swift
+++ b/StoryFlowTests/ImplicitFlowTests.swift
@@ -496,7 +496,7 @@ class ImplicitFlowTests: XCTestCase {
         nav.setViewControllers([to, UIViewController()], animated: false)
 
         let tab = UITabBarController().visible()
-        tab.setViewControllers([from, nav], animated: false)
+        tab.setViewControllers([UINavigationController(rootViewController: from), nav], animated: false)
 
         let output = T()