Handle "current page view" deleted from tile

// FREEBIE

Signal/src/ViewControllers/MediaGalleryViewController.swift

@@ -176,11 +176,11 @@ protocol MediaGalleryDataSource: class {
     func showAllMedia(focusedItem: MediaGalleryItem)
     func dismissMediaDetailViewController(_ mediaDetailViewController: MediaPageViewController, animated isAnimated: Bool, completion: (() -> Void)?)
 
-    func delete(items: [MediaGalleryItem])
+    func delete(items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate)
 }
 
 protocol MediaGalleryDataSourceDelegate: class {
-    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem])
+    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate)
     func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, deletedSections: IndexSet, deletedItems: [IndexPath])
 }
 
@@ -280,6 +280,7 @@ class MediaGalleryViewController: UINavigationController, MediaGalleryDataSource
         self.initialDetailItem = initialDetailItem
 
         let pageViewController = MediaPageViewController(initialItem: initialDetailItem, mediaGalleryDataSource: self, uiDatabaseConnection: self.uiDatabaseConnection, options: self.options)
+        self.addDataSourceDelegate(pageViewController)
 
         self.pageViewController = pageViewController
         self.setViewControllers([pageViewController], animated: false)
@@ -573,10 +574,7 @@ class MediaGalleryViewController: UINavigationController, MediaGalleryDataSource
         let vc = MediaTileViewController(mediaGalleryDataSource: self, uiDatabaseConnection: self.uiDatabaseConnection)
         vc.delegate = self
 
-        // dataSourceDelegate will either be this tile view, or the MessageDetailView, but they should
+        self.addDataSourceDelegate(vc)
-        // be mutually exclusive
-        assert(self.dataSourceDelegate == nil)
-        self.dataSourceDelegate = vc
 
         return vc
     }()
@@ -735,15 +733,21 @@ class MediaGalleryViewController: UINavigationController, MediaGalleryDataSource
         }
     }
 
-    weak var dataSourceDelegate: MediaGalleryDataSourceDelegate?
+    var dataSourceDelegates: [Weak<MediaGalleryDataSourceDelegate>] = []
+    func addDataSourceDelegate(_ dataSourceDelegate: MediaGalleryDataSourceDelegate) {
+        dataSourceDelegates.append(Weak(value: dataSourceDelegate))
+    }
+
     var deletedMessages: Set<TSMessage> = Set()
+    var deletedGalleryItems: Set<MediaGalleryItem> = Set()
 
-    func delete(items: [MediaGalleryItem]) {
+    func delete(items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate) {
         AssertIsOnMainThread()
 
         Logger.info("\(logTag) in \(#function) with items: \(items.map { ($0.attachmentStream, $0.message.timestamp) })")
 
-        dataSourceDelegate?.mediaGalleryDataSource(self, willDelete: items)
+        deletedGalleryItems.formUnion(items)
+        dataSourceDelegates.forEach { $0.value?.mediaGalleryDataSource(self, willDelete: items, initiatedBy: initiatedBy) }
 
         self.editingDatabaseConnection.asyncReadWrite { transaction in
             for item in items {
@@ -812,7 +816,7 @@ class MediaGalleryViewController: UINavigationController, MediaGalleryDataSource
             }
         }
 
-        dataSourceDelegate?.mediaGalleryDataSource(self, deletedSections: deletedSections, deletedItems: deletedIndexPaths)
+        dataSourceDelegates.forEach { $0.value?.mediaGalleryDataSource(self, deletedSections: deletedSections, deletedItems: deletedIndexPaths) }
     }
 
     let kGallerySwipeLoadBatchSize: UInt = 5
@@ -828,7 +832,17 @@ class MediaGalleryViewController: UINavigationController, MediaGalleryDataSource
         }
 
         let index: Int = galleryItems.index(after: currentIndex)
-        return galleryItems[safe: index]
+        guard let nextItem = galleryItems[safe: index] else {
+            // already at last item
+            return nil
+        }
+
+        guard !deletedGalleryItems.contains(nextItem) else {
+            Logger.debug("\(self.logTag) in \(#function) nextItem was deleted - Recursing.")
+            return galleryItem(after: nextItem)
+        }
+
+        return nextItem
     }
 
     internal func galleryItem(before currentItem: MediaGalleryItem) -> MediaGalleryItem? {
@@ -842,7 +856,17 @@ class MediaGalleryViewController: UINavigationController, MediaGalleryDataSource
         }
 
         let index: Int = galleryItems.index(before: currentIndex)
-        return galleryItems[safe: index]
+        guard let previousItem = galleryItems[safe: index] else {
+            // already at first item
+            return nil
+        }
+
+        guard !deletedGalleryItems.contains(previousItem) else {
+            Logger.debug("\(self.logTag) in \(#function) previousItem was deleted - Recursing.")
+            return galleryItem(before: previousItem)
+        }
+
+        return previousItem
     }
 
     var galleryItemCount: Int {

Signal/src/ViewControllers/MediaPageViewController.swift

@@ -18,7 +18,7 @@ public class GalleryItemBox: NSObject {
     }
 }
 
-fileprivate class Box<A> {
+private class Box<A> {
     var value: A
     init(_ val: A) {
         self.value = val
@@ -31,7 +31,7 @@ fileprivate extension MediaDetailViewController {
     }
 }
 
-class MediaPageViewController: UIPageViewController, UIPageViewControllerDataSource, UIPageViewControllerDelegate, MediaDetailViewControllerDelegate {
+class MediaPageViewController: UIPageViewController, UIPageViewControllerDataSource, UIPageViewControllerDelegate, MediaDetailViewControllerDelegate, MediaGalleryDataSourceDelegate {
 
     private weak var mediaGalleryDataSource: MediaGalleryDataSource?
 
@@ -317,23 +317,49 @@ class MediaPageViewController: UIPageViewController, UIPageViewControllerDataSou
         let deleteAction = UIAlertAction(title: NSLocalizedString("TXT_DELETE_TITLE", comment: ""),
                                          style: .destructive) { _ in
                                             let deletedItem = currentViewController.galleryItem
+                                            mediaGalleryDataSource.delete(items: [deletedItem], initiatedBy: self)
+        }
+        actionSheet.addAction(OWSAlerts.cancelAction)
+        actionSheet.addAction(deleteAction)
+
+        self.present(actionSheet, animated: true)
+    }
+
+    // MARK: MediaGalleryDataSourceDelegate
+
+    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate) {
+        Logger.debug("\(self.logTag) in \(#function)")
+
+        guard let currentItem = self.currentItem else {
+              owsFail("\(logTag) in \(#function) currentItem was unexpectedly nil")
+            return
+        }
+
+        guard items.contains(currentItem) else {
+            Logger.debug("\(self.logTag) in \(#function) irrelevant item")
+            return
+        }
+
+        // If we setCurrentItem with (animated: true) while this VC is in the background, then
+        // the next/previous cache isn't expired, and we're able to swipe back to the just-deleted vc.
+        // So to get the correct behavior, we should only animate these transitions when this
+        // vc is in the foreground
+        let isAnimated = initiatedBy === self
 
         if !self.sliderEnabled {
             // In message details, which doesn't use the slider, so don't swap pages.
-                                            } else if let nextItem = mediaGalleryDataSource.galleryItem(after: deletedItem) {
+        } else if let nextItem = mediaGalleryDataSource.galleryItem(after: currentItem) {
-                                                self.setCurrentItem(nextItem, direction: .forward, animated: true)
+            self.setCurrentItem(nextItem, direction: .forward, animated: isAnimated)
-                                            } else if let previousItem = mediaGalleryDataSource.galleryItem(before: deletedItem) {
+        } else if let previousItem = mediaGalleryDataSource.galleryItem(before: currentItem) {
-                                                self.setCurrentItem(previousItem, direction: .reverse, animated: true)
+            self.setCurrentItem(previousItem, direction: .reverse, animated: isAnimated)
         } else {
             // else we deleted the last piece of media, return to the conversation view
             self.dismissSelf(animated: true)
         }
-                                            mediaGalleryDataSource.delete(items: [deletedItem])
     }
-        actionSheet.addAction(OWSAlerts.cancelAction)
-        actionSheet.addAction(deleteAction)
 
-        self.present(actionSheet, animated: true)
+    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, deletedSections: IndexSet, deletedItems: [IndexPath]) {
+        // no-op
     }
 
     @objc
@@ -512,7 +538,7 @@ class MediaPageViewController: UIPageViewController, UIPageViewControllerDataSou
         }
 
         dismissSelf(animated: true) {
-            mediaGalleryDataSource.delete(items: [galleryItem])
+            mediaGalleryDataSource.delete(items: [galleryItem], initiatedBy: self)
         }
     }
 

Signal/src/ViewControllers/MediaTileViewController.swift

@@ -539,7 +539,7 @@ public class MediaTileViewController: UICollectionViewController, MediaGalleryDa
         }()
 
         let deleteAction = UIAlertAction(title: confirmationTitle, style: .destructive) { _ in
-            mediaGalleryDataSource.delete(items: items)
+            mediaGalleryDataSource.delete(items: items, initiatedBy: self)
             self.endSelectMode()
         }
 
@@ -557,7 +557,7 @@ public class MediaTileViewController: UICollectionViewController, MediaGalleryDa
 
     // MARK: MediaGalleryDataSourceDelegate
 
-    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem]) {
+    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate) {
         Logger.debug("\(self.logTag) in \(#function)")
 
         guard let collectionView = self.collectionView else {

Signal/src/ViewControllers/MessageDetailViewController.swift

@@ -758,7 +758,7 @@ class MessageDetailViewController: OWSViewController, UIScrollViewDelegate, Medi
 
     // MediaGalleryDataSourceDelegate
 
-    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem]) {
+    func mediaGalleryDataSource(_ mediaGalleryDataSource: MediaGalleryDataSource, willDelete items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate) {
         Logger.info("\(self.logTag) in \(#function)")
 
         guard (items.map({ $0.message }) == [self.message]) else {
@@ -785,7 +785,7 @@ class MessageDetailViewController: OWSViewController, UIScrollViewDelegate, Medi
         }
 
         let mediaGalleryViewController = MediaGalleryViewController(thread: self.thread, uiDatabaseConnection: self.uiDatabaseConnection)
-        mediaGalleryViewController.dataSourceDelegate = self
+        mediaGalleryViewController.addDataSourceDelegate(self)
         mediaGalleryViewController.presentDetailView(fromViewController: self, mediaMessage: self.message, replacingView: fromView)
     }
 }