From dc23aca30da77e0e4ea0a34ad0209f52d60d4e04 Mon Sep 17 00:00:00 2001
From: Huy Nguyen <allforone1511@gmail.com>
Date: Mon, 27 Feb 2017 17:07:23 +0000
Subject: [PATCH] Fix crashes in ASElementMap and ASMutableElementMap (#3090)

* Fix crashes in ASElementMap and ASMutableElementMap

* testInitialRangeBounds to have a 10% maximum difference
---
 Source/ASTableView.mm                          |  2 +-
 Source/Private/ASElementMap.m                  |  6 +++---
 Source/Private/ASMultidimensionalArrayUtils.mm |  1 +
 Source/Private/ASMutableElementMap.m           | 11 ++++++++---
 Tests/ASCollectionViewTests.mm                 |  2 +-
 5 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/Source/ASTableView.mm b/Source/ASTableView.mm
index 0778832768..c8263aee93 100644
--- a/Source/ASTableView.mm
+++ b/Source/ASTableView.mm
@@ -548,7 +548,7 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
 {
   // If this is a section index path, we don't currently have a method
   // to do a mapping.
-  if (indexPath.row == NSNotFound) {
+  if (indexPath == nil || indexPath.row == NSNotFound) {
     return indexPath;
   } else {
     NSIndexPath *viewIndexPath = [_dataController.visibleMap convertIndexPath:indexPath fromMap:_dataController.pendingMap];
diff --git a/Source/Private/ASElementMap.m b/Source/Private/ASElementMap.m
index c5664683ea..3fed5e9f48 100644
--- a/Source/Private/ASElementMap.m
+++ b/Source/Private/ASElementMap.m
@@ -54,8 +54,8 @@
       }
       s++;
     }
-    for (NSDictionary *supsOfKind in [_supplementaryElements objectEnumerator]) {
-      [supsOfKind enumerateKeysAndObjectsUsingBlock:^(NSIndexPath *_Nonnull indexPath, ASCollectionElement * _Nonnull element, BOOL * _Nonnull stop) {
+    for (NSDictionary *supplementariesForKind in [_supplementaryElements objectEnumerator]) {
+      [supplementariesForKind enumerateKeysAndObjectsUsingBlock:^(NSIndexPath *_Nonnull indexPath, ASCollectionElement * _Nonnull element, BOOL * _Nonnull stop) {
         [_elementToIndexPathMap setObject:indexPath forKey:element];
       }];
     }
@@ -90,7 +90,7 @@
 
 - (nullable ASCollectionElement *)elementForItemAtIndexPath:(NSIndexPath *)indexPath
 {
-  return ASGetElementInTwoDimensionalArray(_sectionsOfItems, indexPath);
+  return (indexPath != nil) ? ASGetElementInTwoDimensionalArray(_sectionsOfItems, indexPath) : nil;
 }
 
 - (nullable ASCollectionElement *)supplementaryElementOfKind:(NSString *)supplementaryElementKind atIndexPath:(NSIndexPath *)indexPath
diff --git a/Source/Private/ASMultidimensionalArrayUtils.mm b/Source/Private/ASMultidimensionalArrayUtils.mm
index 985b4cd5cb..bfa6c65eb4 100644
--- a/Source/Private/ASMultidimensionalArrayUtils.mm
+++ b/Source/Private/ASMultidimensionalArrayUtils.mm
@@ -226,6 +226,7 @@ NSArray *ASIndexPathsForMultidimensionalArray(NSArray *multidimensionalArray)
 
 id ASGetElementInTwoDimensionalArray(NSArray *array, NSIndexPath *indexPath)
 {
+  ASDisplayNodeCAssertNotNil(indexPath, @"Expected non-nil index path");
   ASDisplayNodeCAssert(indexPath.length == 2, @"Expected index path of length 2. Index path: %@", indexPath);
   NSInteger section = indexPath.section;
   if (array.count <= section) {
diff --git a/Source/Private/ASMutableElementMap.m b/Source/Private/ASMutableElementMap.m
index 9a7de008d9..af087218a1 100644
--- a/Source/Private/ASMutableElementMap.m
+++ b/Source/Private/ASMutableElementMap.m
@@ -72,8 +72,8 @@ typedef NSMutableDictionary<NSString *, NSMutableDictionary<NSIndexPath *, ASCol
     [_sectionsOfItems removeObjectsAtIndexes:sections];
   } else {
     // Supplementaries
-    NSMutableDictionary *supsForKind = _supplementaryElements[kind];
-    [supsForKind removeObjectsForKeys:[sections as_filterIndexPathsBySection:supsForKind]];
+    NSMutableDictionary *supplementariesForKind = _supplementaryElements[kind];
+    [supplementariesForKind removeObjectsForKeys:[sections as_filterIndexPathsBySection:supplementariesForKind]];
   }
 }
 
@@ -90,7 +90,12 @@ typedef NSMutableDictionary<NSString *, NSMutableDictionary<NSIndexPath *, ASCol
   if (kind == nil) {
     [_sectionsOfItems[indexPath.section] insertObject:element atIndex:indexPath.item];
   } else {
-    _supplementaryElements[kind][indexPath] = element;
+    NSMutableDictionary *supplementariesForKind = _supplementaryElements[kind];
+    if (supplementariesForKind == nil) {
+      supplementariesForKind = [NSMutableDictionary dictionary];
+      _supplementaryElements[kind] = supplementariesForKind;
+    }
+    supplementariesForKind[indexPath] = element;
   }
 }
 
diff --git a/Tests/ASCollectionViewTests.mm b/Tests/ASCollectionViewTests.mm
index e30734d54a..af96805f89 100644
--- a/Tests/ASCollectionViewTests.mm
+++ b/Tests/ASCollectionViewTests.mm
@@ -1031,7 +1031,7 @@
     r;
   });
   CGFloat expectedHeight = cn.bounds.size.height * 3;
-  XCTAssertEqualWithAccuracy(CGRectGetHeight(preloadBounds), expectedHeight, 100);
+  XCTAssertEqualWithAccuracy(CGRectGetHeight(preloadBounds), expectedHeight, expectedHeight * 0.1);
   XCTAssertEqual([[cn valueForKeyPath:@"rangeController.currentRangeMode"] integerValue], ASLayoutRangeModeMinimum, @"Expected range mode to be minimum before scrolling begins.");
 }