From aa2ae87c812a7c15828f2040f3d032b94b2e42a0 Mon Sep 17 00:00:00 2001
From: Michael Schneider <mischneider1@gmail.com>
Date: Wed, 2 Mar 2016 19:20:58 -0800
Subject: [PATCH] Add -waitUntilAllUpdatesAreCommitted to ASTableView and
 ASCollectionView

The API allows consumer of ASTableView or ASCollectionCiew to block execution of the main thread until all section and row updates are committed.
---
 AsyncDisplayKit/ASCollectionView.h          |  5 +++++
 AsyncDisplayKit/ASCollectionView.mm         |  6 ++++++
 AsyncDisplayKit/ASTableView.h               |  5 +++++
 AsyncDisplayKit/ASTableView.mm              |  6 ++++++
 AsyncDisplayKit/Details/ASDataController.h  |  2 ++
 AsyncDisplayKit/Details/ASDataController.mm | 17 +++++++++++++++++
 6 files changed, 41 insertions(+)

diff --git a/AsyncDisplayKit/ASCollectionView.h b/AsyncDisplayKit/ASCollectionView.h
index f1475156e4..0bc7959a95 100644
--- a/AsyncDisplayKit/ASCollectionView.h
+++ b/AsyncDisplayKit/ASCollectionView.h
@@ -165,6 +165,11 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)reloadDataImmediately;
 
+/**
+ *  Blocks execution of the main thread until all section and row updates are committed. This method must be called from the main thread.
+ */
+- (void)waitUntilAllUpdatesAreCommitted;
+
 /**
  * Registers the given kind of supplementary node for use in creating node-backed supplementary views.
  *
diff --git a/AsyncDisplayKit/ASCollectionView.mm b/AsyncDisplayKit/ASCollectionView.mm
index 69d5e75da0..86dec91d75 100644
--- a/AsyncDisplayKit/ASCollectionView.mm
+++ b/AsyncDisplayKit/ASCollectionView.mm
@@ -268,6 +268,12 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
   [super reloadData];
 }
 
+- (void)waitUntilAllUpdatesAreCommitted
+{
+  ASDisplayNodeAssertMainThread();
+  [_dataController waitUntilAllUpdatesAreCommitted];
+}
+
 - (void)setDataSource:(id<UICollectionViewDataSource>)dataSource
 {
   // UIKit can internally generate a call to this method upon changing the asyncDataSource; only assert for non-nil.
diff --git a/AsyncDisplayKit/ASTableView.h b/AsyncDisplayKit/ASTableView.h
index c2c50cab4a..28b825403a 100644
--- a/AsyncDisplayKit/ASTableView.h
+++ b/AsyncDisplayKit/ASTableView.h
@@ -163,6 +163,11 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)endUpdatesAnimated:(BOOL)animated completion:(void (^ _Nullable)(BOOL completed))completion;
 
+/**
+ *  Blocks execution of the main thread until all section and row updates are committed. This method must be called from the main thread.
+ */
+- (void)waitUntilAllUpdatesAreCommitted;
+
 /**
  * Inserts one or more sections, with an option to animate the insertion.
  *
diff --git a/AsyncDisplayKit/ASTableView.mm b/AsyncDisplayKit/ASTableView.mm
index cc1c6cf686..31647b4b21 100644
--- a/AsyncDisplayKit/ASTableView.mm
+++ b/AsyncDisplayKit/ASTableView.mm
@@ -394,6 +394,12 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
   [_dataController endUpdatesAnimated:animated completion:completion];
 }
 
+- (void)waitUntilAllUpdatesAreCommitted
+{
+  ASDisplayNodeAssertMainThread();
+  [_dataController waitUntilAllUpdatesAreCommitted];
+}
+
 - (void)layoutSubviews
 {
   if (_nodesConstrainedWidth != self.bounds.size.width) {
diff --git a/AsyncDisplayKit/Details/ASDataController.h b/AsyncDisplayKit/Details/ASDataController.h
index 6648275d9b..66a03abbed 100644
--- a/AsyncDisplayKit/Details/ASDataController.h
+++ b/AsyncDisplayKit/Details/ASDataController.h
@@ -181,6 +181,8 @@ FOUNDATION_EXPORT NSString * const ASDataControllerRowNodeKind;
 
 - (void)reloadDataImmediatelyWithAnimationOptions:(ASDataControllerAnimationOptions)animationOptions;
 
+- (void)waitUntilAllUpdatesAreCommitted;
+
 /** @name Data Querying */
 
 - (NSUInteger)numberOfSections;
diff --git a/AsyncDisplayKit/Details/ASDataController.mm b/AsyncDisplayKit/Details/ASDataController.mm
index b37df27c1b..8c48e81732 100644
--- a/AsyncDisplayKit/Details/ASDataController.mm
+++ b/AsyncDisplayKit/Details/ASDataController.mm
@@ -462,6 +462,23 @@ static void *kASSizingQueueContext = &kASSizingQueueContext;
   }];
 }
 
+- (void)waitUntilAllUpdatesAreCommitted
+{
+  ASDisplayNodeAssertMainThread();
+  ASDisplayNodeAssert(_batchUpdateCounter == 0, @"Should not be called between beginUpdate or endUpdate");
+  
+  // This should never be called in a batch update, return immediately therefore
+  if (_batchUpdateCounter > 0) { return; }
+  
+  [_editingTransactionQueue waitUntilAllOperationsAreFinished];
+  
+  // Schedule block in main serial queue to wait until all operations are finished that are
+  // where scheduled while waiting for the _editingTransactionQueue to finish
+  [_mainSerialQueue performBlockOnMainThread:^{
+    ASDisplayNodeAssert(_editingTransactionQueue.operationCount == 0, @"No operation should be in the _editingTransactionQueue anymore");
+  }];
+}
+
 #pragma mark - Data Source Access (Calling _dataSource)
 
 /**