fix(database): prevent FUIArray crash on desynced child events

FirebaseDatabaseUI/FirebaseDatabaseUITests/FUIArrayTest.m

@@ -485,4 +485,72 @@ - (void)testRemovesAllElementsWhenInvalidated {
             self.firebaseArray.count);
 }
 
+#pragma mark - Invariant violations
+
+- (void)testRemovingUnknownKeyDoesNotCrash {
+  [self.observable populateWithCount:10];
+  self.snap.key = @"this-key-was-never-added";
+
+  XCTAssertNoThrow(
+    [self.observable sendEvent:FIRDataEventTypeChildRemoved
+                    withObject:self.snap
+                   previousKey:@"9"
+                         error:nil]);
+
+  XCTAssert(self.firebaseArray.count == 10,
+            @"expected count to stay 10 when removing an unknown key, got %ld",
+            self.firebaseArray.count);
+}
+
+- (void)testInsertingWithUnknownPreviousKeyDoesNotCrash {
+  [self.observable populateWithCount:10];
+  self.snap.key = @"new";
+
+  XCTAssertNoThrow(
+    [self.observable sendEvent:FIRDataEventTypeChildAdded
+                    withObject:self.snap
+                   previousKey:@"this-key-was-never-added"
+                         error:nil]);
+
+  XCTAssert(self.firebaseArray.count == 11,
+            @"expected count to become 11 after insert, got %ld",
+            self.firebaseArray.count);
+  XCTAssert([[self.firebaseArray snapshotAtIndex:10].key isEqualToString:@"new"],
+            @"expected the new snapshot to be appended at the end");
+}
+
+- (void)testChangingUnknownKeyDoesNotCrash {
+  [self.observable populateWithCount:10];
+  self.snap.key = @"this-key-was-never-added";
+
+  XCTAssertNoThrow(
+    [self.observable sendEvent:FIRDataEventTypeChildChanged
+                    withObject:self.snap
+                   previousKey:@"9"
+                         error:nil]);
+
+  XCTAssert(self.firebaseArray.count == 11,
+            @"expected count to become 11 after recovering change as insert, got %ld",
+            self.firebaseArray.count);
+  XCTAssert([[self.firebaseArray snapshotAtIndex:10].key isEqualToString:@"this-key-was-never-added"],
+            @"expected the recovered snapshot to be inserted at index 10");
+}
+
+- (void)testMovingUnknownKeyDoesNotCrash {
+  [self.observable populateWithCount:10];
+  self.snap.key = @"this-key-was-never-added";
+
+  XCTAssertNoThrow(
+    [self.observable sendEvent:FIRDataEventTypeChildMoved
+                    withObject:self.snap
+                   previousKey:@"3"
+                         error:nil]);
+
+  XCTAssert(self.firebaseArray.count == 11,
+            @"expected count to become 11 after recovering move as insert, got %ld",
+            self.firebaseArray.count);
+  XCTAssert([[self.firebaseArray snapshotAtIndex:4].key isEqualToString:@"this-key-was-never-added"],
+            @"expected the recovered snapshot to be inserted at index 4");
+}
+
 @end

FirebaseDatabaseUI/FirebaseDatabaseUITests/FUICollectionViewDataSourceTest.m

@@ -59,6 +59,7 @@ - (void)setUp {
   NSLog(@"%lu", (unsigned long)[self.collectionView numberOfItemsInSection:0]);
 
   [self.observable populateWithCount:10];
+  [self waitForDataSourceCount:10];
 }
 
 - (void)tearDown {
@@ -68,6 +69,13 @@ - (void)tearDown {
   [super tearDown];
 }
 
+- (void)waitForDataSourceCount:(NSUInteger)expected {
+  NSDate *timeout = [NSDate dateWithTimeIntervalSinceNow:1.0];
+  while (self.dataSource.count != expected && timeout.timeIntervalSinceNow > 0) {
+    [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:0.01]];
+  }
+}
+
 - (void)testItHasACount {
   NSUInteger count = self.dataSource.count;
   XCTAssert(count == 10, @"expected data source to have 10 elements after 10 insertions, but got %lu", count);
@@ -119,4 +127,19 @@ - (void)testItMovesCells {
             found %@ at index 4 instead", cell.accessibilityValue);
 }
 
+- (void)testUnbindDoesNotLeakDataSourceWithPendingItems {
+  __weak FUICollectionViewDataSource *weakDataSource = self.dataSource;
+
+  [self.dataSource unbind];
+  self.dataSource = nil;
+
+  NSDate *timeout = [NSDate dateWithTimeIntervalSinceNow:1.0];
+  while (weakDataSource != nil && timeout.timeIntervalSinceNow > 0) {
+    [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:0.01]];
+  }
+
+  XCTAssertNil(weakDataSource,
+               @"expected data source to be deallocated after unbind with pending items, found a retain cycle");
+}
+
 @end

FirebaseDatabaseUI/Sources/FUIArray.m

@@ -191,18 +191,9 @@ - (NSUInteger)indexForKey:(NSString *)key {
 - (void)insertSnapshot:(FIRDataSnapshot *)snap withPreviousChildKey:(NSString *)previous {
   NSUInteger index = 0;
   if (previous != nil) {
-    NSInteger previousChildIndex = (NSInteger)[self indexForKey:previous];
-
-    if (previousChildIndex == NSNotFound) {
-      NSString *reason = [NSString stringWithFormat:@"Attempted to insert snapshot with unknown"
-                          @" previousChildKey %@ into array: %@", previous, self.snapshots];
-      NSException *exception = [NSException exceptionWithName:NSInternalInconsistencyException
-                                                       reason:reason
-                                                     userInfo:nil];
-      @throw exception;
-    }
-
-    index = previousChildIndex + 1;
+    NSUInteger previousChildIndex = [self indexForKey:previous];
+    index = (previousChildIndex == NSNotFound) ? self.snapshots.count
+                                               : previousChildIndex + 1;
   }
 
   [self.snapshots insertObject:snap atIndex:index];
@@ -217,12 +208,7 @@ - (void)removeSnapshot:(FIRDataSnapshot *)snap withPreviousChildKey:(NSString *)
   NSUInteger index = [self indexForKey:snap.key];
 
   if (index == NSNotFound) {
-    NSString *reason = [NSString stringWithFormat:@"Attempted to remove snapshot with unknown"
-                        @" key %@ from array: %@", snap.key, self.snapshots];
-    NSException *exception = [NSException exceptionWithName:NSInternalInconsistencyException
-                                                     reason:reason
-                                                   userInfo:nil];
-    @throw exception;
+    return;
   }
 
   [self.snapshots removeObjectAtIndex:index];
@@ -237,12 +223,8 @@ - (void)changeSnapshot:(FIRDataSnapshot *)snap withPreviousChildKey:(NSString *)
   NSUInteger index = [self indexForKey:snap.key];
 
   if (index == NSNotFound) {
-    NSString *reason = [NSString stringWithFormat:@"Attempted to replace snapshot with unknown"
-                        @" key %@ in array: %@", snap.key, self.snapshots];
-    NSException *exception = [NSException exceptionWithName:NSInternalInconsistencyException
-                                                     reason:reason
-                                                   userInfo:nil];
-    @throw exception;
+    [self insertSnapshot:snap withPreviousChildKey:previous];
+    return;
   }
 
   [self.snapshots replaceObjectAtIndex:index withObject:snap];
@@ -257,12 +239,8 @@ - (void)moveSnapshot:(FIRDataSnapshot *)snap withPreviousChildKey:(NSString *)pr
   NSUInteger fromIndex = [self indexForKey:snap.key];
 
   if (fromIndex == NSNotFound) {
-    NSString *reason = [NSString stringWithFormat:@"Attempted to remove snapshot with unknown"
-                        @" key %@ from array: %@", snap.key, self.snapshots];
-    NSException *exception = [NSException exceptionWithName:NSInternalInconsistencyException
-                                                     reason:reason
-                                                   userInfo:nil];
-    @throw exception;
+    [self insertSnapshot:snap withPreviousChildKey:previous];
+    return;
   }
 
   [self.snapshots removeObjectAtIndex:fromIndex];

FirebaseDatabaseUI/Sources/FUICollectionViewDataSource.m

@@ -38,6 +38,12 @@ @interface FUICollectionViewDataSource () <FUICollectionDelegate>
 @property (strong, nonatomic, readonly) UICollectionViewCell *(^populateCellAtIndexPath)
   (UICollectionView *collectionView, NSIndexPath *indexPath, FIRDataSnapshot *object);
 
+@property (nonatomic, strong) NSMutableArray<void (^)(dispatch_block_t)> *pendingUpdates;
+
+@property (nonatomic, assign) BOOL isApplyingBatchUpdate;
+
+@property (nonatomic, strong) NSMutableArray<FIRDataSnapshot *> *displayedSnapshots;
+
 @end
 
 @implementation FUICollectionViewDataSource
@@ -55,6 +61,7 @@ - (instancetype)initWithCollection:(id<FUICollection>)collection
     _populateCellAtIndexPath = populateCell;
     _count = 0; // This is zero because RTDB arrays start out at zero
                 // and send initial items as a series of adds.
+    _displayedSnapshots = [NSMutableArray array];
   }
   return self;
 }
@@ -72,11 +79,11 @@ - (NSUInteger)count {
 }
 
 - (NSArray<FIRDataSnapshot *> *)items {
-  return self.collection.items;
+  return [self.displayedSnapshots copy];
 }
 
 - (FIRDataSnapshot *)snapshotAtIndex:(NSInteger)index {
-  return [self.collection snapshotAtIndex:index];
+  return self.displayedSnapshots[index];
 }
 
 - (void)bindToView:(UICollectionView *)view {
@@ -96,30 +103,71 @@ - (void)unbind {
 // performBatchUpdates: is used for single updates because of this radar:
 // https://openradar.appspot.com/26484150
 - (void)array:(FUIArray *)array didAddObject:(id)object atIndex:(NSUInteger)index {
-  [self.collectionView performBatchUpdates:^{
-    self.count = array.count;
-    [self.collectionView
-     insertItemsAtIndexPaths:@[ [NSIndexPath indexPathForItem:index inSection:0] ]];
-  } completion:^(BOOL finished) {}];
+  [self enqueueUpdate:^(dispatch_block_t done) {
+    if (self.collectionView == nil) {
+      done();
+      return;
+    }
+    [self.collectionView performBatchUpdates:^{
+      [self.displayedSnapshots insertObject:object atIndex:index];
+      self.count = self.displayedSnapshots.count;
+      [self.collectionView
+       insertItemsAtIndexPaths:@[ [NSIndexPath indexPathForItem:index inSection:0] ]];
+    } completion:^(BOOL finished) {
+      done();
+    }];
+  }];
 }
 
 - (void)array:(FUIArray *)array didChangeObject:(id)object atIndex:(NSUInteger)index {
-  [self.collectionView
-   reloadItemsAtIndexPaths:@[ [NSIndexPath indexPathForItem:index inSection:0] ]];
+  [self enqueueUpdate:^(dispatch_block_t done) {
+    if (self.collectionView == nil) {
+      done();
+      return;
+    }
+    [self.collectionView performBatchUpdates:^{
+      self.displayedSnapshots[index] = object;
+      [self.collectionView
+       reloadItemsAtIndexPaths:@[ [NSIndexPath indexPathForItem:index inSection:0] ]];
+    } completion:^(BOOL finished) {
+      done();
+    }];
+  }];
 }
 
 - (void)array:(FUIArray *)array didRemoveObject:(id)object atIndex:(NSUInteger)index {
-  [self.collectionView performBatchUpdates:^{
-    self.count = array.count;
-    [self.collectionView
-     deleteItemsAtIndexPaths:@[ [NSIndexPath indexPathForItem:index inSection:0] ]];
-  } completion:^(BOOL finished) {}];
+  [self enqueueUpdate:^(dispatch_block_t done) {
+    if (self.collectionView == nil) {
+      done();
+      return;
+    }
+    [self.collectionView performBatchUpdates:^{
+      [self.displayedSnapshots removeObjectAtIndex:index];
+      self.count = self.displayedSnapshots.count;
+      [self.collectionView
+       deleteItemsAtIndexPaths:@[ [NSIndexPath indexPathForItem:index inSection:0] ]];
+    } completion:^(BOOL finished) {
+      done();
+    }];
+  }];
 }
 
 - (void)array:(FUIArray *)array didMoveObject:(id)object
     fromIndex:(NSUInteger)fromIndex toIndex:(NSUInteger)toIndex {
-  [self.collectionView moveItemAtIndexPath:[NSIndexPath indexPathForItem:fromIndex inSection:0]
-                               toIndexPath:[NSIndexPath indexPathForItem:toIndex inSection:0]];
+  [self enqueueUpdate:^(dispatch_block_t done) {
+    if (self.collectionView == nil) {
+      done();
+      return;
+    }
+    [self.collectionView performBatchUpdates:^{
+      [self.displayedSnapshots removeObjectAtIndex:fromIndex];
+      [self.displayedSnapshots insertObject:object atIndex:toIndex];
+      [self.collectionView moveItemAtIndexPath:[NSIndexPath indexPathForItem:fromIndex inSection:0]
+                                   toIndexPath:[NSIndexPath indexPathForItem:toIndex inSection:0]];
+    } completion:^(BOOL finished) {
+      done();
+    }];
+  }];
 }
 
 - (void)array:(id<FUICollection>)array queryCancelledWithError:(NSError *)error {
@@ -128,11 +176,36 @@ - (void)array:(id<FUICollection>)array queryCancelledWithError:(NSError *)error
   }
 }
 
+- (void)enqueueUpdate:(void (^)(dispatch_block_t done))update {
+  if (self.pendingUpdates == nil) {
+    self.pendingUpdates = [NSMutableArray array];
+  }
+  [self.pendingUpdates addObject:update];
+  [self flushPendingUpdatesIfNeeded];
+}
+
+- (void)flushPendingUpdatesIfNeeded {
+  if (self.isApplyingBatchUpdate || self.pendingUpdates.count == 0) {
+    return;
+  }
+
+  void (^next)(dispatch_block_t) = self.pendingUpdates.firstObject;
+  [self.pendingUpdates removeObjectAtIndex:0];
+  self.isApplyingBatchUpdate = YES;
+
+  next(^{
+    self.isApplyingBatchUpdate = NO;
+    dispatch_async(dispatch_get_main_queue(), ^{
+      [self flushPendingUpdatesIfNeeded];
+    });
+  });
+}
+
 #pragma mark - UICollectionViewDataSource methods
 
 - (nonnull UICollectionViewCell *)collectionView:(nonnull UICollectionView *)collectionView
                           cellForItemAtIndexPath:(nonnull NSIndexPath *)indexPath {
-  FIRDataSnapshot *snap = [self.collection.items objectAtIndex:indexPath.item];
+  FIRDataSnapshot *snap = self.displayedSnapshots[indexPath.item];
 
   UICollectionViewCell *cell = self.populateCellAtIndexPath(collectionView, indexPath, snap);
 

samples/swift/FirebaseUI-demo-swift/Samples/Chat/ChatViewController.swift

@@ -64,7 +64,7 @@ class ChatViewController: UIViewController, UICollectionViewDelegateFlowLayout {
           self.collectionView.bind(to: self.query!) { (view, indexPath, snap) -> UICollectionViewCell in
             let cell = view.dequeueReusableCell(withReuseIdentifier: ChatViewController.reuseIdentifier,
                                                 for: indexPath) as! ChatCollectionViewCell
-            let chat = Chat(snapshot: snap)!
+            guard let chat = Chat(snapshot: snap) else { return cell }
             cell.populateCellWithChat(chat, user: self.user, maxWidth: self.view.frame.size.width)
             return cell
       }
@@ -184,7 +184,9 @@ class ChatViewController: UIViewController, UICollectionViewDelegateFlowLayout {
 
     let width = self.view.frame.size.width
     let blob = self.collectionViewDataSource.snapshot(at: indexPath.item)
-    let text = Chat(snapshot: blob)!.text
+    guard let text = Chat(snapshot: blob)?.text else {
+      return CGSize(width: width, height: 0)
+    }
 
     let rect = ChatCollectionViewCell.boundingRectForText(text, maxWidth: width)