From 92d28e468e31fc7dff79f26c7cc4a7feef5031c1 Mon Sep 17 00:00:00 2001
From: Francisco Ryan Tolmasky I <francisco@280north.com>
Date: Mon, 15 Feb 2010 22:34:50 -0800
Subject: [PATCH] Autogenerate the index parameter for Fast Enumeration, and
 add Fast Enumeration support to CPIndexSet.

Reviewed by me.
---
 Foundation/CPArray.j    | 13 +------------
 Foundation/CPIndexSet.j | 24 +++++++++++++++++++++++
 Objective-J/Runtime.js  | 42 +++++++++++++++++++++++++++++++++--------
 3 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/Foundation/CPArray.j b/Foundation/CPArray.j
index de693148a..e070cae94 100755
--- a/Foundation/CPArray.j
+++ b/Foundation/CPArray.j
@@ -1215,20 +1215,9 @@
     if (aState.state >= count)
         return 0;
 
-    aState.items0 = self;
+    aState.items = self;
     aState.state = count;
 
-    if (aState.assigneeCount > 1)
-    {
-        var index = 0,
-            indexes = [];
-
-        for (; index < count; ++index)
-            indexes[index] = index;
-
-        aState.items1 = indexes;
-    }
-
     return count;
 }
 
diff --git a/Foundation/CPIndexSet.j b/Foundation/CPIndexSet.j
index 03b427d34..02d6424c2 100644
--- a/Foundation/CPIndexSet.j
+++ b/Foundation/CPIndexSet.j
@@ -834,6 +834,30 @@ var CPIndexSetCountKey              = @"CPIndexSetCountKey",
 
 @end
 
+@implementation CPIndexSet (CPFastEnumeration)
+
+- (int)countByEnumeratingWithState:(id)aState objects:(id)objects count:(id)aCount
+{
+    var rangeIndex = aState.state;
+
+    if (aState.state >= _ranges.length)
+        return 0;
+
+    var range = _ranges[rangeIndex],
+        index = 0,
+        start = range.location,
+        end = CPMaxRange(range);
+
+    for (; start < end; ++start, ++index)
+        objects[index] = start;
+
+    ++aState.state;
+
+    return range.length;
+}
+
+@end
+
 /*!
     @class CPMutableIndexSet
     @ingroup compatability
diff --git a/Objective-J/Runtime.js b/Objective-J/Runtime.js
index c87783b74..644c02358 100644
--- a/Objective-J/Runtime.js
+++ b/Objective-J/Runtime.js
@@ -549,6 +549,7 @@ function objj_fastEnumerator(/*Object*/ anObject, /*Integer*/ anAssigneeCount)
         this._target = anObject;
 
     this._state = { state:0, assigneeCount:anAssigneeCount };
+    this._index = 0;
 
     // Nothing to iterate in this case.
     if (!anObject)
@@ -569,27 +570,52 @@ objj_fastEnumerator.prototype.e = function()
     var state = this._state,
         index = state.assigneeCount;
 
-    this.items = nil;
-    this.itemsPtr = nil;
+    // Clear out all the old state.
+    state.items = nil;
+    state.itemsPtr = nil;
 
     while (index--)
         state["items" + index] = nil;
 
-    // This is safer, but possibly slower.
     this.o0 = [];
     this.i = 0;
     this.l = objj_msgSend(this._target, fastEnumerationSelector, state, this.o0, 10);
 
+    // We're flexible on this.
     this.o0 = state.items || state.itemsPtr || state.items0 || this.o0;
 
-    index = state.assigneeCount;
-
-    while (index--)
-        this["o" + index] = state["items" + index] || [];
-
+    // We allow the user to not explictly return anything in countByEnumeratingWithState:objects:count:
     if (this.l === undefined)
         this.l = this.o0.length;
 
+    var assigneeCount = state.assigneeCount;
+
+    index = assigneeCount - 1;
+
+    // Handle all items from [1 .. assigneeCount - 1]
+    while (index-- > 1)
+        this["o" + index] = state["items" + index] || [];
+
+    var lastAssigneeIndex = assigneeCount - 1;
+
+    // Autogenerate the indexes if this was left blank.
+    if (lastAssigneeIndex > 0)
+
+        if (state["items" + lastAssigneeIndex])
+            this["o" + lastAssigneeIndex] = state["items" + lastAssigneeIndex];
+
+        else
+        {
+            var count = this.l,
+                indexIndex = 0,
+                indexes = new Array(count)
+
+            for (; indexIndex < count; ++indexIndex, ++this._index)
+                indexes[indexIndex] = this._index;
+
+            this["o" + lastAssigneeIndex] = indexes;
+        }
+
     return this.l > 0;
 }