1 files changed, 246 insertions, 232 deletions

diff --git a/Source/SPDataStorage.m b/Source/SPDataStorage.m
index 1b3d1cba..c2119032 100644
--- a/Source/SPDataStorage.m
+++ b/Source/SPDataStorage.m
@@ -31,45 +31,89 @@
 //  More info at <http://code.google.com/p/sequel-pro/>
 
 #import "SPDataStorage.h"
+#import "SPObjectAdditions.h"
+#import <SPMySQL/SPMySQLStreamingResultStore.h>
 
-@interface SPDataStorage (PrivateAPI)
+@interface SPDataStorage (Private_API)
 
-- (void) _ensureCapacityForAdditionalRowCount:(NSUInteger)numExtraRows;
-- (void) _increaseCapacity;
+- (void) _checkNewRow:(NSMutableArray *)aRow;
 
 @end
 
 @implementation SPDataStorage
 
-static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorage* self, NSUInteger numExtraRows) 
+static inline NSMutableArray* SPDataStorageGetEditedRow(NSPointerArray* rowStore, NSUInteger rowIndex)
 {
-	typedef void (*SPDSEnsureCapacityMethodPtr)(SPDataStorage*, SEL, NSUInteger);
-	static SPDSEnsureCapacityMethodPtr SPDSEnsureCapacity;
-	if (!SPDSEnsureCapacity) SPDSEnsureCapacity = (SPDSEnsureCapacityMethodPtr)[self methodForSelector:@selector(_ensureCapacityForAdditionalRowCount:)];
-	SPDSEnsureCapacity(self, @selector(_ensureCapacityForAdditionalRowCount:), numExtraRows);
+	typedef NSMutableArray* (*SPDSGetEditedRowMethodPtr)(NSPointerArray*, SEL, NSUInteger);
+	static SPDSGetEditedRowMethodPtr SPDSGetEditedRow;
+	if (!SPDSGetEditedRow) SPDSGetEditedRow = (SPDSGetEditedRowMethodPtr)[rowStore methodForSelector:@selector(pointerAtIndex:)];
+	return SPDSGetEditedRow(rowStore, @selector(pointerAtIndex:), rowIndex);
 }
 
+#pragma mark - Setting result store
+
+/**
+ * Set the underlying MySQL data storage.
+ * This will clear all edited rows and unloaded column tracking.
+ */
+- (void) setDataStorage:(SPMySQLStreamingResultStore *)newDataStorage updatingExisting:(BOOL)updateExistingStore
+{
+	NSUInteger i;
+	[editedRows release], editedRows = nil;
+	if (unloadedColumns) free(unloadedColumns), unloadedColumns = NULL;
+
+	if (dataStorage) {
+
+		// If the table is reloading data, link to the current data store for smoother loads
+		if (updateExistingStore) {
+			[newDataStorage replaceExistingResultStore:dataStorage];
+		}
+
+		[dataStorage release], dataStorage = nil;
+	}
+
+	dataStorage = [newDataStorage retain];
+	[dataStorage setDelegate:self];
+
+	numberOfColumns = [dataStorage numberOfFields];
+	editedRows = [NSPointerArray new];
+	if ([dataStorage dataDownloaded]) {
+		[self resultStoreDidFinishLoadingData:dataStorage];
+	}
+
+	unloadedColumns = malloc(numberOfColumns * sizeof(BOOL));
+	for (i = 0; i < numberOfColumns; i++) {
+		unloadedColumns[i] = NO;
+	}
+}
+
+
 #pragma mark -
 #pragma mark Retrieving rows and cells
 
 /**
  * Return a mutable array containing the data for a specified row.
  */
-- (NSMutableArray *) rowContentsAtIndex:(NSUInteger)index
+- (NSMutableArray *) rowContentsAtIndex:(NSUInteger)anIndex
 {
 
-	// Throw an exception if the index is out of bounds
-	if (index >= numRows) [NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)index, (unsigned long long)numRows];
+	// If an edited row exists for the supplied index, return it
+	NSMutableArray *editedRow = SPDataStorageGetEditedRow(editedRows, anIndex);
+	if (editedRow != NULL) {
+		return editedRow;
+	}
 
-	// Construct the NSMutableArray
-	NSMutableArray *rowArray = [NSMutableArray arrayWithCapacity:numColumns];
-	id *row = dataStorage[index];
-	NSUInteger i;
-	for (i = 0; i < numColumns; i++) {
-		CFArrayAppendValue((CFMutableArrayRef)rowArray, row[i]);
+	// Otherwise, prepare to return the underlying storage row
+	NSMutableArray *dataArray = SPMySQLResultStoreGetRow(dataStorage, anIndex);
+
+	// Modify unloaded cells as appropriate
+	for (NSUInteger i = 0; i < numberOfColumns; i++) {
+		if (unloadedColumns[i]) {
+			CFArraySetValueAtIndex((CFMutableArrayRef)dataArray, i, [SPNotLoaded notLoaded]);
+		}
 	}
 
-	return rowArray;
+	return dataArray;
 }
 
 /**
@@ -78,11 +122,78 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
 - (id) cellDataAtRow:(NSUInteger)rowIndex column:(NSUInteger)columnIndex
 {
 
-	// Throw an exception if the row or column index is out of bounds
-	if (rowIndex >= numRows || columnIndex >= numColumns) [NSException raise:NSRangeException format:@"Requested storage index (row %llu, col %llu) beyond bounds (%llu, %llu)", (unsigned long long)rowIndex, (unsigned long long)columnIndex, (unsigned long long)numRows, (unsigned long long)numColumns];
+	// If an edited row exists at the supplied index, return it
+	NSMutableArray *editedRow = SPDataStorageGetEditedRow(editedRows, rowIndex);
+	if (editedRow != NULL) {
+		return CFArrayGetValueAtIndex((CFArrayRef)editedRow, columnIndex);
+	}
+
+	// Throw an exception if the column index is out of bounds
+	if (columnIndex >= numberOfColumns) {
+		[NSException raise:NSRangeException format:@"Requested storage column (col %llu) beyond bounds (%llu)", (unsigned long long)columnIndex, (unsigned long long)numberOfColumns];
+	}
+
+	// If the specified column is not loaded, return a SPNotLoaded reference
+	if (unloadedColumns[columnIndex]) {
+		return [SPNotLoaded notLoaded];
+	}
 
 	// Return the content
-	return dataStorage[rowIndex][columnIndex];
+	return SPMySQLResultStoreObjectAtRowAndColumn(dataStorage, rowIndex, columnIndex);
+}
+
+/**
+ * Return a preview of the data at a specified row and column index, limited
+ * to approximately the supplied length.
+ */
+- (id) cellPreviewAtRow:(NSUInteger)rowIndex column:(NSUInteger)columnIndex previewLength:(NSUInteger)previewLength
+{
+
+	// If an edited row exists at the supplied index, return it
+	NSMutableArray *editedRow = SPDataStorageGetEditedRow(editedRows, rowIndex);
+	if (editedRow != NULL) {
+		id anObject = CFArrayGetValueAtIndex((CFArrayRef)editedRow, columnIndex);
+		if ([anObject isKindOfClass:[NSString class]] && [(NSString *)anObject length] > 150) {
+			return ([NSString stringWithFormat:@"%@...", [anObject substringToIndex:147]]);
+		}
+		return anObject;
+	}
+
+	// Throw an exception if the column index is out of bounds
+	if (columnIndex >= numberOfColumns) {
+		[NSException raise:NSRangeException format:@"Requested storage column (col %llu) beyond bounds (%llu)", (unsigned long long)columnIndex, (unsigned long long)numberOfColumns];
+	}
+
+	// If the specified column is not loaded, return a SPNotLoaded reference
+	if (unloadedColumns[columnIndex]) {
+		return [SPNotLoaded notLoaded];
+	}
+
+	// Return the content
+	return SPMySQLResultStorePreviewAtRowAndColumn(dataStorage, rowIndex, columnIndex, previewLength);
+}
+
+/**
+ * Returns whether the data at a specified row and column index is NULL or unloaded
+ */
+- (BOOL) cellIsNullOrUnloadedAtRow:(NSUInteger)rowIndex column:(NSUInteger)columnIndex
+{
+	// If an edited row exists at the supplied index, check it for a NULL.
+	NSMutableArray *editedRow = SPDataStorageGetEditedRow(editedRows, rowIndex);
+	if (editedRow != NULL) {
+		return [(id)CFArrayGetValueAtIndex((CFArrayRef)editedRow, columnIndex) isNSNull];
+	}
+
+	// Throw an exception if the column index is out of bounds
+	if (columnIndex >= numberOfColumns) {
+		[NSException raise:NSRangeException format:@"Requested storage column (col %llu) beyond bounds (%llu)", (unsigned long long)columnIndex, (unsigned long long)numberOfColumns];
+	}
+
+	if (unloadedColumns[columnIndex]) {
+		return YES;
+	}
+
+	return [dataStorage cellIsNullAtRow:rowIndex column:columnIndex];
 }
 
 #pragma mark -
@@ -90,38 +201,37 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
 
 /**
  * Implementation of the NSFastEnumeration protocol.
- * Note that this currently doesn't implement mutation guards.
+ * Note that rows are currently retrieved individually to avoid mutation and locking issues,
+ * although this could be improved on.
  */
 - (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state objects:(id *)stackbuf count:(NSUInteger)len
 {
 
 	// If the start index is out of bounds, return 0 to indicate end of results
-	if (state->state >= numRows) return 0;
-
-	// Determine how many objects to return - 128, len, or all items remaining
-	NSUInteger itemsToReturn = 128;
-	if (len < 128) itemsToReturn = len;
-	if (numRows - state->state < itemsToReturn) {
-		itemsToReturn = numRows - state->state;
-	}
-
-	// Construct the arrays to return
-	NSUInteger i, j;
-	NSMutableArray *rowArray;
-	id *row;
-	for (i = 0; i < itemsToReturn; i++) {
-		row = dataStorage[state->state + i];
-		rowArray = [NSMutableArray arrayWithCapacity:numColumns];
-		for (j = 0; j < numColumns; j++) {
-			CFArrayAppendValue((CFMutableArrayRef)rowArray, row[j]);
+	if (state->state >= SPMySQLResultStoreGetRowCount(dataStorage)) return 0;
+
+	// If an edited row exists for the supplied index, use that; otherwise use the underlying
+	// storage row
+	NSMutableArray *targetRow = SPDataStorageGetEditedRow(editedRows, state->state);
+	if (targetRow == NULL) {
+		targetRow = SPMySQLResultStoreGetRow(dataStorage, state->state);
+
+		// Modify unloaded cells as appropriate
+		for (NSUInteger i = 0; i < numberOfColumns; i++) {
+			if (unloadedColumns[i]) {
+				CFArraySetValueAtIndex((CFMutableArrayRef)targetRow, i, [SPNotLoaded notLoaded]);
+			}
 		}
-		stackbuf[i] = rowArray;
 	}
 
-	state->state += itemsToReturn;
+	// Add the item to the buffer and return the appropriate state
+	stackbuf[0] = targetRow;
+
+	state->state += 1;
 	state->itemsPtr = stackbuf;
-	state->mutationsPtr = (unsigned long *)&numRows;
-	return itemsToReturn;
+	state->mutationsPtr = (unsigned long *)self;
+
+	return 1;
 }
 
 #pragma mark -
@@ -132,29 +242,17 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
  * of objects.  Note that the supplied objects are retained as a reference
  * rather than copied.
  */
-- (void) addRowWithContents:(NSArray *)row
+- (void) addRowWithContents:(NSMutableArray *)aRow
 {
-	
-	// Ensure that sufficient capacity is available
-	SPDataStorageEnsureCapacityForAdditionalRowCount(self, 1);
-
-	// Add an empty row array to the data store
-	id *newRow = (id *)malloc(columnPointerByteSize);
-	dataStorage[numRows] = newRow;
-	numRows++;
-
-	// Copy over references to the array contents, and retain the objects
-	NSUInteger cellsCopied = 0;
-	for (id cellData in row) {
-		if (cellData) newRow[cellsCopied] = (id)CFRetain(cellData);
-		else newRow[cellsCopied] = nil;
-		if (++cellsCopied == numColumns) break;
-	}
 
-	// If an array shorter than the row width was added, pad with nils
-	if (cellsCopied < numColumns) {
-		for ( ; cellsCopied <= numColumns; cellsCopied++) newRow[cellsCopied] = nil;
-	}
+	// Verify the row is of the correct length
+	[self _checkNewRow:aRow];
+
+	// Add the new row to the editable store
+	[editedRows addPointer:aRow];
+
+	// Update the underlying store as well to keep counts correct
+	[dataStorage addDummyRow];
 }
 
 /**
@@ -162,109 +260,70 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
  * all later rows the next index.  Note that the supplied objects within the
  * array are retained as a reference rather than copied.
  */
-- (void) insertRowContents:(NSArray *)row atIndex:(NSUInteger)index
+- (void) insertRowContents:(NSMutableArray *)aRow atIndex:(NSUInteger)anIndex
 {
+	unsigned long long numberOfRows = SPMySQLResultStoreGetRowCount(dataStorage);
+
+	// Verify the row is of the correct length
+	[self _checkNewRow:aRow];
 
 	// Throw an exception if the index is out of bounds
-	if (index > numRows) [NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)index, (unsigned long long)numRows];
+	if (anIndex > numberOfRows) {
+		[NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)anIndex, numberOfRows];
+	}
 
 	// If "inserting" at the end of the array just add a row
-	if (index == numRows) return SPDataStorageAddRow(self, row);
-
-	// Ensure that sufficient capacity is available to hold all the rows
-	SPDataStorageEnsureCapacityForAdditionalRowCount(self, 1);
-
-	// Renumber the specified index, and all subsequent indices, to create a gap
-	for (NSUInteger j = numRows - 1; j >= index; j--) {
-		dataStorage[j + 1] = dataStorage[j];
+	if (anIndex == numberOfRows) {
+		return [self addRowWithContents:aRow];
 	}
 
-	// Add a new instantiated row array to the data store at the specified point
-	id *newRow = (id *)malloc(columnPointerByteSize);
-	dataStorage[index] = newRow;
-	numRows++;
-
-	// Copy over references to the array contents, and retain the objects
-	NSUInteger cellsCopied = 0;
-	for (id cellData in row) {
-		if (cellData) newRow[cellsCopied] = (id)CFRetain(cellData);
-		else newRow[cellsCopied] = nil;
-		if (++cellsCopied == numColumns) break;
-	}
+	// Add the new row to the editable store
+	[editedRows insertPointer:aRow atIndex:anIndex];
 
-	// If an array shorter than the row width was inserted, pad with nils
-	if (cellsCopied < numColumns) {
-		for ( ; cellsCopied <= numColumns; cellsCopied++) newRow[cellsCopied] = nil;
-	}
+	// Update the underlying store to keep counts and indices correct
+	[dataStorage insertDummyRowAtIndex:anIndex];
 }
 
 /**
  * Replace a row with contents of the supplied NSArray.
  */
-- (void) replaceRowAtIndex:(NSUInteger)index withRowContents:(NSArray *)row
+- (void) replaceRowAtIndex:(NSUInteger)anIndex withRowContents:(NSMutableArray *)aRow
 {
-	NSUInteger cellsProcessed = 0;
-
-	// Throw an exception if the index is out of bounds
-	if (index >= numRows) [NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)index, (unsigned long long)numRows];
-
-	id *storageRow = dataStorage[index];
-
-	// Iterate through the row replacing the objects
-	for (id cellData in row) {
-		if (storageRow[cellsProcessed]) CFRelease(storageRow[cellsProcessed]);
-		if (cellData) storageRow[cellsProcessed] = (id)CFRetain(cellData);
-		else storageRow[cellsProcessed] = nil;
-		if (++cellsProcessed == numColumns) break;
-	}
-
-	// Ensure all cells are correctly updated if an array shorter than the row width was supplied
-	if (cellsProcessed < numColumns) {
-		for ( ; cellsProcessed <= numColumns; cellsProcessed++) {
-			if (storageRow[cellsProcessed]) CFRelease(storageRow[cellsProcessed]);
-			storageRow[cellsProcessed] = nil;
-		}
-	}
+	[self _checkNewRow:aRow];
+	[editedRows replacePointerAtIndex:anIndex withPointer:aRow];
 }
 
 /**
  * Replace the contents of a single cell with a supplied object.
  */
-- (void) replaceObjectInRow:(NSUInteger)rowIndex column:(NSUInteger)columnIndex withObject:(id)object
+- (void) replaceObjectInRow:(NSUInteger)rowIndex column:(NSUInteger)columnIndex withObject:(id)anObject
 {
 
-	// Throw an exception of either index is out of bounds
-	if (rowIndex >= numRows || columnIndex >= numColumns) [NSException raise:NSRangeException format:@"Requested storage index (row %llu, col %llu) beyond bounds (%llu, %llu)", (unsigned long long)rowIndex, (unsigned long long)columnIndex, (unsigned long long)numRows, (unsigned long long)numColumns];
+	// Make sure that the row in question is editable
+	NSMutableArray *editableRow = SPDataStorageGetEditedRow(editedRows, rowIndex);
+	if (editableRow == NULL) {
+		editableRow = [self rowContentsAtIndex:rowIndex];
+		[editedRows replacePointerAtIndex:rowIndex withPointer:editableRow];
+	}
 
-	// Release the old object and retain the new one
-	if (dataStorage[rowIndex][columnIndex]) CFRelease(dataStorage[rowIndex][columnIndex]);
-	if (object) dataStorage[rowIndex][columnIndex] = (id)CFRetain(object);
-	else dataStorage[rowIndex][columnIndex] = nil;
+	// Modify the cell
+	[editableRow replaceObjectAtIndex:columnIndex withObject:anObject];
 }
 
 /**
  * Remove a row, renumbering all elements beyond index.
  */
-- (void) removeRowAtIndex:(NSUInteger)index
+- (void) removeRowAtIndex:(NSUInteger)anIndex
 {
 
 	// Throw an exception if the index is out of bounds
-	if (index >= numRows) [NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)index, (unsigned long long)numRows];
-
-	// Free the row
-	NSUInteger j = numColumns;
-	id *row = dataStorage[index];
-	while (j > 0) {
-		if (row[--j]) CFRelease(row[j]);
+	if (anIndex >= SPMySQLResultStoreGetRowCount(dataStorage)) {
+		[NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)anIndex, SPMySQLResultStoreGetRowCount(dataStorage)];
 	}
-	free(row);
-	numRows--;
 
-	// Renumber all subsequent indices to fill the gap
-	for (j = index; j < numRows; j++) {
-		dataStorage[j] = dataStorage[j + 1];
-	}
-	dataStorage[numRows] = NULL;
+	// Remove the row from the edited list and underlying storage
+	[editedRows removePointerAtIndex:anIndex];
+	[dataStorage removeRowAtIndex:anIndex];
 }
 
 /**
@@ -275,27 +334,16 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
 {
 
 	// Throw an exception if the range is out of bounds
-	if (rangeToRemove.location + rangeToRemove.length > numRows) [NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)(rangeToRemove.location + rangeToRemove.length), (unsigned long long)numRows];
-
-	// Free rows in the range
-	NSUInteger i, j = numColumns;
-	id *row;
-	for (i = rangeToRemove.location; i < rangeToRemove.location + rangeToRemove.length; i++) {
-		row = dataStorage[i];
-		while (j > 0) {
-			if (row[--j]) CFRelease(row[j]);
-		}
-		free(row);	
+	if (rangeToRemove.location + rangeToRemove.length > SPMySQLResultStoreGetRowCount(dataStorage)) {
+		[NSException raise:NSRangeException format:@"Requested storage index (%llu) beyond bounds (%llu)", (unsigned long long)(rangeToRemove.location + rangeToRemove.length), SPMySQLResultStoreGetRowCount(dataStorage)];
 	}
-	numRows -= rangeToRemove.length;
 
-	// Renumber all subsequent indices to fill the gap
-	for (i = rangeToRemove.location + rangeToRemove.length - 1; i < numRows; i++) {
-		dataStorage[i] = dataStorage[i + rangeToRemove.length];
-	}
-	for (i = numRows; i < numRows + rangeToRemove.length; i++) {
-		dataStorage[i] = NULL;
+	// Remove the rows from the edited list and underlying storage
+	NSUInteger i = rangeToRemove.location + rangeToRemove.length;
+	while (--i >= rangeToRemove.location) {
+		[editedRows removePointerAtIndex:i];
 	}
+	[dataStorage removeRowsInRange:rangeToRemove];
 }
 
 /**
@@ -303,81 +351,58 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
  */
 - (void) removeAllRows
 {
-	NSUInteger j;
-	id *row;
-
-	// Free all the data
-	while (numRows > 0) {
-		row = dataStorage[--numRows];
-		j = numColumns;
-		while (j > 0) {
-			if (row[--j]) CFRelease(row[j]);
-		}
-		free(row);
-	}
+	[editedRows setCount:0];
+	[dataStorage removeAllRows];
+}
 
-	numRows = 0;
+#pragma mark - Unloaded columns
+
+/**
+ * Mark a column as unloaded; SPNotLoaded placeholders will be returned for cells requested
+ * from this store which haven't had their value updated from elsewhere.
+ */
+- (void) setColumnAsUnloaded:(NSUInteger)columnIndex
+{
+	if (columnIndex >= numberOfColumns) {
+		[NSException raise:NSRangeException format:@"Invalid column set as unloaded; requested column index (%llu) beyond bounds (%llu)", (unsigned long long)columnIndex, (unsigned long long)numberOfColumns];
+	}
+	unloadedColumns[columnIndex] = true;
 }
 
-#pragma mark -
-#pragma mark Basic information
+#pragma mark - Basic information
 
 /**
  * Returns the number of rows currently held in data storage.
  */
 - (NSUInteger) count
 {
-	return numRows;
+	return (NSUInteger)[dataStorage numberOfRows];
 }
 
 /**
- * Set the number of columns represented by the data storage.
+ * Return the number of columns represented by the data storage.
  */
-- (void) setColumnCount:(NSUInteger)columnCount
+- (NSUInteger) columnCount
 {
-	columnPointerByteSize = columnCount * sizeof(id);
-
-	// If there are rows present in the storage, and the number of
-	// columns has changed, amend the existing rows to match.
-	if (columnCount != numColumns && numRows) {
-		NSUInteger i = numRows, j;
-		id *row;
-
-		// If the new column count is higher than the old count, iterate through the existing rows
-		// and pad with nils
-		if (columnCount > numColumns) {
-			while (i-- > 0) {
-				dataStorage[i] = (id *)realloc(dataStorage[i], columnPointerByteSize);
-				j = numColumns;
-				while (j < columnCount) {
-					dataStorage[i][j++] = nil;
-				}
-			}
-
-		// If the new column count is lower than the old count, iterate through the existing rows
-		// freeing any extra objects
-		} else {
-			while (i > 0) {
-				row = dataStorage[--i];
-				j = numColumns;
-				while (j > columnCount) {
-					if (row[--j]) CFRelease(row[j]);
-				}
-				dataStorage[i] = (id *)realloc(row, columnPointerByteSize);
-			}
-		}
-	}
+	return numberOfColumns;
+}
 
-	// Update the column count
-	numColumns = columnCount;
+/**
+ * Return whether all the data has been downloaded into the underlying result store.
+ */
+- (BOOL) dataDownloaded
+{
+	return [dataStorage dataDownloaded];
 }
 
+#pragma mark - Delegate callback methods
+
 /**
- * Return the number of columns represented by the data storage.
+ * When the underlying result store finishes downloading, update the row store to match
  */
-- (NSUInteger) columnCount
+- (void)resultStoreDidFinishLoadingData:(SPMySQLStreamingResultStore *)resultStore
 {
-	return numColumns;
+	[editedRows setCount:(NSUInteger)[resultStore numberOfRows]];
 }
 
 /**
@@ -387,20 +412,20 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
 
 - (id) init {
 	if ((self = [super init])) {
-		numColumns = 0;
-		columnPointerByteSize = 0;
-		numRows = 0;
+		dataStorage = nil;
+		editedRows = nil;
+		unloadedColumns = NULL;
 
-		// Initialise the array, initially with space for 100 rows
-		numRowsCapacity = 100;
-		dataStorage = (id **)malloc(numRowsCapacity * sizeof(id *));
+		numberOfColumns = 0;
 	}
 	return self;
 }
 
 - (void) dealloc {
-	[self removeAllRows];
-	free(dataStorage);
+	[dataStorage release], dataStorage = nil;
+	[editedRows release], editedRows = nil;
+	if (unloadedColumns) free(unloadedColumns), unloadedColumns = NULL;
+
 	[super dealloc];
 }
 
@@ -408,23 +433,12 @@ static inline void SPDataStorageEnsureCapacityForAdditionalRowCount(SPDataStorag
 
 @implementation SPDataStorage (PrivateAPI)
 
-/**
- * Private method to ensure the array always has sufficient capacity
- * to store any additional rows required.
- */
-- (void) _ensureCapacityForAdditionalRowCount:(NSUInteger)numExtraRows
+- (void) _checkNewRow:(NSMutableArray *)aRow
 {
-	while (numRows + numExtraRows > numRowsCapacity) [self _increaseCapacity];
+	if ([aRow count] != numberOfColumns) {
+		[NSException raise:NSInternalInconsistencyException format:@"New row length (%llu) does not match store column	count (%llu)", (unsigned long long)[aRow count], (unsigned long long)numberOfColumns];
+	}
 }
 
-/**
- * Private method to increase the storage available for the array;
- * currently doubles the capacity as boundaries are reached.
- */
-- (void) _increaseCapacity
-{
-	numRowsCapacity *= 2;
-	dataStorage = (id **)realloc(dataStorage, numRowsCapacity * sizeof(id *));
-}
 
 @end