2626#import " ASCollectionViewLayoutFacilitatorProtocol.h"
2727#import " ASSectionContext.h"
2828#import " ASCollectionView+Undeprecated.h"
29+ #import " _ASHierarchyChangeSet.h"
2930
3031/* *
3132 * A macro to get self.collectionNode and assign it to a local variable, or return
@@ -196,7 +197,17 @@ @interface ASCollectionView () <ASRangeControllerDataSource, ASRangeControllerDe
196197 * (0 sections) we always check at least once after each update (initial reload is the first update.)
197198 */
198199 BOOL _hasEverCheckedForBatchFetchingDueToUpdate;
199-
200+
201+ /* *
202+ * The change set that we're currently building, if any.
203+ */
204+ _ASHierarchyChangeSet *_changeSet;
205+
206+ /* *
207+ * Counter used to keep track of nested batch updates.
208+ */
209+ NSInteger _batchUpdateCount;
210+
200211 struct {
201212 unsigned int scrollViewDidScroll:1 ;
202213 unsigned int scrollViewWillBeginDragging:1 ;
@@ -352,6 +363,8 @@ - (instancetype)_initWithFrame:(CGRect)frame collectionViewLayout:(UICollectionV
352363- (void )dealloc
353364{
354365 ASDisplayNodeAssertMainThread ();
366+ ASDisplayNodeCAssert (_batchUpdateCount == 0 , @" ASCollectionView deallocated in the middle of a batch update." );
367+
355368 // Sometimes the UIKit classes can call back to their delegate even during deallocation, due to animation completion blocks etc.
356369 _isDeallocating = YES ;
357370 [self setAsyncDelegate: nil ];
@@ -402,6 +415,12 @@ - (void)relayoutItems
402415- (void )waitUntilAllUpdatesAreCommitted
403416{
404417 ASDisplayNodeAssertMainThread ();
418+ if (_batchUpdateCount > 0 ) {
419+ // This assertion will be enabled soon.
420+ // ASDisplayNodeFailAssert(@"Should not call %@ during batch update", NSStringFromSelector(_cmd));
421+ return ;
422+ }
423+
405424 [_dataController waitUntilAllUpdatesAreCommitted ];
406425}
407426
@@ -761,15 +780,43 @@ - (ASDataController *)dataController
761780 return _dataController;
762781}
763782
764- - (void )performBatchAnimated : ( BOOL ) animated updates : ( void (^)()) updates completion : ( void (^)( BOOL )) completion
783+ - (void )beginUpdates
765784{
766785 ASDisplayNodeAssertMainThread ();
786+ // _changeSet must be available during batch update
787+ ASDisplayNodeAssertTrue ((_batchUpdateCount > 0 ) == (_changeSet != nil ));
767788
768- [_dataController beginUpdates ];
789+ if (_batchUpdateCount == 0 ) {
790+ _changeSet = [[_ASHierarchyChangeSet alloc ] initWithOldData: [_dataController itemCountsFromDataSource ]];
791+ }
792+ _batchUpdateCount++;
793+ }
794+
795+ - (void )endUpdatesAnimated : (BOOL )animated completion : (nullable void (^)(BOOL ))completion
796+ {
797+ ASDisplayNodeAssertMainThread ();
798+ ASDisplayNodeAssertNotNil (_changeSet, @" _changeSet must be available when batch update ends" );
799+
800+ _batchUpdateCount--;
801+ // Prevent calling endUpdatesAnimated:completion: in an unbalanced way
802+ NSAssert (_batchUpdateCount >= 0 , @" endUpdatesAnimated:completion: called without having a balanced beginUpdates call" );
803+
804+ [_changeSet addCompletionHandler: completion];
805+
806+ if (_batchUpdateCount == 0 ) {
807+ [_dataController updateWithChangeSet: _changeSet animated: animated];
808+ _changeSet = nil ;
809+ }
810+ }
811+
812+ - (void )performBatchAnimated : (BOOL )animated updates : (void (^)())updates completion : (void (^)(BOOL ))completion
813+ {
814+ ASDisplayNodeAssertMainThread ();
815+ [self beginUpdates ];
769816 if (updates) {
770817 updates ();
771818 }
772- [_dataController endUpdatesAnimated: animated completion: completion];
819+ [self endUpdatesAnimated: animated completion: completion];
773820}
774821
775822- (void )performBatchUpdates : (void (^)())updates completion : (void (^)(BOOL ))completion
@@ -789,27 +836,35 @@ - (void)insertSections:(NSIndexSet *)sections
789836{
790837 ASDisplayNodeAssertMainThread ();
791838 if (sections.count == 0 ) { return ; }
792- [_dataController insertSections: sections withAnimationOptions: kASCollectionViewAnimationNone ];
839+ [self performBatchUpdates: ^{
840+ [_changeSet insertSections: sections animationOptions: kASCollectionViewAnimationNone ];
841+ } completion: nil ];
793842}
794843
795844- (void )deleteSections : (NSIndexSet *)sections
796845{
797846 ASDisplayNodeAssertMainThread ();
798847 if (sections.count == 0 ) { return ; }
799- [_dataController deleteSections: sections withAnimationOptions: kASCollectionViewAnimationNone ];
848+ [self performBatchUpdates: ^{
849+ [_changeSet deleteSections: sections animationOptions: kASCollectionViewAnimationNone ];
850+ } completion: nil ];
800851}
801852
802853- (void )reloadSections : (NSIndexSet *)sections
803854{
804855 ASDisplayNodeAssertMainThread ();
805856 if (sections.count == 0 ) { return ; }
806- [_dataController reloadSections: sections withAnimationOptions: kASCollectionViewAnimationNone ];
857+ [self performBatchUpdates: ^{
858+ [_changeSet reloadSections: sections animationOptions: kASCollectionViewAnimationNone ];
859+ } completion: nil ];
807860}
808861
809862- (void )moveSection : (NSInteger )section toSection : (NSInteger )newSection
810863{
811864 ASDisplayNodeAssertMainThread ();
812- [_dataController moveSection: section toSection: newSection withAnimationOptions: kASCollectionViewAnimationNone ];
865+ [self performBatchUpdates: ^{
866+ [_changeSet moveSection: section toSection: newSection animationOptions: kASCollectionViewAnimationNone ];
867+ } completion: nil ];
813868}
814869
815870- (id <ASSectionContext>)contextForSection : (NSInteger )section
@@ -822,27 +877,35 @@ - (void)insertItemsAtIndexPaths:(NSArray *)indexPaths
822877{
823878 ASDisplayNodeAssertMainThread ();
824879 if (indexPaths.count == 0 ) { return ; }
825- [_dataController insertRowsAtIndexPaths: indexPaths withAnimationOptions: kASCollectionViewAnimationNone ];
880+ [self performBatchUpdates: ^{
881+ [_changeSet insertItems: indexPaths animationOptions: kASCollectionViewAnimationNone ];
882+ } completion: nil ];
826883}
827884
828885- (void )deleteItemsAtIndexPaths : (NSArray *)indexPaths
829886{
830887 ASDisplayNodeAssertMainThread ();
831888 if (indexPaths.count == 0 ) { return ; }
832- [_dataController deleteRowsAtIndexPaths: indexPaths withAnimationOptions: kASCollectionViewAnimationNone ];
889+ [self performBatchUpdates: ^{
890+ [_changeSet deleteItems: indexPaths animationOptions: kASCollectionViewAnimationNone ];
891+ } completion: nil ];
833892}
834893
835894- (void )reloadItemsAtIndexPaths : (NSArray *)indexPaths
836895{
837896 ASDisplayNodeAssertMainThread ();
838897 if (indexPaths.count == 0 ) { return ; }
839- [_dataController reloadRowsAtIndexPaths: indexPaths withAnimationOptions: kASCollectionViewAnimationNone ];
898+ [self performBatchUpdates: ^{
899+ [_changeSet reloadItems: indexPaths animationOptions: kASCollectionViewAnimationNone ];
900+ } completion: nil ];
840901}
841902
842903- (void )moveItemAtIndexPath : (NSIndexPath *)indexPath toIndexPath : (NSIndexPath *)newIndexPath
843904{
844905 ASDisplayNodeAssertMainThread ();
845- [_dataController moveRowAtIndexPath: indexPath toIndexPath: newIndexPath withAnimationOptions: kASCollectionViewAnimationNone ];
906+ [self performBatchUpdates: ^{
907+ [_changeSet moveItemAtIndexPath: indexPath toIndexPath: newIndexPath animationOptions: kASCollectionViewAnimationNone ];
908+ } completion: nil ];
846909}
847910
848911#pragma mark -
0 commit comments