From 1df3158050357d0330cf43987e50b0c2d0b9a5fd Mon Sep 17 00:00:00 2001
From: Soner YUKSEL <grkmyksl@gmail.com>
Date: Sat, 2 Sep 2023 03:15:55 -0400
Subject: [PATCH] Fix #7996: Switching Private Mode Tab Tray cause position
 loss in Normal Mode (#8015)

---
 .../Brave/Frontend/Browser/TabManager.swift   |  1 +
 .../Tabs/TabTray/TabTrayController.swift      | 32 +++++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/Sources/Brave/Frontend/Browser/TabManager.swift b/Sources/Brave/Frontend/Browser/TabManager.swift
index afa99c1a5f4..b6e62555c65 100644
--- a/Sources/Brave/Frontend/Browser/TabManager.swift
+++ b/Sources/Brave/Frontend/Browser/TabManager.swift
@@ -85,6 +85,7 @@ class TabManager: NSObject {
     return _selectedIndex
   }
   var normalTabSelectedIndex: Int = 0
+  var privateTabSelectedIndex: Int = 0
   var tempTabs: [Tab]?
   private weak var rewards: BraveRewards?
   private weak var tabGeneratorAPI: BraveTabGeneratorAPI?
diff --git a/Sources/Brave/Frontend/Browser/Tabs/TabTray/TabTrayController.swift b/Sources/Brave/Frontend/Browser/Tabs/TabTray/TabTrayController.swift
index d32002cbdbc..2fd5b819ba4 100644
--- a/Sources/Brave/Frontend/Browser/Tabs/TabTray/TabTrayController.swift
+++ b/Sources/Brave/Frontend/Browser/Tabs/TabTray/TabTrayController.swift
@@ -621,8 +621,10 @@ class TabTrayController: AuthenticationController {
       return
     }
 
-    // Record the slected index before private mode navigation
-    if !privateMode {
+    // Record the selected index before mode navigation
+    if privateMode {
+      tabManager.privateTabSelectedIndex = Preferences.Privacy.persistentPrivateBrowsing.value ? tabManager.selectedIndex : 0
+    } else {
       tabManager.normalTabSelectedIndex = tabManager.selectedIndex
     }
     
@@ -646,8 +648,19 @@ class TabTrayController: AuthenticationController {
           tabManager.addTabAndSelect(isPrivate: true)
         }
         
+        let privateModeTabSelected = tabManager.allTabs[safe: tabManager.privateTabSelectedIndex]
+
+        if Preferences.Privacy.persistentPrivateBrowsing.value {
+          tabManager.selectTab(privateModeTabSelected)
+        }
         tabTrayView.hidePrivateModeInfo()
         tabTrayView.collectionView.reloadData()
+        
+        // When you go back from normal mode, current tab should be selected
+        // in case private tabs are persistent
+        if Preferences.Privacy.persistentPrivateBrowsing.value {
+          scrollToSelectedTab(privateModeTabSelected)
+        }
         navigationController?.setNavigationBarHidden(false, animated: false)
       }
     } else {
@@ -655,8 +668,12 @@ class TabTrayController: AuthenticationController {
       
       // When you go back from private mode, a previous current tab is selected
       // Reloding the collection view in order to mark the selecte the tab
-      tabManager.selectTab(tabManager.tabsForCurrentMode[safe: tabManager.normalTabSelectedIndex])
+      let normalModeTabSelected = tabManager.allTabs[safe: tabManager.normalTabSelectedIndex]
+      
+      tabManager.selectTab(normalModeTabSelected)
       tabTrayView.collectionView.reloadData()
+      
+      scrollToSelectedTab(normalModeTabSelected)
       navigationController?.setNavigationBarHidden(false, animated: false)
     }
     
@@ -724,6 +741,15 @@ class TabTrayController: AuthenticationController {
     present(settingsNavigationController, animated: true)
   }
   
+  private func scrollToSelectedTab(_ tab: Tab?) {
+    if let selectedTab = tab,
+       let selectedIndexPath = dataSource.indexPath(for: selectedTab) {
+      DispatchQueue.main.async {
+        self.tabTrayView.collectionView.scrollToItem(at: selectedIndexPath, at: .centeredVertically, animated: false)
+      }
+    }
+  }
+  
   @objc private func tappedCollectionViewBackground() {
     if traitCollection.horizontalSizeClass == .compact {
       doneAction()