PSYEnumerator-Examples.m

/*
 Example.m
 Created by Remy "Psy" Demarest on 21/04/2012.
 
 Copyright (c) 2012. Remy "Psy" Demarest

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 copies of the Software, and to permit persons to whom the Software is
 furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included in
 all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 */

#import <Foundation/Foundation.h>
#import "PSYEnumerator.h"

int main(int argc, char *argv[])
{
    @autoreleasepool
    {
        NSDictionary *dict1 = [NSDictionary dictionaryWithObjectsAndKeys:@"dict1 ", @"a", @"dict1 ", @"b", @"dict1 ", @"c", @"dict1 ", @"d", @"dict1 ", @"e", @"dict1 ", @"f", @"dict1 ", @"g", @"dict1 ", @"h", @"dict1 ", @"i", nil];
        NSDictionary *dict2 = [NSDictionary dictionaryWithObjectsAndKeys:@"dict2 ", @"a", @"dict2 ", @"c", @"dict2 ", @"e", @"dict2 ", @"g", @"dict2 ", @"i", @"dict2 ", @"j", @"dict2 ", @"k", @"dict2 ", @"l", @"dict2 ", @"m", nil];
        NSDictionary *dict3 = [NSDictionary dictionaryWithObjectsAndKeys:@"dict3 ", @"a", @"dict3 ", @"e", @"dict3 ", @"i", @"dict3 ", @"k", @"dict3 ", @"m", @"dict3 ", @"n", @"dict3 ", @"o", @"dict3 ", @"p", @"dict3 ", @"q", nil];

        NSLog(@"Dictionary enumeration, stop as soon as all keys common to all dictionaries are enumerated:");
        [@[dict1, dict2, dict3] enumerateDictionariesWithOptions:0 usingBlock:
         ^(id key, __unsafe_unretained id const *objects, BOOL *stop)
         {
             NSLog(@"dict[%@] = %@ %@ %@", key, objects[0], objects[1], objects[2]);
         }];
        NSLog(@"--------------");
        
        NSLog(@"Dictionary enumeration, stop when all keys in all dictionaries are enumerated:");
        [@[dict1, dict2, dict3] enumerateDictionariesWithOptions:PSYMultiEnumerationFinishAllCollections usingBlock:
         ^(id key, __unsafe_unretained id const *objects, BOOL *stop)
         {
             NSLog(@"dict[%@] = %@ %@ %@", key, objects[0], objects[1], objects[2]);
         }];
        NSLog(@"--------------");

        NSDictionary *dict  = [NSDictionary dictionaryWithObjectsAndKeys:@"a", @"b", @"c", @"d", @"e", @"f", @"g", @"h", @"i", @"j", nil];
        NSArray      *array = [NSArray arrayWithObjects:@"k", @"l", @"m", @"n", @"o", @"p", @"q", nil];
        NSSet        *set   = [NSSet setWithObjects:@"r", @"s", @"t", @"u", @"v", @"w", @"x", @"y", @"z", nil];

        NSLog(@"Collection enumeration, stop as soon as one of the collection runs out of objects:");
        [@[dict, array, set] enumerateCollectionsWithOptions:0 usingBlock:
         ^(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop)
         {
             NSLog(@"Objects[%ld]: %@ %@ %@", idx, objects[0], objects[1], objects[2]);
         }];
        NSLog(@"--------------");

        NSLog(@"Collection enumeration, stop when all objects of all collections were enumerated:");
        [@[dict, array, set] enumerateCollectionsWithOptions:PSYMultiEnumerationFinishAllCollections usingBlock:
         ^(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop)
         {
             NSLog(@"Objects[%ld]: %@ %@ %@", idx, objects[0], objects[1], objects[2]);
         }];

        NSArray *enumerables = @[dict, array, set];
        NSUInteger count = [enumerables count];
        __block NSUInteger latest = 0;

        NSLog(@"--------------");

        NSLog(@"Collection enumeration, only enumerate a collection if its slot is set to nil by the block and stop when all objects of all collections were enumerated:");
        [enumerables enumerateCollectionsWithOptions:PSYMultiEnumerationFinishAllCollections | PSYMultiEnumerationConsuming usingBlock:
         ^(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop)
         {
             latest %= count;

             while(objects[latest] == nil)
             {
                 latest++;
                 latest %= count;
             }

             NSLog(@"Consumed Object: enum[%ld][%ld] = %@", latest, idx, objects[latest]);
             objects[latest] = nil;

             latest++;
         }];
        NSLog(@"--------------");
        NSLog(@"Merge and sort using collection zipping in consuming mode:");

        NSArray *airNomads = @[ @"Aang (Air)", @"Afiko (Air)", @"Appa (Air)", @"Gyatso (Air)", @"Iio (Air)", @"Jinju (Air)", @"Malu (Air)", @"Momo (Air)", @"Oogi (Air)", @"Pasang (Air)", @"Tashi (Air)", @"Yangchen (Air)" ];
        NSArray *earthKingdom = @[ @"Bosco (Earth)", @"Bumi (Earth)", @"Cabbage merchant (Earth)", @"Chin (Earth)", @"Chong (Earth)", @"Flopsie (Earth)", @"Fluffykins (Earth)", @"Fong (Earth)", @"Foo Foo Cuddlypoops (Earth)", @"Fung (Earth)", @"Gansu (Earth)", @"Ghashiun (Earth)", @"Gow (Earth)", @"Haru (Earth)", @"Headhunter (Earth)", @"Herbalist (Earth)", @"Hope (Earth)", @"How (Earth)", @"Jet (Earth)", @"Jin (Earth)", @"Jin Wei (Earth)", @"Joo Dee (Earth)", @"June (Earth)", @"Kenji (Earth)", @"Kuei (Earth)", @"Kyoshi (Earth)", @"Lao Beifong (Earth)", @"Lee (Earth)", @"Lily (Earth)", @"Lin Yee (Earth)", @"Long Feng (Earth)", @"Longshot (Earth)", @"Macmu-Ling (Earth)", @"Meng (Earth)", @"Miyuki (Earth)", @"Moku (Earth)", @"Nyla (Earth)", @"Oh (Earth)", @"Oyaji (Earth)", @"Pao (Earth)", @"Pipsqueak (Earth)", @"Pong (Earth)", @"Poppy Beifong (Earth)", @"Quon (Earth)", @"Sela (Earth)", @"Sensu (Earth)", @"Sha-Mo (Earth)", @"Shu (Earth)", @"Smellerbee (Earth)", @"Song (Earth)", @"Star (Earth)", @"Sud (Earth)", @"Suki (Earth)", @"Sung (Earth)", @"Tahn (Earth)", @"Teo (Earth)", @"The Big Bad Hippo (Earth)", @"The Boulder (Earth)", @"The Duke (Earth)", @"The Gecko (Earth)", @"The Gopher (Earth)", @"Tong (Earth)", @"Toph Beifong (Earth)", @"Tycho (Earth)", @"Tyro (Earth)", @"Wu (Earth)", @"Xin Fu (Earth)", @"Yu (Earth)", @"Yung (Earth)", @"Zei (Earth)" ];
        NSArray *fireNation = @[ @"Azula (Fire)", @"Azulon (Fire)", @"Bujing (Fire)", @"Chan (Fire)", @"Chey (Fire)", @"Chit Sang (Fire)", @"Combustion Man (Fire)", @"Ding (Fire)", @"Dock (Fire)", @"Fang (Fire)", @"Fat (Fire)", @"Ham Ghao (Fire)", @"Hawky (Fire)", @"Hide (Fire)", @"Iroh (Fire)", @"Jee (Fire)", @"Jeong Jeong (Fire)", @"Kahchi (Fire)", @"Kwan (Fire)", @"Li (Fire)", @"Lu Ten (Fire)", @"Mai (Fire)", @"Mak (Fire)", @"Mongke (Fire)", @"Mung (Fire)", @"Ogodei (Fire)", @"On Ji (Fire)", @"Ozai (Fire)", @"Poon (Fire)", @"Qin (Fire)", @"Roku (Fire)", @"Shinu (Fire)", @"Shoji (Fire)", @"Shyu (Fire)", @"Sozin (Fire)", @"Ty Lee (Fire)", @"Ursa (Fire)", @"Vachir (Fire)", @"Yao (Fire)", @"Yeh-Lu (Fire)", @"Yon Rha (Fire)", @"Zhao (Fire)", @"Zuko (Fire)" ];
        NSArray *waterTribe = @[ @"Amon (Water)", @"Arnook (Water)", @"Bato (Water)", @"Due (Water)", @"Hahn (Water)", @"Hakoda (Water)", @"Hama (Water)", @"Huu (Water)", @"Kanna (Water)", @"Katara (Water)", @"Korra (Water)", @"Kuruk (Water)", @"Kya (Water)", @"La (Water)", @"Naga (Water)", @"Pakku (Water)", @"Sangok (Water)", @"Senna (Water)", @"Slim (Water)", @"Sokka (Water)", @"Tarrlok (Water)", @"Tho (Water)", @"Tonraq (Water)", @"Tui (Water)", @"Ummi (Water)", @"Wei (Water)", @"Yakone (Water)", @"Yue (Water)", @"Yugoda (Water)" ];

        NSArray *arrays = @[ airNomads, earthKingdom, fireNation, waterTribe ];
        count = [arrays count];
        NSArray *result = [arrays arrayByEnumeratingCollectionsWithOptions:PSYMultiEnumerationConsuming | PSYMultiEnumerationFinishAllCollections usingBlock:
                           ^ NSString * (NSUInteger idx, __unsafe_unretained NSString **objects, BOOL *stop)
                           {
                               NSString *result = nil;
                               NSUInteger found = NSNotFound;
                               for(NSUInteger i = 0; i < count; i++)
                               {
                                   NSString *current = objects[i];
                                   if(current == nil) continue;

                                   if(result == nil || [result caseInsensitiveCompare:current] == NSOrderedDescending)
                                   {
                                       result = current;
                                       found = i;
                                   }
                               }

                               if(found != NSNotFound)
                               {
                                   result = objects[found];
                                   objects[found] = nil;
                               }

                               return result;
                           }];

        NSLog(@"Result: %@", result);
    }

    return 0;
}

PSYEnumerator.h

/*
 PSYEnumerator.h
 Created by Remy "Psy" Demarest on 28/04/2012.
 
 Copyright (c) 2012. Remy "Psy" Demarest

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 copies of the Software, and to permit persons to whom the Software is
 furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included in
 all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 */

typedef NS_OPTIONS(NSUInteger, PSYMultiEnumerationOptions) {
    PSYMultiEnumerationFinishAllCollections = 1 << 0,
    PSYMultiEnumerationConsuming            = 1 << 1, // Only works with Collections enumeration
};

@interface NSArray (PSYMutliEnumeration)
- (void)enumerateCollectionsWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(void(^)(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop))block;
- (NSArray *)arrayByEnumeratingCollectionsWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(id(^)(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop))block;

- (void)enumerateDictionariesWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(void(^)(id key, __unsafe_unretained const id *objects, BOOL *stop))block;
- (NSDictionary *)dictionaryByEnumeratingDictionariesWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(id(^)(id key, __unsafe_unretained const id *objects, BOOL *stop))block;
@end

PSYEnumerator.m

/*
 PSYEnumerator.m
 Created by Remy "Psy" Demarest on 21/04/2012.
 
 Copyright (c) 2012. Remy "Psy" Demarest

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 copies of the Software, and to permit persons to whom the Software is
 furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included in
 all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 */

#import "PSYEnumerator.h"

#define OBJECT_BUFFER_SIZE 8 /* I don't necessary want to allocate too much memory there... */

typedef struct _PSYMultiEnumState {
    __unsafe_unretained id objectsBuf[OBJECT_BUFFER_SIZE]; 
    NSFastEnumerationState state;
    unsigned long          mutationPtrValue;
    NSUInteger             count;
    NSUInteger             position;
} PSYMultiEnumState;

@implementation NSArray (PSYMutliEnumeration)

- (void)enumerateDictionariesWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(void(^)(id key, __unsafe_unretained const id *objects, BOOL *stop))block
{
    const NSUInteger count = [self count];
    if(count == 0) return;

    BOOL finishAll = !!(options & PSYMultiEnumerationFinishAllCollections);

    __unsafe_unretained id *objects = (__unsafe_unretained id *)calloc(count, sizeof(*objects));
    __unsafe_unretained NSDictionary **enumerables = (__unsafe_unretained NSDictionary **)calloc(count, sizeof(NSDictionary *));

    [self getObjects:enumerables range:NSMakeRange(0, count)];

    NSMutableSet *alreadyDone = finishAll ? [NSMutableSet setWithCapacity:[enumerables[0] count]] : nil;

    PSYMultiEnumState currentState = { 0 };

    BOOL stop = NO;

    for(NSUInteger baseIdx = 0; baseIdx < count; baseIdx++)
    {
        currentState = (PSYMultiEnumState){ 0 };
        while((currentState.count = [enumerables[baseIdx] countByEnumeratingWithState:&currentState.state objects:currentState.objectsBuf count:OBJECT_BUFFER_SIZE]))
        {
            for(NSUInteger i = 0; i < currentState.count; i++)
            {
                id key = currentState.state.itemsPtr[i];

                if([alreadyDone containsObject:key]) continue;
                [alreadyDone addObject:key];

                BOOL shouldEnumerate = YES;

                for(NSUInteger j = baseIdx; j < count; j++)
                {
                    objects[j] = [enumerables[j] objectForKey:key];

                    if(!finishAll && objects[j] == nil)
                    {
                        shouldEnumerate = NO;
                        break;
                    }
                }

                if(shouldEnumerate) block(key, objects, &stop);

                if(stop) break;
            }

            if(stop) break;
        }

        if(stop || !finishAll) break;

        // This value at this index will remain nil until the end of the whole enumeration
        objects[baseIdx] = nil;
    }

    free(objects);
    free(enumerables);
}

- (void)enumerateCollectionsWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(void(^)(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop))block
{
    const NSUInteger count = [self count];
    if(count == 0) return;

    const BOOL finishAll = !!(options & PSYMultiEnumerationFinishAllCollections);
    const BOOL consuming = !!(options & PSYMultiEnumerationConsuming);

    PSYMultiEnumState      *states  = calloc(count, sizeof(*states));
    __unsafe_unretained id *objects = (__unsafe_unretained id *)calloc(count, sizeof(*objects));

    __unsafe_unretained id<NSObject, NSFastEnumeration> *enumerables = (__unsafe_unretained id<NSObject, NSFastEnumeration> *)calloc(count, sizeof(id<NSObject, NSFastEnumeration>));
    [self getObjects:enumerables range:NSMakeRange(0, count)];

    BOOL stop = NO, isFirstLoop = YES;
    NSUInteger idx = 0;

    while(YES)
    {
        BOOL hasNonZeroCount = NO;

        // Update each enumeration states
        for(NSUInteger i = 0; i < count; i++)
        {
            if((!consuming || objects[i] == nil) && states[i].position >= states[i].count)
            {
                // Sentinel to avoid calling the count method over and over
                if(states[i].count == 0 && states[i].position == NSNotFound) continue;

                states[i].count = [enumerables[i] countByEnumeratingWithState:&states[i].state objects:states[i].objectsBuf count:OBJECT_BUFFER_SIZE];

                if(states[i].count == 0)
                {
                    states[i].position = NSNotFound;

                    if(!finishAll)
                    {
                        hasNonZeroCount = NO;
                        break;
                    }
                }
                else
                {
                    states[i].position = 0;

                    if(!isFirstLoop && states[i].mutationPtrValue != *states[i].state.mutationsPtr)
                        @throw [NSException exceptionWithName:NSGenericException reason:[NSString stringWithFormat:@"Collection <%@: %p> was mutated while being enumerated.", [enumerables[i] class], enumerables[i]] userInfo:nil];

                    hasNonZeroCount = YES;
                }
            }
            else hasNonZeroCount = YES;

            if((!consuming || objects[i] == nil)) objects[i] = (states[i].position < states[i].count ? states[i].state.itemsPtr[states[i].position++] : nil);
        }

        // All enumeratable objects returned 0, we're done enumerating
        if(!hasNonZeroCount) break;

        block(idx, objects, &stop);

        if(stop) break;

        idx++;
    }

    free(states);
    free(objects);
    free(enumerables);
}

- (NSArray *)arrayByEnumeratingCollectionsWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(id(^)(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop))block
{
    if([self count] == 0) return @[];

    NSMutableArray *result = [NSMutableArray arrayWithCapacity:[self[0] count]];

    [self enumerateCollectionsWithOptions:options usingBlock:
     ^(NSUInteger idx, __unsafe_unretained id *objects, BOOL *stop)
     {
         id value = block(idx, objects, stop);

         if(value != nil) [result addObject:value];
     }];

    return result;
}

- (NSDictionary *)dictionaryByEnumeratingDictionariesWithOptions:(PSYMultiEnumerationOptions)options usingBlock:(id(^)(id key, __unsafe_unretained const id *objects, BOOL *stop))block
{
    if([self count] == 0) return @{};

    NSMutableDictionary *result = [NSMutableDictionary dictionaryWithCapacity:[self[0] count]];

    [self enumerateDictionariesWithOptions:options usingBlock:
     ^(id key, __unsafe_unretained id const *objects, BOOL *stop)
     {
         id value = block(key, objects, stop);

         if(value != nil) result[key] = value;
     }];

    return result;
}

@end