Implement fast enumeration for PostgresKit's result class. Based on SPMySQL's implementation by Rowan.

2 files changed, 31 insertions, 1 deletions

diff --git a/Frameworks/PostgresKit/Source/FLXPostgresResult.h b/Frameworks/PostgresKit/Source/FLXPostgresResult.h
index 168f8f15..6f10c3e5 100644
--- a/Frameworks/PostgresKit/Source/FLXPostgresResult.h
+++ b/Frameworks/PostgresKit/Source/FLXPostgresResult.h
@@ -30,7 +30,7 @@ typedef enum
 } 
 FLXPostgresResultRowType;
 
-@interface FLXPostgresResult : NSObject
+@interface FLXPostgresResult : NSObject <NSFastEnumeration>
 {
 	void *_result;
 	void **_typeHandlers;
diff --git a/Frameworks/PostgresKit/Source/FLXPostgresResult.m b/Frameworks/PostgresKit/Source/FLXPostgresResult.m
index 06950a3b..02375650 100644
--- a/Frameworks/PostgresKit/Source/FLXPostgresResult.m
+++ b/Frameworks/PostgresKit/Source/FLXPostgresResult.m
@@ -175,6 +175,36 @@
 }
 
 #pragma mark -
+#pragma mark Fast enumeration implementation
+
+- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state objects:(id *)stackbuf count:(NSUInteger)len
+{
+	if (state->state >= _numberOfRows) return 0;
+	
+	if (state->state != _row) [self seekToRow:state->state];
+	
+	// Determine how many objects to return - 128, len, or all items remaining
+	NSUInteger itemsToReturn = 128;
+	
+	if (len < 128) itemsToReturn = len;
+	
+	if (_numberOfRows - state->state < itemsToReturn) {
+		itemsToReturn = (unsigned long)_numberOfRows - state->state;
+	}
+	
+	for (NSUInteger i = 0; i < itemsToReturn; i++) 
+	{
+		stackbuf[i] = [self rowAsType:_defaultRowType];
+	}
+	
+	state->state += itemsToReturn;
+	state->itemsPtr = stackbuf;
+	state->mutationsPtr = (unsigned long *)self;
+	
+	return itemsToReturn;
+}
+
+#pragma mark -
 #pragma mark Private API
 
 /**