Implementing countByEnumeratingWithState:objects:count:

Some sample NSFastEnumeration implementations

You need to implement NSFastEnumeration if you want to use your own classes in the Objective-C 2.0 "for...in" fast-enumeration language feature.

The documentation for its only method (countByEnumeratingWithState:objects:count:) is pretty cryptic and overwhelming — mostly because the method needs to be flexible and Apple don't want to tell you how to write your code. I have no such qualms and since there aren't many examples on the Internets at the moment I thought I'd show you some easy approaches for implementing it in your own programs.

There are two different ways to enumerate. The first is where your class already has, or is willing to create, a C array of Objective-C id values which point to the objects being enumerated. The second is where you don't have this storage and want to use storage passed to you.

Warning: These examples do not use the mutationsPtr and are therefore only safe for use with immutable collections. If the enumerated collection is mutable, you will need to point mutationsPtr to an appropriate mutation guard value (normally a mutations count value).

First case: already have a C array of storage

Our class looks like this:

@interface SimpleStringArray : NSObject <NSFastEnumeration>
{
    NSString *stringArray[ARRAY_LENGTH];
}

In this case, the implementation of countByEnumeratingWithState:objects:count: would look like this:

- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state objects:(id *)stackbuf count:(NSUInteger)len
{
    if (state->state >= ARRAY_LENGTH)
    {
        return 0;
    }

    state->itemsPtr = stringArray;
    state->state = ARRAY_LENGTH;
    state->mutationsPtr = (unsigned long *)self;
    
    return ARRAY_LENGTH;
}

Quick explanation:

• we've ignored stackbuf and count (because we already have storage)
• as required, we've set state->state to a non zero value (the iteration index after the current items are iterated)
• as required, we've set state->mutationsPtr to a non zero value (the self pointer since we have no "array has changed" flag that we can point it to)
• we've returned the complete length of the array, so it will all be iterated in one pass (on the second pass, state->state will equal ARRAY_LENGTH and we'll return "0", ending the loop)

Second case: need storage

If our class' data is a collection of NSString objects stored in a list of linearly connected C structs, then we might have no C array value to return.

Assume our list is declared like this:

typedef struct
{
    NSString * stringPtr;
    struct MyList * nextNode;
} MyList;

and our class is declared like this:

@interface StringList : NSObject <NSFastEnumeration>
{
    MyList * _startOfListNode;
    MyList * _endOfListPlusOneNode;
}

This is more complicated example because we need to actually gather the data from the list and store our traversal state between iterations. Here is how it might look:

- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state objects:(id *)stackbuf count:(NSUInteger)len
{
    MyList *currentNode;
    if (state->state == 0)
    {
        // Set the starting point. _startOfListNode is assumed to be our
        // object's instance variable that points to the start of the list.
        currentNode = _startOfListNode;
    }
    else
    {
        // Subsequent iterations, get the current progress out of state->state
        currentNode = (struct MyList *)state->state;
    }
    
    // Accumulate nodes from the list until we reach the object's
    // _endOfListPlusOneNode
    NSUInteger batchCount = 0;
    while (currentNode != _endOfListPlusOneNode && batchCount < len)
    {
        stackbuf[batchCount] = currentNode->stringPtr;
        currentNode = currentNode->nextNode;
        batchCount++;
    }

    state->state = (unsigned long)currentNode;
    state->itemsPtr = stackbuf;
    state->mutationsPtr = (unsigned long *)self;

    return batchCount;
}

Explanation of this example:

• We now use stackbuf to store the data accumulated from the list and we return it in state->itemsPtr.
• len is used now because it is the maximum number of objects we can accumulate in stackbuf each time.
• When executed the first time (when state->state == 0), we set the currentNode to our object's private _startOfListNode — that we assume we have correctly set to the start of our list.
• Every other time we are run, state->state will hold our currentNode, saved from the last iteration.
• The fast enumeration will repeatedly invoke this method until we reach a currentNode == _endOfListPlusOneNode at the end of the list — remember that state->state is not allowed to be nil, which is why we use this non-nil end marker node. If you use a list that ends in nil, then you should detect this and set state->state to a special non-nil "end" flag so that you won't get trapped in an infinite loop.