From 3386b29ed04e863caee438fe02800e3d187f2c16 Mon Sep 17 00:00:00 2001
From: Joe Turner <joturner@google.com>
Date: Wed, 30 Sep 2015 07:29:55 -0700
Subject: [PATCH] Fixes an issue where NITableViewModel and
 NICollectionViewModel use immutable section arrays, which can get into the
 mutable subclasses. So now we have a method where the sections array is
 created, al lowing us to make sure it is mutable when needed.

---
 src/collections/src/NICollectionViewModel+Private.h |  1 +
 src/collections/src/NICollectionViewModel.m         | 10 +++++++---
 src/collections/src/NIMutableCollectionViewModel.m  |  8 ++++++++
 src/models/src/NIMutableTableViewModel.m            | 12 ++++++++++--
 src/models/src/NITableViewModel+Private.h           |  1 +
 src/models/src/NITableViewModel.m                   | 10 +++++++---
 6 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/collections/src/NICollectionViewModel+Private.h b/src/collections/src/NICollectionViewModel+Private.h
index 6e104834c..c094c9f2e 100644
--- a/src/collections/src/NICollectionViewModel+Private.h
+++ b/src/collections/src/NICollectionViewModel+Private.h
@@ -25,6 +25,7 @@
 - (void)_resetCompiledData;
 - (void)_compileDataWithListArray:(NSArray *)listArray;
 - (void)_compileDataWithSectionedArray:(NSArray *)sectionedArray;
+- (void)_setSectionsWithArray:(NSArray *)sectionsArray;
 
 @end
 
diff --git a/src/collections/src/NICollectionViewModel.m b/src/collections/src/NICollectionViewModel.m
index 4d5dac5fb..948a70bbb 100644
--- a/src/collections/src/NICollectionViewModel.m
+++ b/src/collections/src/NICollectionViewModel.m
@@ -59,7 +59,7 @@ - (id)init {
 
 
 - (void)_resetCompiledData {
-  self.sections = nil;
+  [self _setSectionsWithArray:nil];
   self.sectionIndexTitles = nil;
   self.sectionPrefixToSectionIndex = nil;
 }
@@ -70,7 +70,7 @@ - (void)_compileDataWithListArray:(NSArray *)listArray {
   if (nil != listArray) {
     NICollectionViewModelSection* section = [NICollectionViewModelSection section];
     section.rows = listArray;
-    self.sections = [NSArray arrayWithObject:section];
+    [self _setSectionsWithArray:@[ section ]];
   }
 }
 
@@ -132,7 +132,11 @@ - (void)_compileDataWithSectionedArray:(NSArray *)sectionedArray {
   currentSectionRows = nil;
 
   // Update the compiled information for this data source.
-  self.sections = sections;
+  [self _setSectionsWithArray:sections];
+}
+
+- (void)_setSectionsWithArray:(NSArray *)sectionsArray {
+  self.sections = sectionsArray;
 }
 
 #pragma mark - UICollectionViewDataSource
diff --git a/src/collections/src/NIMutableCollectionViewModel.m b/src/collections/src/NIMutableCollectionViewModel.m
index 7676d05f6..962d1ba54 100644
--- a/src/collections/src/NIMutableCollectionViewModel.m
+++ b/src/collections/src/NIMutableCollectionViewModel.m
@@ -121,6 +121,14 @@ - (NICollectionViewModelSection *)_insertSectionAtIndex:(NSUInteger)index {
   return section;
 }
 
+- (void)_setSectionsWithArray:(NSArray *)sectionsArray {
+  if ([sectionsArray isKindOfClass:[NSMutableArray class]]) {
+    self.sections = (NSMutableArray *)sectionsArray;
+  } else {
+    self.sections = [sectionsArray mutableCopy];
+  }
+}
+
 @end
 
 
diff --git a/src/models/src/NIMutableTableViewModel.m b/src/models/src/NIMutableTableViewModel.m
index cb4a04283..2a4143366 100644
--- a/src/models/src/NIMutableTableViewModel.m
+++ b/src/models/src/NIMutableTableViewModel.m
@@ -99,7 +99,7 @@ - (void)updateSectionIndex {
 
 - (NITableViewModelSection *)_appendSection {
   if (nil == self.sections) {
-    self.sections = [NSMutableArray array];
+    [self _setSectionsWithArray:[NSMutableArray array]];
   }
   NITableViewModelSection* section = nil;
   section = [[NITableViewModelSection alloc] init];
@@ -110,7 +110,7 @@ - (NITableViewModelSection *)_appendSection {
 
 - (NITableViewModelSection *)_insertSectionAtIndex:(NSUInteger)index {
   if (nil == self.sections) {
-    self.sections = [NSMutableArray array];
+    [self _setSectionsWithArray:[NSMutableArray array]];
   }
   NITableViewModelSection* section = nil;
   section = [[NITableViewModelSection alloc] init];
@@ -120,6 +120,14 @@ - (NITableViewModelSection *)_insertSectionAtIndex:(NSUInteger)index {
   return section;
 }
 
+- (void)_setSectionsWithArray:(NSArray *)sectionsArray {
+  if ([sectionsArray isKindOfClass:[NSMutableArray class]]) {
+    self.sections = (NSMutableArray *)sectionsArray;
+  } else {
+    self.sections = [sectionsArray mutableCopy];
+  }
+}
+
 #pragma mark - UITableViewDataSource
 
 
diff --git a/src/models/src/NITableViewModel+Private.h b/src/models/src/NITableViewModel+Private.h
index 854ccf4cf..5215759de 100644
--- a/src/models/src/NITableViewModel+Private.h
+++ b/src/models/src/NITableViewModel+Private.h
@@ -26,6 +26,7 @@
 - (void)_compileDataWithListArray:(NSArray *)listArray;
 - (void)_compileDataWithSectionedArray:(NSArray *)sectionedArray;
 - (void)_compileSectionIndex;
+- (void)_setSectionsWithArray:(NSArray *)sectionsArray;
 
 @end
 
diff --git a/src/models/src/NITableViewModel.m b/src/models/src/NITableViewModel.m
index b553d9d36..cedc58fe0 100644
--- a/src/models/src/NITableViewModel.m
+++ b/src/models/src/NITableViewModel.m
@@ -65,7 +65,7 @@ - (id)init {
 
 
 - (void)_resetCompiledData {
-  self.sections = nil;
+  [self _setSectionsWithArray:nil];
   self.sectionIndexTitles = nil;
   self.sectionPrefixToSectionIndex = nil;
 }
@@ -76,7 +76,7 @@ - (void)_compileDataWithListArray:(NSArray *)listArray {
   if (nil != listArray) {
     NITableViewModelSection* section = [NITableViewModelSection section];
     section.rows = [listArray mutableCopy];
-    self.sections = [NSMutableArray arrayWithObject:section];
+    [self _setSectionsWithArray:@[ section ]];
   }
 }
 
@@ -138,7 +138,7 @@ - (void)_compileDataWithSectionedArray:(NSArray *)sectionedArray {
   currentSectionRows = nil;
 
   // Update the compiled information for this data source.
-  self.sections = sections;
+  [self _setSectionsWithArray:sections];
 }
 
 - (void)_compileSectionIndex {
@@ -222,6 +222,10 @@ - (void)_compileSectionIndex {
   self.sectionPrefixToSectionIndex = sectionPrefixToSectionIndex;
 }
 
+- (void)_setSectionsWithArray:(NSArray *)sectionsArray {
+  self.sections = sectionsArray;
+}
+
 #pragma mark - UITableViewDataSource