APEFunctionalExtension

APEFunctionalExtension.h

//
// Created by Chenyu Lan on 9/4/14.
// Copyright (c) 2014 Fenbi. All rights reserved.
//

#import <Foundation/Foundation.h>

@interface NSArray (APEFunctionalExtension)

- (NSArray *)map:(id (^)(id value))transformBlock;
- (NSArray *)filter:(BOOL (^)(id value))transformBlock;
- (NSArray *)reject:(BOOL (^)(id value))transformBlock;
- (id)reduce:(id)initial combineBlock:(id (^)(id value0, id value1))combineBlock;

@end

@interface NSSet (APEFunctionalExtension)

- (NSSet *)map:(id (^)(id value))transformBlock;
- (NSSet *)filter:(BOOL (^)(id value))transformBlock;
- (NSSet *)reject:(BOOL (^)(id value))transformBlock;
- (id)reduce:(id)initial combineBlock:(id (^)(id value0, id value1))combineBlock;

@end

@interface APEFunctionalExtensionPrivate : NSObject

+ (id)ape_filterCollection:(id<NSFastEnumeration>)collection
                     block:(BOOL (^)(id value))transformBlock
              filterToggle:(BOOL)toggle;

+ (id)ape_mapCollection:(id<NSFastEnumeration>)collection
                  block:(id (^)(id value))transformBlock;

+ (id)ape_reduceCollection:(id<NSFastEnumeration>)collection
                   initial:(id)initial
                     block:(id (^)(id value0, id value1))combineBlock;

@end

APEFunctionalExtension.m

//
// Created by Chenyu Lan on 9/4/14.
// Copyright (c) 2014 Fenbi. All rights reserved.
//

#import "APEFunctionalExtension.h"

@implementation NSArray (APEFunctionalExtension)

- (NSArray *)map:(id (^)(id value))transformBlock {
    return [APEFunctionalExtensionPrivate ape_mapCollection:self block:transformBlock];
}

- (NSArray *)filter:(BOOL (^)(id value))transformBlock {
    return [APEFunctionalExtensionPrivate ape_filterCollection:self block:transformBlock filterToggle:YES];
}

- (NSArray *)reject:(BOOL (^)(id value))transformBlock {
    return [APEFunctionalExtensionPrivate ape_filterCollection:self block:transformBlock filterToggle:NO];
}

- (id)reduce:(id)initial combineBlock:(id (^)(id value0, id value1))combineBlock {
    return [APEFunctionalExtensionPrivate ape_reduceCollection:self initial:initial block:combineBlock];
}

@end

@implementation NSSet (APEFunctionalExtension)

- (NSSet *)map:(id (^)(id value))transformBlock {
    return [APEFunctionalExtensionPrivate ape_mapCollection:self block:transformBlock];
}

- (NSSet *)filter:(BOOL (^)(id value))transformBlock {
    return [APEFunctionalExtensionPrivate ape_filterCollection:self block:transformBlock filterToggle:YES];
}

- (NSSet *)reject:(BOOL (^)(id value))transformBlock {
    return [APEFunctionalExtensionPrivate ape_filterCollection:self block:transformBlock filterToggle:NO];
}

- (id)reduce:(id)initial combineBlock:(id (^)(id value0, id value1))combineBlock {
    return [APEFunctionalExtensionPrivate ape_reduceCollection:self initial:initial block:combineBlock];
}

@end

@implementation APEFunctionalExtensionPrivate

+ (id)ape_filterCollection:(id<NSFastEnumeration>)collection
                     block:(BOOL (^)(id value))transformBlock
              filterToggle:(BOOL)toggle {
    Class class = [self ape_classFromCollection:collection];
    id result = [class new];
    for (id value in collection) {
        if (transformBlock(value) == toggle) {
            // cast to NSMutableArray to perform addObject: and avoid compile error
            [(NSMutableArray *)result addObject:value];
        }
    }
    return result;
}

+ (id)ape_mapCollection:(id<NSFastEnumeration>)collection
                  block:(id (^)(id value))transformBlock {
    Class class = [self ape_classFromCollection:collection];
    id result = [class new];
    for (id value in collection) {
        // cast to NSMutableArray to perform addObject: and avoid compile error
        [(NSMutableArray *)result addObject:transformBlock(value)];
    }
    return result;
}

+ (id)ape_reduceCollection:(id<NSFastEnumeration>)collection initial:(id)initial block:(id (^)(id, id))combineBlock {
    id result = initial;
    for (id value in collection) {
        result = combineBlock(result, value);
    }
    return result;
}


+ (Class)ape_classFromCollection:(id)collection {
    Class class;
    if ([collection isKindOfClass:[NSArray class]]) {
        class = [NSMutableArray class];
    } else if ([collection isKindOfClass:[NSSet class]]) {
        class = [NSMutableSet class];
    }
    NSAssert(class != NULL, @"APEFunctionalExtension Error: Unsupport Class: %@", NSStringFromClass([collection class]));
    return class;
}

@end