From 7840a35e2940a97f6880689c147efa48214f7e88 Mon Sep 17 00:00:00 2001
From: Justin Mazzocchi <2831158+jzzocc@users.noreply.github.com>
Date: Fri, 30 Oct 2020 17:53:57 -0700
Subject: [PATCH] Fix scroll position maintenance

---
 DB/Sources/DB/Entities/CollectionItem.swift       | 15 +++++++++++++++
 Data Sources/TableViewDataSource.swift            |  8 ++++++++
 View Controllers/TableViewController.swift        |  8 ++++----
 .../ViewModels/CollectionItemsViewModel.swift     |  6 +++---
 .../ViewModels/Entities/CollectionUpdate.swift    |  2 +-
 5 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/DB/Sources/DB/Entities/CollectionItem.swift b/DB/Sources/DB/Entities/CollectionItem.swift
index addae33..92856af 100644
--- a/DB/Sources/DB/Entities/CollectionItem.swift
+++ b/DB/Sources/DB/Entities/CollectionItem.swift
@@ -10,6 +10,8 @@ public enum CollectionItem: Hashable {
 }
 
 public extension CollectionItem {
+    typealias Id = String
+
     struct StatusConfiguration: Hashable {
         public let showContentToggled: Bool
         public let showAttachmentsToggled: Bool
@@ -32,6 +34,19 @@ public extension CollectionItem {
             self.hasReplyFollowing = hasReplyFollowing
         }
     }
+
+    var itemId: Id? {
+        switch  self {
+        case let .status(status, _):
+            return status.id
+        case .loadMore:
+            return nil
+        case let .account(account):
+            return account.id
+        case let .notification(notification, _):
+            return notification.id
+        }
+    }
 }
 
 public extension CollectionItem.StatusConfiguration {
diff --git a/Data Sources/TableViewDataSource.swift b/Data Sources/TableViewDataSource.swift
index c6f21a3..a95c857 100644
--- a/Data Sources/TableViewDataSource.swift	
+++ b/Data Sources/TableViewDataSource.swift	
@@ -42,3 +42,11 @@ final class TableViewDataSource: UITableViewDiffableDataSource<Int, CollectionIt
         }
     }
 }
+
+extension TableViewDataSource {
+    func indexPath(itemId: CollectionItem.Id) -> IndexPath? {
+        guard let item = snapshot().itemIdentifiers.first(where: { $0.itemId == itemId }) else { return nil }
+
+        return indexPath(for: item)
+    }
+}
diff --git a/View Controllers/TableViewController.swift b/View Controllers/TableViewController.swift
index 908cc7b..f0eefa7 100644
--- a/View Controllers/TableViewController.swift	
+++ b/View Controllers/TableViewController.swift	
@@ -260,8 +260,8 @@ private extension TableViewController {
         var offsetFromNavigationBar: CGFloat?
 
         if
-            let item = update.maintainScrollPosition,
-            let indexPath = dataSource.indexPath(for: item),
+            let itemId = update.maintainScrollPositionItemId,
+            let indexPath = dataSource.indexPath(itemId: itemId),
             let navigationBar = navigationController?.navigationBar {
             let navigationBarMaxY = tableView.convert(navigationBar.bounds, from: navigationBar).maxY
             offsetFromNavigationBar = tableView.rectForRow(at: indexPath).origin.y - navigationBarMaxY
@@ -271,8 +271,8 @@ private extension TableViewController {
             guard let self = self else { return }
 
             if
-                let item = update.maintainScrollPosition,
-                let indexPath = self.dataSource.indexPath(for: item) {
+                let itemId = update.maintainScrollPositionItemId,
+                let indexPath = self.dataSource.indexPath(itemId: itemId) {
                 if self.viewModel.shouldAdjustContentInset {
                     self.tableView.contentInset.bottom = max(
                         0,
diff --git a/ViewModels/Sources/ViewModels/CollectionItemsViewModel.swift b/ViewModels/Sources/ViewModels/CollectionItemsViewModel.swift
index 74173f1..ad08994 100644
--- a/ViewModels/Sources/ViewModels/CollectionItemsViewModel.swift
+++ b/ViewModels/Sources/ViewModels/CollectionItemsViewModel.swift
@@ -16,7 +16,7 @@ final public class CollectionItemsViewModel: ObservableObject {
     private let eventsSubject = PassthroughSubject<CollectionItemEvent, Never>()
     private let loadingSubject = PassthroughSubject<Bool, Never>()
     private let expandAllSubject: CurrentValueSubject<ExpandAllState, Never>
-    private var maintainScrollPosition: CollectionItem?
+    private var maintainScrollPositionItemId: CollectionItem.Id?
     private var topVisibleIndexPath = IndexPath(item: 0, section: 0)
     private let lastReadId = CurrentValueSubject<String?, Never>(nil)
     private var lastSelectedLoadMore: LoadMore?
@@ -57,7 +57,7 @@ extension CollectionItemsViewModel: CollectionViewModel {
     public var updates: AnyPublisher<CollectionUpdate, Never> {
         items.map { [weak self] in
             CollectionUpdate(items: $0,
-                             maintainScrollPosition: self?.maintainScrollPosition)
+                             maintainScrollPositionItemId: self?.maintainScrollPositionItemId)
         }
         .eraseToAnyPublisher()
     }
@@ -267,7 +267,7 @@ private extension CollectionItemsViewModel {
     }
 
     func process(items: [[CollectionItem]]) {
-        maintainScrollPosition = itemForScrollPositionMaintenance(newItems: items)
+        maintainScrollPositionItemId = itemForScrollPositionMaintenance(newItems: items)?.itemId
         self.items.send(items)
 
         let itemsSet = Set(items.reduce([], +))
diff --git a/ViewModels/Sources/ViewModels/Entities/CollectionUpdate.swift b/ViewModels/Sources/ViewModels/Entities/CollectionUpdate.swift
index 41eb994..15b9260 100644
--- a/ViewModels/Sources/ViewModels/Entities/CollectionUpdate.swift
+++ b/ViewModels/Sources/ViewModels/Entities/CollectionUpdate.swift
@@ -2,5 +2,5 @@
 
 public struct CollectionUpdate: Hashable {
     public let items: [[CollectionItem]]
-    public let maintainScrollPosition: CollectionItem?
+    public let maintainScrollPositionItemId: String?
 }