From b581a3e2661bf5f6ef10ad5d324f28f6a316fc51 Mon Sep 17 00:00:00 2001
From: Ryo Aoyama <r.fe51028.r@gmail.com>
Date: Tue, 10 Dec 2019 15:06:12 +0900
Subject: [PATCH] Add completion block to apply function

---
 .../AppKit/CocoaCollectionViewDiffableDataSource.swift |  6 +++++-
 Sources/Internal/DiffableDataSourceCore.swift          | 10 +++++++---
 Sources/UIKit/CollectionViewDiffableDataSource.swift   |  8 ++++++--
 Sources/UIKit/TableViewDiffableDataSource.swift        |  8 ++++++--
 4 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/Sources/AppKit/CocoaCollectionViewDiffableDataSource.swift b/Sources/AppKit/CocoaCollectionViewDiffableDataSource.swift
index bd6ed6b..aa5eb7e 100644
--- a/Sources/AppKit/CocoaCollectionViewDiffableDataSource.swift
+++ b/Sources/AppKit/CocoaCollectionViewDiffableDataSource.swift
@@ -33,6 +33,8 @@ open class CocoaCollectionViewDiffableDataSource<SectionIdentifierType: Hashable
     ///   - snapshot: A snapshot object to be applied to data model.
     ///   - animatingDifferences: A Boolean value indicating whether to update with
     ///                           diffing animation.
+    ///   - completion: An optional completion block which is called when the complete
+    ///                 performing updates.
     public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true) {
         core.apply(
             snapshot,
@@ -40,7 +42,9 @@ open class CocoaCollectionViewDiffableDataSource<SectionIdentifierType: Hashable
             animatingDifferences: animatingDifferences,
             performUpdates: { collectionView, changeset, setSections in
                 collectionView.reload(using: changeset, setData: setSections)
-        })
+        },
+            completion: completion
+        )
     }
 
     /// Returns a new snapshot object of current state.
diff --git a/Sources/Internal/DiffableDataSourceCore.swift b/Sources/Internal/DiffableDataSourceCore.swift
index 11a0f1c..b45b757 100644
--- a/Sources/Internal/DiffableDataSourceCore.swift
+++ b/Sources/Internal/DiffableDataSourceCore.swift
@@ -13,7 +13,8 @@ final class DiffableDataSourceCore<SectionIdentifierType: Hashable, ItemIdentifi
         _ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>,
         view: View?,
         animatingDifferences: Bool,
-        performUpdates: @escaping (View, StagedChangeset<[Section]>, @escaping ([Section]) -> Void) -> Void
+        performUpdates: @escaping (View, StagedChangeset<[Section]>, @escaping ([Section]) -> Void) -> Void,
+        completion: (() -> Void)?
         ) {
         dispatcher.dispatch { [weak self] in
             guard let self = self else {
@@ -35,15 +36,18 @@ final class DiffableDataSourceCore<SectionIdentifierType: Hashable, ItemIdentifi
                 }
             }
 
+            CATransaction.begin()
+            CATransaction.setCompletionBlock(completion)
+
             if animatingDifferences {
                 performDiffingUpdates()
             }
             else {
-                CATransaction.begin()
                 CATransaction.setDisableActions(true)
                 performDiffingUpdates()
-                CATransaction.commit()
             }
+
+            CATransaction.commit()
         }
     }
 
diff --git a/Sources/UIKit/CollectionViewDiffableDataSource.swift b/Sources/UIKit/CollectionViewDiffableDataSource.swift
index d978f6d..789aedf 100644
--- a/Sources/UIKit/CollectionViewDiffableDataSource.swift
+++ b/Sources/UIKit/CollectionViewDiffableDataSource.swift
@@ -39,14 +39,18 @@ open class CollectionViewDiffableDataSource<SectionIdentifierType: Hashable, Ite
     ///   - snapshot: A snapshot object to be applied to data model.
     ///   - animatingDifferences: A Boolean value indicating whether to update with
     ///                           diffing animation.
-    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true) {
+    ///   - completion: An optional completion block which is called when the complete
+    ///                 performing updates.
+    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true, completion: (() -> Void)? = nil) {
         core.apply(
             snapshot,
             view: collectionView,
             animatingDifferences: animatingDifferences,
             performUpdates: { collectionView, changeset, setSections in
                 collectionView.reload(using: changeset, setData: setSections)
-        })
+        },
+            completion: completion
+        )
     }
 
     /// Returns a new snapshot object of current state.
diff --git a/Sources/UIKit/TableViewDiffableDataSource.swift b/Sources/UIKit/TableViewDiffableDataSource.swift
index 67ce2d7..fcc7a1f 100644
--- a/Sources/UIKit/TableViewDiffableDataSource.swift
+++ b/Sources/UIKit/TableViewDiffableDataSource.swift
@@ -36,14 +36,18 @@ open class TableViewDiffableDataSource<SectionIdentifierType: Hashable, ItemIden
     ///   - snapshot: A snapshot object to be applied to data model.
     ///   - animatingDifferences: A Boolean value indicating whether to update with
     ///                           diffing animation.
-    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true) {
+    ///   - completion: An optional completion block which is called when the complete
+    ///                 performing updates.
+    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true, completion: (() -> Void)? = nil) {
         core.apply(
             snapshot,
             view: tableView,
             animatingDifferences: animatingDifferences,
             performUpdates: { tableView, changeset, setSections in
                 tableView.reload(using: changeset, with: self.defaultRowAnimation, setData: setSections)
-        })
+        },
+            completion: completion
+        )
     }
 
     /// Returns a new snapshot object of current state.