Timeline state restoration

Tusker/Scenes/MainSceneDelegate.swift

@@ -95,8 +95,15 @@ class MainSceneDelegate: UIResponder, UIWindowSceneDelegate, TuskerSceneDelegate
     }
     
     func stateRestorationActivity(for scene: UIScene) -> NSUserActivity? {
-        if let mastodonController =  window?.windowScene?.session.mastodonController {
+        if let mastodonController = window?.windowScene?.session.mastodonController {
-            return UserActivityManager.mainSceneActivity(accountID: mastodonController.accountInfo!.id)
+            if let vcActivity = rootViewController?.stateRestorationActivity() {
+                vcActivity.isStateRestorationActivity = true
+                stateRestorationLogger.info("MainSceneDelegate returning stateRestorationActivity of type \(vcActivity.activityType, privacy: .public) from VC")
+                return vcActivity
+            } else {
+                // need to have an activity to make sure the same account is used
+                return UserActivityManager.mainSceneActivity(accountID: mastodonController.accountInfo!.id)
+            }
         } else {
             return nil
         }
@@ -144,7 +151,6 @@ class MainSceneDelegate: UIResponder, UIWindowSceneDelegate, TuskerSceneDelegate
         }
     }
     
-    
     func showAppOrOnboardingUI(session: UISceneSession? = nil) {
         let session = session ?? window!.windowScene!.session
         if LocalData.shared.onboardingComplete {
@@ -162,9 +168,12 @@ class MainSceneDelegate: UIResponder, UIWindowSceneDelegate, TuskerSceneDelegate
             
             activateAccount(account, animated: false)
             
-            if let activity = launchActivity,
+            if let activity = launchActivity {
-               activity.activityType != UserActivityType.mainScene.rawValue {
+                if activity.isStateRestorationActivity {
-                _ = activity.handleResume(manager: UserActivityManager(scene: window!.windowScene!))
+                    rootViewController?.restoreActivity(activity)
+                } else if activity.activityType != UserActivityType.mainScene.rawValue {
+                    _ = activity.handleResume(manager: UserActivityManager(scene: window!.windowScene!))
+                }
             }
         } else {
             window!.rootViewController = createOnboardingUI()

Tusker/Screens/Main/AccountSwitchingContainerViewController.swift

@@ -87,6 +87,16 @@ extension AccountSwitchingContainerViewController {
 }
 
 extension AccountSwitchingContainerViewController: TuskerRootViewController {
+    func stateRestorationActivity() -> NSUserActivity? {
+        loadViewIfNeeded()
+        return root.stateRestorationActivity()
+    }
+    
+    func restoreActivity(_ activity: NSUserActivity) {
+        loadViewIfNeeded()
+        root.restoreActivity(activity)
+    }
+    
     func presentCompose() {
         loadViewIfNeeded()
         root.presentCompose()

Tusker/Screens/Main/Duckable+Root.swift

@@ -11,6 +11,14 @@ import Duckable
 
 @available(iOS 16.0, *)
 extension DuckableContainerViewController: TuskerRootViewController {
+    func stateRestorationActivity() -> NSUserActivity? {
+        (child as? TuskerRootViewController)?.stateRestorationActivity()
+    }
+    
+    func restoreActivity(_ activity: NSUserActivity) {
+        (child as? TuskerRootViewController)?.restoreActivity(activity)
+    }
+    
     func presentCompose() {
         (child as? TuskerRootViewController)?.presentCompose()
     }

Tusker/Screens/Main/MainSplitViewController.swift

@@ -83,6 +83,14 @@ class MainSplitViewController: UISplitViewController {
         secondaryNavController.viewControllers = getOrCreateNavigationStack(item: item)
     }
     
+    func navigationStackFor(item: MainSidebarViewController.Item) -> [UIViewController]? {
+        if sidebar.selectedItem == item {
+            return secondaryNavController.viewControllers
+        } else {
+            return navigationStacks[item]
+        }
+    }
+    
     func getOrCreateNavigationStack(item: MainSidebarViewController.Item) -> [UIViewController] {
         if let existing = navigationStacks[item], existing.count > 0 {
             return existing
@@ -378,6 +386,36 @@ extension MainSplitViewController: TuskerNavigationDelegate {
 }
 
 extension MainSplitViewController: TuskerRootViewController {
+    func stateRestorationActivity() -> NSUserActivity? {
+        if traitCollection.horizontalSizeClass == .compact {
+            return tabBarViewController.stateRestorationActivity()
+        } else {
+            if let timelinePages = navigationStackFor(item: .tab(.timelines))?.first as? TimelinesPageViewController {
+                let timeline = timelinePages.pageControllers[timelinePages.currentIndex] as! TimelineViewController
+                return timeline.stateRestorationActivity()
+            } else {
+                stateRestorationLogger.fault("MainSplitViewController: Unable to create state restoration activity")
+                return nil
+            }
+        }
+    }
+    
+    func restoreActivity(_ activity: NSUserActivity) {
+        if traitCollection.horizontalSizeClass == .compact {
+            tabBarViewController.restoreActivity(activity)
+        } else {
+            if activity.activityType == UserActivityType.showTimeline.rawValue {
+                guard let timelinePages = navigationStackFor(item: .tab(.timelines))?.first as? TimelinesPageViewController else {
+                    stateRestorationLogger.fault("MainSplitViewController: Unable to restore timeline activity, couldn't find VC")
+                    return
+                }
+                timelinePages.restoreActivity(activity)
+            } else {
+                stateRestorationLogger.fault("MainSplitViewController: Unable to restore activity of unexpected type \(activity.activityType, privacy: .public)")
+            }
+        }
+    }
+    
     @objc func presentCompose() {
         self.compose()
     }

Tusker/Screens/Main/MainTabBarViewController.swift

@@ -233,6 +233,30 @@ extension MainTabBarViewController: TuskerNavigationDelegate {
 }
 
 extension MainTabBarViewController: TuskerRootViewController {
+    func stateRestorationActivity() -> NSUserActivity? {
+        let nav = viewController(for: .timelines) as! UINavigationController
+        guard let timelinePages = nav.viewControllers.first as? TimelinesPageViewController,
+              let timelineVC = timelinePages.pageControllers[timelinePages.currentIndex] as? TimelineViewController else {
+            stateRestorationLogger.fault("MainTabBarViewController: Unable to create state restoration actiivty, couldn't find timeline/page VC")
+            return nil
+        }
+        return timelineVC.stateRestorationActivity()
+    }
+    
+    func restoreActivity(_ activity: NSUserActivity) {
+        if activity.activityType == UserActivityType.showTimeline.rawValue {
+            let nav = viewController(for: .timelines) as! UINavigationController
+            guard let timelinePages = nav.viewControllers.first as? TimelinesPageViewController,
+                  let timelineVC = timelinePages.pageControllers[timelinePages.currentIndex] as? TimelineViewController else {
+                stateRestorationLogger.fault("MainTabBarViewController: Unable to restore timeline activity, couldn't find VC")
+                return
+            }
+            timelineVC.restoreActivity(activity)
+        } else {
+            stateRestorationLogger.fault("MainTabBarViewController: Unable to restore activity of unexpected type \(activity.activityType, privacy: .public)")
+        }
+    }
+    
     @objc func presentCompose() {
         compose()
     }

Tusker/Screens/Main/TuskerRootViewController.swift

@@ -9,6 +9,8 @@
 import UIKit
 
 protocol TuskerRootViewController: UIViewController, StatusBarTappableViewController {
+    func stateRestorationActivity() -> NSUserActivity?
+    func restoreActivity(_ activity: NSUserActivity)
     func presentCompose()
     func select(tab: MainTabBarViewController.Tab)
     func getTabController(tab: MainTabBarViewController.Tab) -> UIViewController?

Tusker/Screens/Timeline/TimelineViewController.swift

@@ -23,6 +23,7 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
     private(set) var dataSource: UICollectionViewDiffableDataSource<Section, Item>!
     
     private var contentOffsetObservation: NSKeyValueObservation?
+    private var activityToRestore: NSUserActivity?
     
     init(for timeline: Timeline, mastodonController: MastodonController!) {
         self.timeline = timeline
@@ -88,7 +89,7 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
 #endif
         
         contentOffsetObservation = collectionView.observe(\.contentOffset) { [weak self] _, _ in
-           if let indexPath = self?.dataSource.indexPath(for: .gap),
+            if let indexPath = self?.dataSource.indexPath(for: .gap),
               let cell = self?.collectionView.cellForItem(at: indexPath) as? TimelineGapCollectionViewCell {
                cell.update()
            }
@@ -156,9 +157,15 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
             collectionView.deselectItem(at: $0, animated: true)
         }
         
-        Task {
+        if case .notLoadedInitial = controller.state {
-            if case .notLoadedInitial = controller.state {
+            if doRestore() {
-                await controller.loadInitial()
+                Task {
+                    await checkPresent()
+                }
+            } else {
+                Task {
+                    await controller.loadInitial()
+                }
             }
         }
     }
@@ -176,6 +183,108 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
         }
     }
     
+    func stateRestorationActivity() -> NSUserActivity? {
+        let visible = collectionView.indexPathsForVisibleItems.sorted()
+        let snapshot = dataSource.snapshot()
+        guard let currentAccountID = mastodonController.accountInfo?.id,
+              !visible.isEmpty,
+              let statusesSection = snapshot.sectionIdentifiers.firstIndex(of: .statuses),
+              let firstVisible = visible.first(where: { $0.section == statusesSection }),
+              let lastVisible = visible.last(where: { $0.section == statusesSection }) else {
+            return nil
+        }
+        let allItems = snapshot.itemIdentifiers(inSection: .statuses)
+        
+        let startIndex = max(0, firstVisible.row - 20)
+        let endIndex = min(allItems.count - 1, lastVisible.row + 20)
+        
+        let firstVisibleItem: Item
+        var items = allItems[startIndex...endIndex]
+        if let gapIndex = items.firstIndex(of: .gap) {
+            // if the gap is above the top visible item, we take everything below the gap
+            // otherwise, we take everything above the gap
+            if gapIndex <= firstVisible.row {
+                items = allItems[(gapIndex + 1)...endIndex]
+                if gapIndex == firstVisible.row {
+                    firstVisibleItem = allItems.first!
+                } else {
+                    assert(items.indices.contains(firstVisible.row))
+                    firstVisibleItem = allItems[firstVisible.row]
+                }
+            } else {
+                items = allItems[startIndex..<gapIndex]
+                firstVisibleItem = allItems[firstVisible.row]
+            }
+        } else {
+            firstVisibleItem = allItems[firstVisible.row]
+        }
+        let ids = items.map {
+            if case .status(id: let id, state: _) = $0 {
+                return id
+            } else {
+                fatalError()
+            }
+        }
+        let firstVisibleID: String
+        if case .status(id: let id, state: _) = firstVisibleItem {
+            firstVisibleID = id
+        } else {
+            fatalError()
+        }
+        stateRestorationLogger.debug("TimelineViewController: creating state restoration activity with topID \(firstVisibleID)")
+        
+        let activity = UserActivityManager.showTimelineActivity(timeline: timeline, accountID: currentAccountID)!
+        activity.addUserInfoEntries(from: [
+            "statusIDs": ids,
+            "topID": firstVisibleID,
+        ])
+        activity.isEligibleForPrediction = false
+        return activity
+    }
+    
+    func restoreActivity(_ activity: NSUserActivity) {
+        self.activityToRestore = activity
+    }
+    
+    private func doRestore() -> Bool {
+        guard let activity = activityToRestore,
+              let statusIDs = activity.userInfo?["statusIDs"] as? [String] else {
+            stateRestorationLogger.fault("TimelineViewController: activity missing statusIDs")
+            return false
+        }
+        activityToRestore = nil
+        loadViewIfNeeded()
+        controller.restoreInitial {
+            var snapshot = dataSource.snapshot()
+            snapshot.appendSections([.statuses])
+            let items = statusIDs.map { Item.status(id: $0, state: .unknown) }
+            snapshot.appendItems(items, toSection: .statuses)
+            dataSource.apply(snapshot, animatingDifferences: false) {
+                if let topID = activity.userInfo?["topID"] as? String,
+                   let index = statusIDs.firstIndex(of: topID),
+                   let indexPath = self.dataSource.indexPath(for: items[index]) {
+                    // it sometimes takes multiple attempts to convert on the right scroll position
+                    // since we're dealing with a bunch of unmeasured cells, so just try a few times in a loop
+                    var count = 0
+                    while count < 5 {
+                        count += 1
+                        let origOffset = self.collectionView.contentOffset
+                        self.collectionView.layoutIfNeeded()
+                        self.collectionView.scrollToItem(at: indexPath, at: .top, animated: false)
+                        let newOffset = self.collectionView.contentOffset
+                        if abs(origOffset.y - newOffset.y) <= 1 {
+                            break
+                        }
+                    }
+                    stateRestorationLogger.fault("TimelineViewController: restored statuses with top ID \(topID)")
+                } else {
+                    stateRestorationLogger.fault("TimelineViewController: restored statuses, but couldn't find top ID")
+                }
+            }
+        }
+        return true
+    }
+    
     private func removeTimelineDescriptionCell() {
         var snapshot = dataSource.snapshot()
         snapshot.deleteSections([.header])
@@ -190,10 +299,7 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
             return
         }
         Task {
-            if case .idle = controller.state,
+            await checkPresent()
-               let presentItems = try? await loadInitial() {
-                insertPresentItemsIfNecessary(presentItems)
-            }
         }
     }
     
@@ -214,43 +320,27 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
         }
     }
     
+    private func checkPresent() async {
+        if case .idle = controller.state,
+           let presentItems = try? await loadInitial() {
+            insertPresentItemsIfNecessary(presentItems)
+        }
+    }
+    
     private func insertPresentItemsIfNecessary(_ presentItems: [String]) {
         var snapshot = dataSource.snapshot()
         let currentItems = snapshot.itemIdentifiers(inSection: .statuses)
         if case .status(id: let firstID, state: _) = currentItems.first,
            // if there's no overlap between presentItems and the existing items in the data source, insert the present items and prompt the user
            !presentItems.contains(firstID) {
-            let applySnapshotBeforeScrolling: Bool
             
             // remove any existing gap, if there is one
             if let index = currentItems.lastIndex(of: .gap) {
                 snapshot.deleteItems(Array(currentItems[index...]))
-                
-                let statusesSection = snapshot.indexOfSection(.statuses)!
-                if collectionView.indexPathsForVisibleItems.contains(IndexPath(row: index, section: statusesSection)) {
-                    // the gap cell is on screen
-                    applySnapshotBeforeScrolling = false
-                } else if let topMostVisibleCell = collectionView.indexPathsForVisibleItems.first(where: { $0.section == statusesSection }),
-                          index < topMostVisibleCell.row {
-                    // the gap cell is above the top, so applying the snapshot would remove the currently-viewed statuses
-                    applySnapshotBeforeScrolling = false
-                } else {
-                    // the gap cell is below the bottom of the screen
-                    applySnapshotBeforeScrolling = true
-                }
-            } else {
-                // there is no existing gap
-                applySnapshotBeforeScrolling = true
             }
             snapshot.insertItems([.gap], beforeItem: currentItems.first!)
             snapshot.insertItems(presentItems.map { .status(id: $0, state: .unknown) }, beforeItem: .gap)
             
-            if applySnapshotBeforeScrolling {
-                let firstVisibleIndexPath = collectionView.indexPathsForVisibleItems.min()!
-                let firstVisibleItem = dataSource.itemIdentifier(for: firstVisibleIndexPath)!
-                applySnapshot(snapshot, maintainingBottomRelativeScrollPositionOf: firstVisibleItem)
-            }
-            
             var config = ToastConfiguration(title: "Jump to present")
             config.edge = .top
             config.systemImageName = "arrow.up"
@@ -258,9 +348,7 @@ class TimelineViewController: UIViewController, TimelineLikeCollectionViewContro
             config.action = { [unowned self] toast in
                 toast.dismissToast(animated: true)
                 
-                if !applySnapshotBeforeScrolling {
+                self.dataSource.apply(snapshot, animatingDifferences: false)
-                    self.dataSource.apply(snapshot, animatingDifferences: false)
-                }
                 
                 self.collectionView.scrollToTop()
             }

Tusker/Screens/Timeline/TimelinesPageViewController.swift

@@ -45,5 +45,23 @@ class TimelinesPageViewController: SegmentedPageViewController {
     required init?(coder: NSCoder) {
         fatalError("init(coder:) has not been implemented")
     }
+    
+    func restoreActivity(_ activity: NSUserActivity) {
+        guard let timeline = UserActivityManager.getTimeline(from: activity) else {
+            return
+        }
+        switch timeline {
+        case .home:
+            selectPage(at: 0, animated: false)
+        case .public(local: false):
+            selectPage(at: 1, animated: false)
+        case .public(local: true):
+            selectPage(at: 2, animated: false)
+        default:
+            return
+        }
+        let timelineVC = pageControllers[currentIndex] as! TimelineViewController
+        timelineVC.restoreActivity(activity)
+    }
 
 }

Tusker/Shortcuts/NSUserActivity+Extensions.swift

@@ -21,6 +21,18 @@ extension NSUserActivity {
             userInfo!["displaysAuxiliaryScene"] = newValue
         }
     }
+    
+    var isStateRestorationActivity: Bool {
+        get {
+            (userInfo?["isStateRestorationActivity"] as? Bool) ?? false
+        }
+        set {
+            if userInfo == nil {
+                userInfo = [:]
+            }
+            userInfo!["isStateRestorationActivity"] = newValue
+        }
+    }
 
     convenience init(type: UserActivityType) {
         self.init(activityType: type.rawValue)

Tusker/TimelineLikeController.swift

@@ -79,6 +79,16 @@ class TimelineLikeController<Item> {
         }
     }
     
+    /// Used to indicate to the controller that the initial set of posts have been restored externally.
+    func restoreInitial(doRestore: () -> Void) {
+        guard state == .notLoadedInitial else {
+            return
+        }
+        state = .restoringInitial
+        doRestore()
+        state = .idle
+    }
+    
     func loadNewer() async {
         guard state == .idle else {
             return
@@ -188,6 +198,7 @@ class TimelineLikeController<Item> {
     enum State: Equatable, CustomDebugStringConvertible {
         case notLoadedInitial
         case idle
+        case restoringInitial
         case loadingInitial(LoadAttemptToken, hasAddedLoadingIndicator: Bool)
         case loadingNewer(LoadAttemptToken)
         case loadingOlder(LoadAttemptToken, hasAddedLoadingIndicator: Bool)
@@ -199,6 +210,8 @@ class TimelineLikeController<Item> {
                 return "notLoadedInitial"
             case .idle:
                 return "idle"
+            case .restoringInitial:
+                return "restoringInitial"
             case .loadingInitial(let token, let hasAddedLoadingIndicator):
                 return "loadingInitial(\(ObjectIdentifier(token)), hasAddedLoadingIndicator: \(hasAddedLoadingIndicator))"
             case .loadingNewer(let token):
@@ -214,7 +227,7 @@ class TimelineLikeController<Item> {
             switch self {
             case .notLoadedInitial:
                 switch to {
-                case .loadingInitial(_, _):
+                case .restoringInitial, .loadingInitial(_, _):
                     return true
                 default:
                     return false
@@ -226,6 +239,8 @@ class TimelineLikeController<Item> {
                 default:
                     return false
                 }
+            case .restoringInitial:
+                return to == .idle
             case .loadingInitial(let token, let hasAddedLoadingIndicator):
                 return to == .notLoadedInitial || to == .idle || (!hasAddedLoadingIndicator && to == .loadingInitial(token, hasAddedLoadingIndicator: true))
             case .loadingNewer(_):